SSD-pytorch 训练过程全记录
SSD-pytorch 训练过程全记录
文章目录
- SSD-pytorch 训练过程全记录
-
- 准备工作
- 数据集处理
- 训练过程
-
- train.py的修改
- 训练错误1：loss为NAN
- 训练错误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.
- 训练错误3：StopIteration
- 验证过程
修改内容适用于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)
验证过程
eval.py中代码主要是针对VOC数据的，COCO数据使用的代码在末尾提供。
eval.py中需要修改的部分比较少，首先是要验证的模型的位置。
parser.add_argument(’–trained_model’, default=‘weights/ssd300_COCO_15000.pth’, type=str,help=‘Trained state_dict file path to open’)
此外各种路径的位置也需要注意。建议在运行代码前都检查一遍。
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)