SSD-pytorch训练

源码：https://github.com/amdegroot/ssd.pytorch

1、根据readme下载好所需资料放到指定文件夹
2、设置模型参数：
在train.py中找到

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

设置批大小，根据自己显卡能力设置，默认为32，这里我改为8
3、修改部分源码

• 在train.py中

（1）将第165行

images, targets = next(batch_iterator)

改为

try:
    images, targets = next(batch_iterator)
except StopIteration as e:
    batch_iterator = iter(data_loader)
    images, targets = next(batch_iterator)

（2）将第183行、184行、188行

loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]

中的data[0]改为item（）

• 在mutibox_loss.py中

（1）将第97行和98行交换顺序

loss_c = loss_c.view(num, -1)
loss_c[pos] = 0 

（2）将第114行

N = num_pos.data.sum()

改为

N = num_pos.data.sum().double()

4、开始训练即可

主要是记录一下自己平时遇到的问题，和大家分享一下
如有侵犯，请联系我
点个赞支持一下吧