在Win10和linux下跑通SSD目标检测算法，训练自己的数据集

本文链接：https://blog.csdn.net/skylake_/article/details/107314440

本文介绍了在Windows 10和Linux环境下，如何配置和运行基于PyTorch的SSD目标检测算法。首先，确保拥有正确的运行环境，如Python 3.7、CUDA 9.0+和PyTorch 1.0+。接着，从GitHub克隆SSD代码并进行必要的修改以适应SSD300和SSD500。详细步骤包括数据集的准备，如VOC数据集的处理和自定义数据集的格式转换。在训练过程中，针对不同数据集调整配置参数，并展示训练效果和实验结果。

摘要由CSDN通过智能技术生成

运行环境

python 3.7
cuda 9.0及以上
pytorch 1.0及以上

注意事项：

如果有GPU，请注意检查cuda的路径配置，nvcc -V 命令是否好使！
nvcc -V
查看pytorch的版本（推荐使用conda安装，避免与pip混用，容易造成环境冲突）：

imoort torch
print(torch.__version_)

查看cuda版本

import torch
print(torch.version.cuda)

在这里插入图片描述
环境正确是一切代码运行ok的基本条件。

下载pytorch版本SSD代码

SSD300，请从该链接clone代码：https://github.com/amdegroot/ssd.pytorch
这个代码存在问题，需要修改，首先是因为它使用的torch版本比较低，所以有些代码不兼容，需要更改，其次是这个代码是SSD300的代码，不支持512的裁剪，因此需要进行扩展。
支持SSD500的代码，链接如下：https://github.com/midasklr/SSD.Pytorch这里代码同样有一些错误，直接运行会报错，这里说明如何进行修改，然后训练自己的数据集。

下载好了之后，需要先跑通SSD300的代码，其中train.py中的代码如下：

from data import *
from utils.augmentations import SSDAugmentation
from layers.modules import MultiBoxLoss
from ssd import build_ssd
import os
import sys
import time
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torch.nn.init as init
import torch.utils.data as data
import numpy as np
import argparse
import visdom


viz = visdom.Visdom()  # 可视化控件


def str2bool(v):
    return v.lower() in ("yes", "true", "t", "1")


parser = argparse.ArgumentParser(description='Single Shot MultiBox Detector Training With Pytorch')
train_set = parser.add_mutually_exclusive_group()
parser.add_argument('--dataset', default='FACEAI', choices=['VOC', 'COCO', 'FACEAI'], type=str, help='VOC or COCO')
parser.add_argument('--dataset_root', default=FACEAI_ROOT, help='Dataset root directory path')
parser.add_argument('--basenet', default='vgg16_reducedfc.pth', help='Pretrained base model')
parser.add_argument('--batch_size', default=1, 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=True, 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()


if torch.cuda.is_available():
    if args.cuda:
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
    if not args.cuda:
        print("WARNING: It looks like you have a CUDA device, but aren't " +
              "using CUDA.\nRun with --cuda for optimal training speed.")
        torch.set_default_tensor_type('torch.FloatTensor')
else:
    torch.set_default_tensor_type('torch.FloatTensor')

if not os.path.exists(args.save_folder):
    os.mkdir(args.save_folder)


def train():
    if args.dataset == 'COCO':
        pass
        # only use VOC0712
        # if args.dataset_root == VOC_ROOT:
        #     if not os.path.exists(COCO_ROOT):
        #         parser.error('Must specify dataset_root if specifying dataset')
        #     print("WARNING: Using default COCO dataset_root because " + "--dataset_root was not specified.")
        #     args.dataset_root = COCO_ROOT
        # cfg = coco
        # dataset = COCODetection(root=args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], MEANS))
    elif args.dataset == 'VOC':
        # if args.dataset_root == COCO_ROOT:
        #     parser.error('Must specify dataset if specifying dataset_root')
        cfg = voc  # cfg为配置文件，包含了指定数据集的一些信息
        dataset = VOCDetection(root=args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], MEANS))
    elif args.dataset == 'FACEAI':
        cfg = faceai  # cfg为配置文件，包含了指定数据集的一些信息
        dataset = FACEAIDetection(args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], MEANS))

    ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
    net = ssd_net

    if args.cuda:
        net = torch.nn.DataParallel(ssd_net)
        cudnn.benchmark = True

    if args.resume:
        print('Resuming training, loading {}...'.format(args.resume))
        ssd_net.load_weights(args.resume)
    else:
        vgg_weights = torch.load(args.save_folder + args.basenet)
        print('Loading base network...')
        ssd_net.vgg.load_state_dict(vgg_weights)

    if args.cuda:
        net = net.cuda()

    # 如果resume为None
    if not args.resume:
        print('Initializing weights...')
        # initialize newly added layers' weights with xavier method
        ssd_net.extras.apply(weights_init)
        ssd_net.loc.apply(weights_init)
        ssd_net.conf.apply(weights_init)

    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    # SSD网络之MultiBoxLoss层, 属于自定义损失
    criterion