本章节既是指南如何训练测试模型,同时也是检验上一章所提的平台搭建是否成功。

目录

一、下载预训练模型
1、在ssd_pytorch文件夹下新建weights文件夹

cd ..
mkdir weights
cd weights

2、下载预训练模型

wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

上面这个官方的下载指令下载速度巨慢,我用了整整一个半小时才下载下来,嫌慢的小伙伴可以用下面的链接下载。
网盘链接:https://pan.baidu.com/s/1WzKPph3Xf0K6G3Tqc_o0fQ
提取码:gvh6
二、修改multibox_loss.py代码
原本的train.py直接运行会报错,所以需要进行修改。
0、可能出现的问题(我遇到的)

RuntimeError: The shape of the mask [2, 8732] at index 0 does not match the shape of the indexed tensor [17464, 1] at index 0
RuntimeError: Expected object of type torch.cuda.FloatTensor but found type torch.cuda.LongTensor for argument #3 'other’

(更多错误与解决方法,参考:ssd.pytorch源码迁移实现(pytorch版本从0.2迁移到0.4(1.0)))

为解决上述两个问题,需依次对ssd_pytorch/layers/modules/multibox_loss.py进行修改

1、找到代码的第97行

pytorch vgg16预训练 pytorch预训练模型下载_迭代


交换两行代码的顺序就可以解决第一个error

pytorch vgg16预训练 pytorch预训练模型下载_平台搭建_02


2、找到代码的最后一部分

原本是

N = num_pos.data.sum()
loss_l /= N
loss_c /= N
return loss_l, loss_c

中间加上一句,改为

N = num_pos.data.sum()
N=N.type(torch.cuda.FloatTensor)
loss_l /= N
loss_c /= N
return loss_l, loss_c

就可以解决第二个error
三、修改迭代次数
本章节主要是为了指导如何搭建平台,代码能否成功运行为判断依据,因此需要降低训练用时。
打开ssd_pytorch/data/config.py找到第17行,将迭代次数(max_iter)从120000改为200(建议迭代次数不要太小,太小会查不出Stoplteration异常我用的是batch_size=32,max_iter=200的组合)。

# SSD300 CONFIGS
voc = {
    'num_classes': 21,
    'lr_steps': (80000, 100000, 120000),
    'max_iter': 120000,                                                                  #迭代次数
    'feature_maps': [38, 19, 10, 5, 3, 1],
    'min_dim': 300,
    'steps': [8, 16, 32, 64, 100, 300],
    'min_sizes': [30, 60, 111, 162, 213, 264],
    'max_sizes': [60, 111, 162, 213, 264, 315],
    'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
    'variance': [0.1, 0.2],
    'clip': True,
    'name': 'VOC',                                                                         
}

四、修改train.py代码
这一步骤,按照个人情况自行选择是否操作。
1、指定GPU
在train.py代码头部加一行指令(注:加在import os之后),“1”代表选择了1号GPU

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

2、修改batch_size
官方默认的batch_size为32,如果GPU空间不够可以调小,在train.py代码的第32行进行修改

parser.add_argument('--batch_size', default=4, type=int,
                    help='Batch size for training')

3、在迭代过程中,每次执行batch张图片,通过images, targets = next(batch_iterator)读取图片时,如果next()中没有数据后会触发Stoplteration异常,使用下面语句替换 images, targets = next(batch_iterator)将解决这种异常问题。

while True:
    try:
        # 获得下一个值:
        images, targets = next(batch_iterator)
 
    except StopIteration:
        # 遇到StopIteration就退出循环
        break

五、训练数据集

cd ..
python train.py

六、修改eval.py代码
1、指定GPU
在eval.py代码头部加一行指令(注:加在import os之后),“1”代表选择了1号GPU

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

(注:这一小步可不做)
2、原始代码中默认的是读取官方的训练模型(注:无法下载),而前一步中,我们已经生成了自己的VOC.pth模型,因此需要将其修改。
  找到代码的低38行,将ssd300_mAP_77.43_v2.pth改为VOC.pth

parser.add_argument('--trained_model',
                    default='weights/ssd300_mAP_77.43_v2.pth', type=str,
                    help='Trained state_dict file path to open')

七、测试效果
1、直接运行代码。

python eval.py

2、因为我只训练了200次,所以效果很差。不过运行成功还是说明了平台搭建没有问题。

pytorch vgg16预训练 pytorch预训练模型下载_数据集_03