快速使用Mobilenet SSD进行训练VOC格式的数据集

1.训练步骤

STEP1:下载代码并配置环境

• 下载代码

git clone https://github.com/lufficc/SSD.git
cd SSD

• 修改requirements.txt

torch==1.5
torchvision==0.5
yacs
tqdm
opencv-python
vizer

• 根据requirements.txt完成环境配置

# Required packages: torch torchvision yacs tqdm opencv-python vizer
conda create -n ssd-lufficc python=3.8
pip install -r requirements.txt

• 额外补充安装

pip install tensorboardX
pip install six

STEP2:在当前目录下建立数据集文件夹的软连接或者复制数据集到当前文件夹

• 完成后文件夹结构

├──VOC_data/VOC2007/
    ├── Annotations #可以采用软连接的方式避免对大量数据进行复制
        ├──放置xml文件 #TODO
    ├── JPEGImages
        ├──放置img文件 #TODO
    ├──ImageSets/Main #可以采用软连接的方式避免对大量数据进行复制
        ├── split.py #数据分割脚本,用于生成训练索引文件

• split.py文件内容

import os
import random 
random.seed(0)

xmlfilepath=r'./Annotations'
saveBasePath=r"./ImageSets/Main/"
 
trainval_percent=0.9 # 可以自己设置
train_percent=0.9 # 可以自己设置

temp_xml = os.listdir(xmlfilepath)
total_xml = []
for xml in temp_xml:
    if xml.endswith(".xml"):
        total_xml.append(xml)

num=len(total_xml)  
list=range(num)  
tv=int(num*trainval_percent)  
tr=int(tv*train_percent)  
trainval= random.sample(list,tv)  
train=random.sample(trainval,tr)  
 
print("train and val size",tv)
print("traub suze",tr)
ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')  
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')  
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')  
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')  
 
for i  in list:  
    name=total_xml[i][:-4]+'\n'  
    if i in trainval:  
        ftrainval.write(name)  
        if i in train:  
            ftrain.write(name)  
        else:  
            fval.write(name)  
    else:  
        ftest.write(name)  
  
ftrainval.close()  
ftrain.close()  
fval.close()  
ftest .close()

python split.py

• 配置文件根目录的环境变量

export VOC_ROOT="./VOC_data"  #/path/to/voc_root

STEP3:修改配置文件

 vim configs/mobilenet_v2_ssd320_voc0712.yaml

MODEL:
  NUM_CLASSES: 11  # 修改NUM_CLASSES
  BOX_HEAD:
    PREDICTOR: 'SSDLiteBoxPredictor'
  BACKBONE:
    NAME: 'mobilenet_v2'
    OUT_CHANNELS: (96, 1280, 512, 256, 256, 64)
  PRIORS:
    FEATURE_MAPS: [20, 10, 5, 3, 2, 1]
    STRIDES: [16, 32, 64, 107, 160, 320]
    MIN_SIZES: [60, 105, 150, 195, 240, 285]
    MAX_SIZES: [105, 150, 195, 240, 285, 330]
    ASPECT_RATIOS: [[2, 3], [2, 3], [2, 3], [2, 3], [2, 3], [m2, 3]]
    BOXES_PER_LOCATION: [6, 6, 6, 6, 6, 6]
INPUT:
  IMAGE_SIZE: 320
DATASETS:
  TRAIN: ("voc_2007_trainval", "voc_2012_trainval")  # TODO
  TEST: ("voc_2007_test", )
SOLVER:
  MAX_ITER: 120000
  LR_STEPS: [80000, 100000]
  GAMMA: 0.1
  BATCH_SIZE: 32
  LR: 1e-3

OUTPUT_DIR: 'outputs/mobilenet_v2_ssd320_voc0712

STEP4:修改类别信息

vim ssd/data/datasets/voc.py

class VOCDataset(torch.utils.data.Dataset):
    class_names = ('person','bridgevehicle','luggagevehicle','plane','refuelvehicle','foodvehicle','rubbishvehicle','watervehicle','platformvehicle','tractorvehicle','bridgeconnector') # 改成自己的class注意，类名必须小写
    ······

STEP5: 模型训练

• 修改默认所用的device和batch_size

 vim ssd/config/defaults.py

# 修改以下内容
_C.MODEL.DEVICE = "cuda:1" # cpu/cuda/cuda:1
_C.SOLVER.BATCH_SIZE = 128
_C.TEST.BATCH_SIZE = 32

• 单GPU

# for example, train SSD300:
python train.py --config-file configs/mobilenet_v2_ssd320_voc0712.yaml

• 多GPU

# for example, train SSD300 with 2 GPUs:
export NGPUS=2
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py --config-file configs/mobilenet_v2_ssd320_voc0712.yaml SOLVER.WARMUP_FACTOR 0.03333 SOLVER.WARMUP_ITERS 1000

2.模型评估

• 单GPU

# for example, evaluate SSD300:
python test.py --config-file configs/mobilenet_v2_ssd320_voc0712.yaml

• 多GPU

# for example, evaluate SSD300 with 2 GPUs:
export NGPUS=2
python -m torch.distributed.launch --nproc_per_node=$NGPUS test.py --config-file configs/mobilenet_v2_ssd320_voc0712.yaml

3.模型调用(预测)

#TODO

参考资料

1. https://github.com/qfgaohao/pytorch-ssd