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

jupiter
2021-11-27 / 0 评论 / 846 阅读 / 正在检测是否收录...
温馨提示:
本文最后更新于2021年12月07日,已超过1064天没有更新,若内容或图片失效,请留言反馈。

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
0

评论 (0)

打卡
取消