codes/MMSelfSup_tutorials/【6】在 MMDetection 中使用自监督预训练模型.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/mmselfsup_logo.png\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 在 MMDetection 中使用自监督预训练模型\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/open-mmlab/OpenMMLabCourse/blob/main/codes/MMSelfSup_tutorials/【6】在%20MMDetection%20中使用自监督预训练模型.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**MMDetection Repo**：[https://github.com/open-mmlab/mmdetection](https://github.com/open-mmlab/mmdetection)\n",
    "\n",
    "**MMDetection 官方文档链接**：[https://mmdetection.readthedocs.io/en/latest](https://mmdetection.readthedocs.io/en/latest)\n",
    "\n",
    "**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",
    "\n",
    "\n",
    "**作者**：OpenMMLab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. 任务介绍"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们除了可以直接使用 MMSelfSup 里 `tools/benchmark` 底下的工具对自监督预训练模型进行评估，还可以直接将保存好的模型文件用于 OpenMMLab 的其他库（如 MMDetection、MMSegmentation）中进行训练。\n",
    "\n",
    "本教程将演示：如何在 MMDetection 中使用自监督预训练模型的权值进行目标检测任务的训练。\n",
    "\n",
    "**注意：下游任务的 backbone 必须和自监督预训练模型的 backbone 保持一致，才能进行自监督预训练权值的迁移！**"
   ]
  },
  {
   "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 安装 MMDetection 的依赖库：MMCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U openmim\n",
    "!mim install mmcv-full"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3  安装 MMDetection\n",
    "其他方式的安装详情请参考： [MMDetection 安装文档](https://mmdetection.readthedocs.io/en/latest/get_started.html#best-practices)。"
   ]
  },
  {
   "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/mmdetection.git\n",
    "%cd mmdetection\n",
    "!pip install -v -e ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.5 检查安装是否正确"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mmdet\n",
    "print(mmdet.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 准备数据集"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.0 数据集介绍"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本教程将在 `COCO2017` 数据集上训练目标检测任务。\n",
    "\n",
    "COCO 是一个大规模目标检测、图像分割和图像描述数据集。它包含 80 个物体类别，150 万个物体实例，约 33 万张图像（其中超过 20 万张图像有标注）。\n",
    "\n",
    "数据集官方网址：https://cocodataset.org"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 下载并解压数据集"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "使用该命令可以下载并解压 COCO 数据集。该命令支持下载 COCO，VOC 和 LVIS 数据集。\n",
    "```shell\n",
    "python tools/misc/download_dataset.py --dataset-name ${DATASET_NAME} --unzip --delete\n",
    "```\n",
    "\n",
    "参数：\n",
    "+ DATASET_NAME：支持数据集的名称 `coco2017`，`voc2007` 和 `lvis`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/mmdetection\n",
    "!mkdir -p data/coco\n",
    "%cd data/coco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget http://images.cocodataset.org/zips/train2017.zip\n",
    "!unzip train2017.zip\n",
    "!rm -rf train2017.zip\n",
    "\n",
    "!wget http://images.cocodataset.org/zips/val2017.zip\n",
    "!unzip val2017.zip\n",
    "!rm -rf val2017.zip\n",
    "\n",
    "!wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip\n",
    "!unzip annotations_trainval2017.zip\n",
    "!rm -rf annotations_trainval2017.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 准备自监督预训练模型的 backbone 权值文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 针对自监督预训练过程中保存的 checkpoint 文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以参考前面的教程，使用命令 `tools/model_converters/extract_backbone_weights.py` 来提取自监督预训练模型的 backbone 权值，代码如下：\n",
    "\n",
    "```python\n",
    "python tools/model_converters/extract_backbone_weights.py {CHECKPOINT} {MODEL_FILE}\n",
    "```\n",
    "\n",
    "参数:\n",
    "- CHECKPOINT：自监督预训练过程中保存下来（名为 `epoch_*.pth`）的模型文件路径\n",
    "- MODEL_FILE：输出 backbone 权重文件的保存路径。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 针对 MMSelfSup 模型库里的模型文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**注意：MMSelfSup 的 [模型库](https://github.com/open-mmlab/mmselfsup/blob/master/docs/en/model_zoo.md) 中的模型文件都已经提取过 backbone 权值，不需要再次提取！**我们直接使用即可。\n",
    "\n",
    "在模型库中找到在 SimCLR 的预训练模型文件 `simclr_resnet50_8xb32-coslr-200e_in1k`，下载放在 `checkpoints` 文件夹里 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/mmdetection\n",
    "!mkdir checkpoints\n",
    "!wget https://download.openmmlab.com/mmselfsup/simclr/simclr_resnet50_8xb32-coslr-200e_in1k_20220428-46ef6bb9.pth -P checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 使用自监督预训练模型的权值初始化目标检测的模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 写目标检测任务的配置文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. 新建一个名为 `faster_rcnn_simclr-pretrained_r50_caffe_fpn_9k_coco.py` 的配置文件。\n",
    "\n",
    "\n",
    "2. `faster_rcnn_simclr-pretrained_r50_caffe_fpn_9k_coco.py` 训练配置文件的内容：\n",
    "    1. 继承 [faster_rcnn_r50_caffe_fpn_90k_coco.py](https://github.com/open-mmlab/mmdetection/blob/master/configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_90k_coco.py) 配置文件\n",
    "    2. 将模型的 checkpoint 参数修改为自监督预训练模型的 backbone 所在路径\n",
    "    3. 根据 batch size 调整学习率（调整原则请参考：[这里](https://mmselfsup.readthedocs.io/zh_CN/latest/get_started.html#id2)）\n",
    "    4. 修改总共训练的迭代次数 max_iters、模型训练多少个 iteration 评估一次、模型训练多少个 iteration 保存一次 checkpoint 文件等参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile /content/mmdetection/configs/faster_rcnn/faster_rcnn_simclr-pretrained_r50_caffe_fpn_9k_coco.py\n",
    "_base_ = 'faster_rcnn_r50_caffe_fpn_90k_coco.py'\n",
    "\n",
    "model = dict(\n",
    "    backbone=dict(\n",
    "        frozen_stages=-1,\n",
    "        init_cfg=dict(\n",
    "            type='Pretrained',\n",
    "            checkpoint='checkpoints/simclr_resnet50_8xb32-coslr-200e_in1k_20220428-46ef6bb9.pth')\n",
    "    )\n",
    ")\n",
    "\n",
    "# optimizer\n",
    "optimizer = dict(\n",
    "    lr=0.02 * (1 / 8)\n",
    ")\n",
    "\n",
    "# Runner type\n",
    "runner = dict(_delete_=True, type='IterBasedRunner', max_iters=9000)\n",
    "\n",
    "checkpoint_config = dict(interval=3000)\n",
    "\n",
    "evaluation = dict(interval=3000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 开始训练目标检测模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在训练目标检测模型时，我们使用 [tools/train.py](https://github.com/open-mmlab/mmdetection/blob/master/tools/train.py) 训练工具来启动训练。\n",
    "\n",
    "```shell\n",
    "python tools/train.py \\\n",
    "    ${CONFIG_FILE} \\\n",
    "    [optional arguments]\n",
    "```\n",
    "\n",
    "参数：\n",
    "- CONFIG_FILE：“检测” 评估所使用的配置文件，位于 configs 里对应模型的目录下\n",
    "\n",
    "详情请参考：[文档](https://mmdetection.readthedocs.io/en/latest/1_exist_data_model.html#training-on-a-single-gpu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/mmdetection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python tools/train.py \\\n",
    "configs/faster_rcnn/faster_rcnn_simclr-pretrained_r50_caffe_fpn_9k_coco.py"
   ]
  }
 ],
 "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
}