运行环境
python 3.7
cuda 9.0及以上
pytorch 1.0及以上
注意事项:
如果有GPU,请注意检查cuda的路径配置,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