SSD-pytorch 训练过程全记录

@[toc]
修改内容适用于torch版本0.4.1，其余版本不保证。

准备工作

首先从github上下载源码: ssd.pytorch

git上提供了VGG-16模型的下载链接：vgg16_reducedfc

下载完成模型后，模型放到weights文件夹中。

数据集放入data文件夹中，可以自行下载数据集，也可以准备自己的数据集。

数据集处理

源码可以适应COCO和VOC两种数据格式。两种格式在文件夹中的目录结构如下：

如果使用的是自行准备的COCO数据集，需要自己创建coco_labels.txt。

将数据集准备好后，首先对data目录下的config.py文件进行修改。主要修改的是"num_classes"这一项，类别数量应为"你数据集中的类别数+1"，这里的1代表背景。

完成这一步后，进行对coco.py或voc0712.py的修改。

需要修改XXXX_ROOT，更改为你自己的数据集的位置，还有XXXX_CLASSES更改为你自己的类别。

此外voc0712.py额外需要修改的部分是class VOCDetection(data.Dataset), 初始化中使用了2002和2007两个image_sets，建议自行修改这段部分的代码。

建议确保每个图像中都有标记框，否则训练过程会出现报错。

img, boxes, labels = self.transform(img, target[:, :4], target[:, 4]) IndexError: too many indices for array

训练过程

train.py的修改

parser = argparse.ArgumentParser(description='Single Shot MultiBox Detector Training With Pytorch')

train_set = parser.add_mutually_exclusive_group()

parser.add_argument('--dataset', default='VOC', choices=['VOC', 'COCO'], type=str, help='VOC or COCO') parser.add_argument('--dataset_root', default=VOC_ROOT, help='Dataset root directory path') parser.add_argument('--basenet', default='vgg16_reducedfc.pth', help='Pretrained base model')

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

parser.add_argument('--resume', default=None, type=str, help='Checkpoint state_dict file to resume training from')

parser.add_argument('--start_iter', default=0, type=int,help='Resume training at this iter')

parser.add_argument('--num_workers', default=4, type=int, help='Number of workers used in dataloading')

parser.add_argument('--cuda', default=True, type=str2bool, help='Use CUDA to train model')

parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float,help='initial learning rate')

parser.add_argument('--momentum', default=0.9, type=float, help='Momentum value for optim')

parser.add_argument('--weight_decay', default=5e-4, type=float,help='Weight decay for SGD')

parser.add_argument('--gamma', default=0.1, type=float,help='Gamma update for SGD')

parser.add_argument('--visdom', default=False, type=str2bool,help='Use visdom for loss visualization')

parser.add_argument('--save_folder', default='weights/', help='Directory for saving checkpoint models')

args = parser.parse_args()

标红的部分为可能需要修改的部分，请自行注意。

如果出现未找到文件或者文件夹的报错，请自行溯源。

训练错误1：loss为NAN

如果lr设置过高，可能会导致训练过程中loss出现NAN的情况。

可以尝试把lr缩小一个数量级重新训练。

训练错误2：in layers/modules/multibox_loss.py In line 97, loss_c[pos]=0, the shape of mask [32,8732] ar index 0 not match the shape of the indexed tensor [270424,1] at index 0.

可以参考：Is there a bug in multibox_loss.py? · Issue #421 · amdegroot/ssd.pytorch (github.com)

其中对.item()的修改主要是版本问题。

训练错误3：StopIteration

可以参考：StopIteration · Issue #393 · amdegroot/ssd.pytorch (github.com)

作者代码本身存在问题。迭代器迭代完后再进行next的数据读取就会导致该报错。

改为：

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

验证过程

此外各种路径的位置也需要注意。建议在运行代码前都检查一遍。

annopath = os.path.join(args.voc_root, 'Annotations', '%s.xml')

imgpath = os.path.join(args.voc_root, 'JPEGImages', '%s.jpg')

imgsetpath = os.path.join(args.voc_root, 'ImageSets', 'Main', '{ : s }.txt')

YEAR = '2007'

devkit_path = args.voc_root + 'VOC' + YEAR

set_type = 'test'

这几行都要改成自己要用到数据的位置。'test'是存储了要用来验证的图片名字的"test.txt"的文件的名字，注意检查一下，我在验证的时候就犯了错误，文件里面是空的。

如果出现报错，重新执行前记得删除生成的annotations_cache/ 文件夹。

如果对coco数据集进行检验，可以参考：How to evaluate for coco dataset? · Issue #253 · amdegroot/ssd.pytorch (github.com)