mmdetection训练自己的数据集

一、mmdetection下载

github：open-mmlab/mmdetection: OpenMMLab Detection Toolbox and Benchmark (github.com)

二、mmdetection环境配置

1. 创建一个新环境：

conda create -n mmd python==3.9
activate mmd

2.查看电脑配置:

nvidia-smi

可以看到驱动最高支持的cuda版本为12.2, 电脑已经安装好cuda toolkit11.6:

3.安装torch和cudnn等

根据需求选择安装torch版本是'1.11.0+cu113’，也可以在虚拟环境中安装cudatoolkit和cudnn,以及安装torch1.11GPU版本

pytorch v1.11.0PyTorch官网历史版本链接（ https://pytorch.org/get-started/previous-versions/）里面，找到v1.11.0的安装命令， 此处选择的是cuda11.3对应的PyTorch安装命令。

# CUDA 11.3

pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113

测试： 

python
>>import torch
>>torch.__version__
>>torch.cuda.is_available()

可以看到torch对应的版本

3.环境编译:

cd mmdetection
pip install -r requirements/build.txt
pip install -v -e .  # or "python s

4.安装mmcv:

pip install mmcv
#也可以选择 pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/{cu_version}/{torch_version}/index.html
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.6.0/index.html

 

三、数据集的准备

数据的格式参考yolov3或者maskrcnn的coco格式
目录格式如下:

mmdetection
----data
--------coco_my（自己的数据集命名）
-------------annotations
-----------------instance_train.json
-----------------instance_val.json
-----------------instance_test.json
-------------train（图片）
-------------val（图片）
-------------test（图片）

一般train,val,test按比例分开然后在config文件里面修改数据集格式

#修改数据集相关设置
img_scale = (640, 640)  # width, height
dataset_type = 'COCODataset'
classes = ('person','bicycle','car','UAV',...) #修改为自己数据集的名称
data = dict(
    train=dict(
        img_prefix=r'E:\mmdetection-main\data\coco_4classes\train2017',
        classes=classes,
        ann_file=r'E:\mmdetection-main\datacoco_4classes\coco\annotations\instances_train2017.json'),
    val=dict(
        img_prefix=r'E:\mmdetection-main\data\coco_4classes\val2017',
        classes=classes,
        ann_file=r'E:\mmdetection-main\data\coco_4classes\annotations\instances_val2017.json'),
    test=dict(
        img_prefix=r'E:\mmdetection-main\data\coco_4classes\test2017',
        classes=classes,
        ann_file=r'E:\mmdetection-main\data\coco_4classes\annotations\instances_test2017.json'))

四、使用mmdetection训练

本文以RetinaNet为例，首先打开configs中的centernet目录下的网络文件，将base中的coco注释掉修改为数据集准备所示，下面的dataset_type和data_root也需要进行修改，与目录对应，num_classes需要修改为4，这与自制数据集类别保持一致。

关于配置文件的学习可以看官方文档：学习配置文件 — MMDetection 3.0.0 文档

下面打开Terminal输入指令即可训练，如果代码没有报错并出现了第二张图中损失函数的变化说明模型正在训练：

python tools/train.py configs/retinanet/retinanet_r18_fpn_1x_my.py --work-dir work_dirs/retinanet

成功运行：