codes/MMSelfSup_tutorials/【1】模型自监督预训练 之 SimCLR.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/mmselfsup_logo.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型自监督预训练 之 SimCLR\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/open-mmlab/OpenMMLabCourse/blob/main/codes/MMSelfSup_tutorials/【1】模型自监督预训练%20之%20SimCLR.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**MMSelfSup Repo**：[https://github.com/open-mmlab/mmselfsup](https://github.com/open-mmlab/mmselfsup)\n",
    "\n",
    "**MMSelfSup 官方文档链接**：[https://mmselfsup.readthedocs.io/en/latest](https://mmselfsup.readthedocs.io/en/latest)\n",
    "\n",
    "**MMSelfSup 视频教学**：[https://space.bilibili.com/1293512903/channel/collectiondetail?sid=657287](https://space.bilibili.com/1293512903/channel/collectiondetail?sid=657287)\n",
    "\n",
    "**MMSelfSup 代码库介绍 PPT 获取方式**：关注 OpenMMLab 公众号，后台回复：mmselfsup，即可获取课程 PPT\n",
    "\n",
    "**加入微信社群方式**：关注公众号，选择 “加入我们” -> “微信社区”，即可获取入群二维码。非常期待你的到来呀~\n",
    "\n",
    "**作者**：OpenMMLab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. 自监督预训练方法介绍：SimCLR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**论文地址**：https://arxiv.org/pdf/2002.05709.pdf\n",
    "\n",
    "**SimCLR 基本思想**：对一张图片做两次不同的数据增强操作，增强后的两张图片互为彼此的正样本，同一个 batch 里其他图片的增强结果为这两张增强图片的负样本。SimCLR 要求编码器最大化当前图像与其正样本表示的相似度，最小化当前图像与其负样本表示的相似度。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/SimCLR.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 环境配置"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 查看 Python、PyTorch 和 Torchvision 的版本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check nvcc version\n",
    "!nvcc -V"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check GCC version\n",
    "!gcc --version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check PyTorch installation\n",
    "import torch, torchvision\n",
    "print(torch.__version__)\n",
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 安装 MMSelfSup 的依赖库：MMCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install openmim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mim install mmcv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3  安装 MMSelfSup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/open-mmlab/mmselfsup.git\n",
    "%cd /content/mmselfsup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install MMSelfSup from source\n",
    "!pip install -e . "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 检查安装是否正确"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mmselfsup\n",
    "print(mmselfsup.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 准备数据集"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.0 数据集介绍"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本教程将在 `Tiny ImageNet` 数据集上训练自监督模型 SimCLR。\n",
    "\n",
    "Tiny ImageNet 数据集是 ImageNet 的一个子集。\n",
    "\n",
    "该数据集包含 200 个类别，每个类别有 500 张训练图片、50 张验证图片和 50 张测试图片，共 120,000 张图像。每张图片均为 64×64 彩色图片。\n",
    "\n",
    "数据集官方下载地址：http://cs231n.stanford.edu/tiny-imagenet-200.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 下载数据集"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用 GNU [Wget](https://www.gnu.org/software/wget/) 工具从斯坦福官方网站下载：http://cs231n.stanford.edu/tiny-imagenet-200.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/mmselfsup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir data\n",
    "%cd data\n",
    "!wget http://cs231n.stanford.edu/tiny-imagenet-200.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 解压数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip -q tiny-imagenet-200.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf tiny-imagenet-200.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 查看数据集目录"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check data directory\n",
    "!apt-get install tree\n",
    "!tree -d /content/mmselfsup/data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 准备标注文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为了减少大家重写 `加载数据集` 代码的负担，我们整理好了标注文件，复制到数据集根目录 `mmselfsup/data/tiny-imagenet-200` 下即可。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/mmselfsup/data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/open-mmlab/OpenMMLabCourse/main/codes/MMSelfSup_tutorials/anno_files/train.txt -P tiny-imagenet-200\n",
    "!wget https://raw.githubusercontent.com/open-mmlab/OpenMMLabCourse/main/codes/MMSelfSup_tutorials/anno_files/val.txt -P tiny-imagenet-200"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 写模型自监督预训练的配置文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. 新建一个名为 `simclr_resnet50_1xb32-coslr-1e_tinyin200.py` 的配置文件。（配置文件命名要求 & 含义可参考[这里](https://mmsegmentation.readthedocs.io/zh_CN/latest/tutorials/config.html#id3))\n",
    "\n",
    "\n",
    "\n",
    "2. `simclr_resnet50_1xb32-coslr-1e_tinyin200.py` 训练配置文件的内容：\n",
    "    1. 继承 [simclr_resnet50_8xb32-coslr-200e_in1k.py](https://github.com/open-mmlab/mmselfsup/blob/master/configs/selfsup/simclr/simclr_resnet50_8xb32-coslr-200e_in1k.py) 配置文件\n",
    "    2. 根据需求修改参数 samples_per_gpu（单个 GPU 的 Batch size）和 workers_per_gpu （单个 GPU 分配的数据加载线程数）\n",
    "    3. 修改数据集路径和数据标注文件路径\n",
    "    4. 根据 batch size 调整学习率（调整原则请参考：[这里](https://mmselfsup.readthedocs.io/zh_CN/latest/get_started.html#id2)）\n",
    "    5. 修改训练的总轮数 epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/mmselfsup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile /content/mmselfsup/configs/selfsup/simclr/simclr_resnet50_1xb32-coslr-1e_tinyin200.py\n",
    "\n",
    "_base_ = 'simclr_resnet50_8xb32-coslr-200e_in1k.py'\n",
    "\n",
    "# dataset\n",
    "data = dict(\n",
    "    samples_per_gpu=32, \n",
    "    workers_per_gpu=2,\n",
    "    train=dict(\n",
    "        data_source=dict(\n",
    "            data_prefix='data/tiny-imagenet-200/train',\n",
    "            ann_file='data/tiny-imagenet-200/train.txt',\n",
    "        )\n",
    "    )\n",
    ")\n",
    "\n",
    "# optimizer\n",
    "optimizer = dict(\n",
    "    lr=0.3 * ((32 * 1) / (32 * 8)),\n",
    ")\n",
    "\n",
    "runner = dict(max_epochs=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 模型自监督预训练"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们推荐使用分布式训练工具 [tools/dist_train.sh](https://github.com/open-mmlab/mmselfsup/blob/master/tools/dist_train.sh) 来启动训练任务（即使您只用一张 GPU 进行训练）。\n",
    "因为一些自监督预训练算法需要用多张 GPU 进行训练，为此 MMSelfSup 支持了多卡训练可能会用到的模块，如 `SyncBN` 等。如果算法在训练的过程中使用到了这些模块，但不使用分布式训练，就会报错。\n",
    "\n",
    "```shell\n",
    "bash tools/dist_train.sh ${CONFIG_FILE} ${GPUS} --work-dir ${YOUR_WORK_DIR} [optional arguments]\n",
    "```\n",
    "\n",
    "参数:\n",
    "+ CONFIG_FILE：自监督训练的配置文件所在路径\n",
    "\n",
    "+ GPUS：进行训练时所使用的 GPU 数量\n",
    "\n",
    "+ work-dir：训练过程中产生模型和日志等文件的保存路径\n",
    "\n",
    "其他可选参数 `optional arguments` 可参考[这里](https://mmselfsup.readthedocs.io/zh_CN/latest/get_started.html#id3)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/mmselfsup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash tools/dist_train.sh \\\n",
    "configs/selfsup/simclr/simclr_resnet50_1xb32-coslr-1e_tinyin200.py \\\n",
    "1 \\\n",
    "--work_dir work_dirs/selfsup/simclr/simclr_resnet50_1xb32-coslr-1e_tinyin200/ "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}