前面我們準備好了COCO格式的數據集：將YOLO格式的數據集轉換為mmdetection格式-CSDN博客https://blog.csdn.net/qq_54708219/article/details/148224187?spm=1001.2014.3001.5501

下面我們使用MMdetection開始訓練。

1.創建新的數據集類

首先，在mmdet/datasets里，創建屬于自己的數據集文件。因為數據集是coco格式，可直接復制coco.py，文件名更改為Dataset_depth.py，然后把里面的標簽改掉。我們需要把coco.py中的類名、數據集中的類別和標注的顏色換成自己的數據集的類名。這里，類名我修改為BoxDataset，數據集只有兩個類別：'shallow_box'和'shallow_half_box'，下面的palette是調色板，保持和classes長度一樣即可，里面的顏色可以填成自己喜歡的，這里是直接保留了前兩個后面全部刪除：

2.錄入數據集

在mmdet/datasets/init.py中仿照coco的格式把剛才創建的數據集類進行錄入：

注意：如果注冊完成訓練仍然報KeyError，看一下報錯的文件路徑，可能是導入數據集類的時候一直是conda虛擬環境上的，而不是自己修改后的。解決方案就是順著報錯路徑（conda環境路徑）重新注冊一遍數據集：

3.生成完整的配置文件

在configs文件夾下選擇你想訓練的模型，例如，我這里選擇configs/deformable_detr/deformable-detr_r50_16xb2-50e_coco.py，然后在終端運行以下命令：

python tools/train.py <你所選擇的模型路徑>

python tools/train.py configs/deformable_detr/deformable-detr_r50_16xb2-50e_coco.py

此時運行會報錯，因為他會自動尋找coco數據集而不是自己的（我自己沒有下載coco數據集）：

但是在當前目錄下會生成work_dir文件夾，里面包含了完整的模型配置文件work_dirs/deformable-detr_r50_16xb2-50e_coco/deformable-detr_r50_16xb2-50e_coco.py：

4.修改生成的配置文件

接下來對work_dirs/deformable-detr_r50_16xb2-50e_coco/deformable-detr_r50_16xb2-50e_coco.py進行修改：

（1）數據集位置和數據集類：修改data_root（數據集相對路徑）和dataset_type（之前注冊的數據集類名）：

同時在train_dataloader，test_dataloader，val_dataloader中修改data_root和type：

同時還修改一下val_evaluator和test_evaluator的ann_file：

（2）修改數據集類別數：ctrl+F召喚出搜索框，輸入num_classes，把coco數據集的80類改為自己的數據集類別個數（我的是2類）：

（3）修改batch_size

ctrl+F召喚出搜索框，輸入batch_size，根據顯卡算力設置train_dataloader，test_dataloader和val_dataloader的批量大小：

（4）注冊數據集的metainfo

參考第一步，在配置文件中注冊以下信息：

metainfo = dict(classes=('shallow_box', 'shallow_half_box'),palette=[(220, 20, 60), (119, 11, 32)]
)

然后在train_dataloader，test_dataloader，val_dataloader中都寫入：

（5）修改訓練輪數：ctrl+F召喚出搜索框，輸入max_epochs，這個模型默認的訓練輪數是50

5.訓練

然后把3中的命令改為修改后的配置文件路徑：

python tools/train.py <你所修改的模型配置文件路徑>

python tools/train.py work_dirs/deformable-detr_r50_16xb2-50e_coco/deformable-detr_r50_16xb2-50e_coco.py

看到以下界面說明正常運行：

這個時候模型也會自動保存每一輪的權重文件：