def detect():
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# load net
# cfg = config.voc_af
input_size = [416, 416]
num_classes = len(VOC_CLASSES)
testset = VOCDetection(
args.voc_root, [('2007', 'test')],
BaseTransform(input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)), VOCAnnotationTransform())
# build model
if args.version == 'yolo':
from model.yolo import YOLO
net = YOLO(device,
input_size=input_size,
num_classes=num_classes,
trainable=False)
print('Let us test yolo on the VOC0712 dataset ......')
else:
print('Unknown Version !!!')
exit()
net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.eval()
print('Finished loading model!')
net = net.to(device)
# evaluation
detect_net(net,
device,
testset,
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)),
thresh=args.visual_threshold)
def showData():
dataDir = r'/mnt/hdisk1/trainSet/pytorch/antifollow'
trainData = VOCDetection(root=dataDir,
image_sets=[('2012', 'trainval')],
transform=SSDAugmentationTest(),
target_transform=VOCAnnotationTransform(
None, True),
dataset_name='VOC0712')
epoch_size = len(trainData)
data_loader = torchData.DataLoader(trainData,
2,
num_workers=1,
shuffle=False,
collate_fn=getNumpy_collate,
pin_memory=True)
batch_iterator = iter(data_loader)
picCount = 0
while (True):
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
break
#show img
for i in range(len(images)):
img = images[i]
imWidth = img.shape[1]
imHeight = img.shape[0]
imgColor = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
for rect in targets[i]:
cv2.rectangle(
imgColor,
(int(rect[0] * imWidth), int(rect[1] * imHeight)),
(int(rect[2] * imWidth), int(rect[3] * imHeight)),
VOC_COLOR_ID_MAP[int(rect[4])])
picCount += 1
print(picCount)
cv2.imshow("input", imgColor)
cv2.waitKey()
def init_dataset(self, init_dataloader=True):
if self.dataset == 'COCO':
raise NotImplementedError()
elif self.dataset == 'VOC':
# raise NotImplementedError()
self.cfg = voc
dataset = VOCDetection(root=VOC_ROOT, transform=SSDAugmentation(self.cfg['min_dim'], MEANS))
elif self.dataset == 'helmet':
self.cfg = helmet
dataset = HelmetDetection(root=HELMET_ROOT, transform=SSDAugmentation(self.cfg['min_dim'], MEANS))
else:
raise RuntimeError()
if not init_dataloader:
return dataset
data_loader = DataLoader(dataset, self.batch_size,
num_workers=self.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
return dataset, data_loader
from layers.functions import PriorBox
import torch
import cv2
# cfg = VOC_300
# cfg = VOC_300_3
cfg = VEHICLE_240
# train_sets = [('VehiclePersonV2_00', 'trainval'),
# ('VOC2007', 'trainval'),
# ('ShenzhenClean', 'trainval'),
# #('negative','trainval'),
# ]
train_sets = [('ShenzhenClean', 'trainval')]
dataset = VOCDetection(VOCroot, train_sets, preproc(
cfg['min_dim'], rgb_means, p), AnnotationTransform())
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
num_priors = (priors.size(0))
threshold = 0.5
variance = [0.1, 0.2]
anchor_match = 'NORMAL'
for index in range(len(dataset)):
i = np.random.randint(0, len(dataset))
img, targets = dataset[i]
image = dataset.pull_image(i)
h,w,_ = image.shape
do_python_eval(use_07=False)
if __name__ == '__main__':
# 加载网络
num_classes = len(labelmap) + 1
net = build_ssd('test', 300, num_classes)
net.load_state_dict(torch.load(args.trained_model))
if torch.cuda.is_available():
net = net.cuda()
cudnn.benchmark = True
net.eval()
# 记载数据
dataset = VOCDetection(args.voc_root, [('2007', 'test')],
BaseTransform(300, (104, 117, 123)),
VOCAnnotationTransform())
# 测试结果
num_images = len(dataset)
all_boxes = [[[] for _ in range(num_images)]
for _ in range(len(labelmap) + 1)] # 用于保存所有符合条件的结果
with torch.no_grad():
for i in range(num_images):
img, gt, h, w = dataset.pull_item(i)
img = img.unsqueeze(0)
if torch.cuda.is_available():
img = img.cuda()
detections = net(img) # 得到结果,shape[1,21,200,5]
for j in range(1, detections.shape[1]): # 循环计算每个类别
print('Evaluating detections')
testset.evaluate_detections(all_boxes, save_folder)
if __name__ == '__main__':
args = arg_parse()
weightsfile = args.weights
nms_thresh = args.nms_thresh
input_wh = args.input_wh
cuda = args.cuda
use_pad = args.use_pad
save_folder = args.save_folder
dataset = args.dataset
if dataset[0] == "V":
cfg = voc_config
test_dataset = VOCDetection(cfg["root"], datasets_dict["VOC2007"],
input_wh)
elif dataset[0] == "C":
cfg = coco_config
test_dataset = COCODetection(cfg["root"], datasets_dict["COCOval"],
input_wh)
else:
print("only support VOC and COCO datasets !!!")
print("load test_dataset successfully.....")
with open(cfg["name_path"], "r") as f:
classes = [i.strip() for i in f.readlines()]
net = Yolov3("test", input_wh, cfg["anchors"], cfg["anchors_mask"],
cfg["num_classes"])
state_dict = torch.load(weightsfile)
type=int,
help='threshold for the number of prior boxes.')
parser.add_argument('--gpus',
default='1',
type=str,
help='visible devices for CUDA')
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
dataset = None
if args.dataset == 'VOC':
config = voc
rt = VOC_ROOT if args.dataset_root is None else args.dataset_root
dataset = VOCDetection(rt, transform=BaseTransform(300, (104, 117, 123)))
elif args.dataset == 'COCO18':
config = coco18
rt = COCO_ROOT if args.dataset_root is None else args.dataset_root
dataset = COCODetection(rt,
transform=BaseTransform(300, (104, 117, 123)),
target_transform=COCOAnnotationTransform('COCO18'))
elif args.dataset == 'helmet':
config = helmet
rt = HELMET_ROOT if args.dataset_root is None else args.dataset_root
dataset = HelmetDetection(rt,
transform=BaseTransform(300, (104, 117, 123)))
else:
raise NotImplementedError()
config = generic
def train():
if args.dataset == 'COCO18':
# cfg = coco18
cfg = vococo
rt = args.dataset_root or COCO_ROOT
dataset = COCODetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS),
target_transform=COCOAnnotationTransform('COCO18'))
elif args.dataset == 'COCO':
cfg = coco
# cfg = vococo
rt = args.dataset_root or COCO_ROOT
dataset = COCODetection(root=rt, image_sets=(('2017', 'train'),),
transform=SSDAugmentation(cfg['min_dim'], MEANS))
elif args.dataset == 'VOC':
cfg = voc
# cfg = coco_on_voc
rt = args.dataset_root or VOC_ROOT
dataset = VOCDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS))
elif args.dataset == 'BCCD':
cfg = bccd
rt = args.dataset_root or BCCD_ROOT
dataset = BCCDDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS))
elif args.dataset == 'SHWD':
cfg = shwd
rt = args.dataset_root or SHWD_ROOT
dataset = SHWDDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS))
elif args.dataset == 'helmet':
cfg = helmet
rt = args.dataset_root or HELMET_ROOT
dataset = HelmetDetection(root=rt, transform=SSDAugmentation(cfg['min_dim'], MEANS))
else:
raise RuntimeError()
if args.custom_priors is not None:
apt = IOAdapterSSD(cfg, 'test')
apt.load(*torch.load(args.custom_priors))
custom_priors = apt.fit_output(apt.msks[0])
print('num_boxes = %d ' % custom_priors.size()[0])
custom_mbox = None
# params = torch.load(args.custom_priors)
# # bbox = gen_priors(params, args.prior_types, cfg)
# gen = AdaptivePriorBox(cfg, phase='test')
# custom_priors = gen.forward(params)
# custom_mbox = [p.size(0) for p in params]
if args.cuda:
custom_priors = custom_priors.cuda()
ssd_net = build_ssd('train', cfg, custom_mbox, custom_priors)
else:
# priors = torch.load('anchors/voc_baseline.pth')
# if args.cuda:
# priors = priors.cuda()
ssd_net = build_ssd('train', cfg)
net = ssd_net
if args.cuda:
net = torch.nn.DataParallel(ssd_net).cuda()
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 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)
criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, best_prior_weight=args.k, use_gpu=args.cuda)
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
batch_iterator = iter(data_loader)
for step in cfg['lr_steps']:
if args.start_iter > step:
print('over %d steps, adjust lr' % step)
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
else:
break
for iteration in range(args.start_iter, cfg['max_iter']):
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if args.cuda:
images = images.cuda()
targets = [ann.cuda() for ann in targets]
# targets = targets.cuda()
# else:
#
# targets = [ann for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ')
if iteration != 0 and iteration % 2000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(), ('weights/cache/%s_%s_' % (args.save_name, args.dataset)) +
repr(iteration) + '.pth')
name = '%s_%s' % (args.save_name, args.dataset)
torch.save(ssd_net.state_dict(),
args.save_folder + name + '.pth')
# different datasets, include coco, voc0712 trainval, coco val
datasets_version = {
"VOC": [('0712', '0712_trainval')],
"VOC0712++": [('0712', '0712_trainval_test')],
"VOC2012": [('2012', '2012_trainval')],
"COCO": [('2014', 'train'), ('2014', 'valminusminival')],
"VOC2007": [('0712', "2007_test")],
"COCOval": [('2014', 'minival')]
}
print('Loading Dataset...')
if dataset_name[0] == "V":
cfg = voc_config
train_dataset = VOCDetection(cfg["root"],
datasets_version[dataset_name], input_wh,
batch_size, cfg["multiscale"],
dataset_name)
elif dataset_name[0] == "C":
cfg = coco_config
train_dataset = COCODetection(cfg["root"],
datasets_version[dataset_name], input_wh,
batch_size, cfg["multiscale"],
dataset_name)
else:
print('Unkown dataset!')
# load Yolov3 net
net = Yolov3("train", input_wh, cfg["anchors"], cfg["anchors_mask"],
cfg["num_classes"])
if args.resume_net == None:
net.load_weights(basenet)
# build model
if args.version == 'yolo':
from model.yolo import YOLO
net = YOLO(device,
input_size=input_size,
num_classes=num_classes,
trainable=False)
print('Let us test yolo on the VOC0712 dataset ......')
else:
print('Unknown Version !!!')
exit()
# load net
net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(
args.voc_root, [('2007', set_type)],
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)), VOCAnnotationTransform())
net = net.to(device)
# evaluation
detect_net(args.save_folder,
net,
device,
dataset,
args.top_k,
thresh=args.confidence_threshold)
cfg = config_dict[args.dataset]
# from data.config import coco_on_voc
# cfg = coco_on_voc
net = build_ssd('test', cfg) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
no_anno = args.write_det_results
# load data
if args.dataset == 'helmet':
dataset = HelmetDetection(root, [('s2', set_type)],
BaseTransform(300, dataset_mean),
HelmetAnnotationTransform())
elif args.dataset == 'VOC':
dataset = VOCDetection(root, [('2007', set_type)],
BaseTransform(300, dataset_mean),
VOCAnnotationTransform(),
no_anno=no_anno)
elif args.dataset == 'SHWD':
dataset = SHWDDetection(root, (set_type, ),
BaseTransform(300, dataset_mean))
# elif args.dataset == 'VOC07':
# dataset = VOCDetection(root, [('2007', set_type)],
# BaseTransform(300, dataset_mean),
# VOCAnnotationTransform())
# elif args.dataset == 'VOC-v2':
# dataset = VOCDetection(root, [('2012', set_type)],
# BaseTransform(300, dataset_mean),
# VOCAnnotationTransform())
elif args.dataset == 'COCO18':
dataset = COCODetection(root, [('18', set_type)],
BaseTransform(300, dataset_mean),
def train(config, model):
if config.name == 'COCO':
from data.coco import COCODetection, COCOAnnotationTransform, COCO_CLASSES, get_label_map
elif config.name == 'VOC':
from data.voc0712 import VOCDetection, VOCAnnotationTransform, VOC_CLASSES
if not os.path.exists(config.train['save_folder']):
os.mkdir(config.train['save_folder'])
if config.train['dataset'] == 'COCO':
dataset = COCODetection(root=config.train['dataset_root'],
transform=SSDAugmentation(
config.min_dim, config.MEANS))
elif config.train['dataset'] == 'VOC':
dataset = VOCDetection(root=config.train['dataset_root'],
transform=SSDAugmentation(
config.min_dim, config.MEANS))
if config.train['visdom']:
import visdom
viz = visdom.Visdom()
if config.train['cuda']:
net = torch.nn.DataParallel(model.cuda())
cudnn.benchmark = True
if config.train['resume']:
print('Resuming training, loading {}...'.format(
config.train['resume']))
# model.load_weights(config.train['resume'])
optimizer = optim.SGD(net.parameters(),
lr=config.train['lr'],
momentum=config.train['momentum'],
weight_decay=config.train['weight_decay'])
criterion = MultiBoxLoss(config.variance, config.num_classes, 0.5, True, 0,
True, 3, 0.5, False, config.train['cuda'])
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer=optimizer, milestones=config.train['milestones'])
epochs = config.train['milestones'][-1]
net.train()
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // config.train['batch_size']
print('Training SSD on:', dataset.name)
step_index = 0
if config.train['visdom']:
vis_title = 'SSD.PyTorch on ' + dataset.name
vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
data_loader = data.DataLoader(dataset,
config.train['batch_size'],
num_workers=config.train['num_workers'],
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
print(len(dataset))
global_step = 0
for epoch in range(epochs):
print(config.train['save_folder'] + '' + config.train['dataset'] +
str(epoch) + '.pth')
for step, (images, targets) in enumerate(data_loader):
print(targets)
exit(0)
if config.train['cuda']:
images = images.cuda()
targets = [ann.cuda() for ann in targets]
else:
images = images
targets = [ann for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
optimizer.zero_grad()
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
if step > 0 and step % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('epoch' + repr(epoch) + ' iter ' + repr(step) +
' || Loss: %.4f ||' % (loss.item()),
end=' ')
# if config.train['visdom']:
# update_vis_plot(iteration, loss_l.data[0], loss_c.data[0], iter_plot, epoch_plot, 'append')
if step != 0 and step % 5000 == 0:
print('Saving state, iter:', global_step)
torch.save(
model.state_dict(), 'weights/peleenet_ssd800_COCO_' +
repr(global_step) + '.pth')
global_step += 1
# adjust lr
lr_scheduler.step()
torch.save(
model.state_dict(), config.train['save_folder'] + '' +
config.train['dataset'] + str(epoch) + '.pth')
def evaluate_detections(box_list, output_dir, dataset):
write_voc_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
# load net
num_classes = len(VOC_CLASSES) + 1 # +1 background
net = build_ssd('test', 300, num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(args.voc_root, [(YEAR, set_type)],
BaseTransform(300, dataset_mean),
AnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder,
net,
args.cuda,
dataset,
BaseTransform(net.size, dataset_mean),
args.top_k,
300,
thresh=args.confidence_threshold)
def evaluate_detections(box_list, output_dir, dataset):
write_voc_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
# load net
num_classes = cfg.NUM_CLASSES # +1 for background
net = build_net('test', num_classes) # initialize SSD
net.load_state_dict(torch.load(args.model))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(args.voc_root, [('2007', set_type)],
RefineBasicTransform(320, dataset_mean),
VOCAnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder, net, args.cuda, dataset,
RefineBasicTransform(
cfg.INPUT_SIZE, dataset_mean), args.top_k, 320,
thresh=args.confidence_threshold)
'''
output_dir = get_output_dir('refinedet300_20000', set_type)
do_python_eval(output_dir)'''
import matplotlib.pyplot as plt
from data.voc0712 import VOC_ROOT, VOCDetection, VOC_CLASSES
from data.config import MEANS, voc
from utils.augmentations import SSDAugmentation
import pickle
import math
import torch
import numpy as np
from math import sqrt
from itertools import product
rt = '..\\data\\VOCdevkit'
dataset = VOCDetection(root=rt,
transform=SSDAugmentation(voc['min_dim'], MEANS))
voc_config_points = []
config = [
(30, [1, 2, 1 / 2]),
(sqrt(30 * 60), [1]),
(60, [1, 2, 1 / 2, 3, 1 / 3]),
(sqrt(111 * 60), [1]),
(111, [1, 2, 1 / 2, 3, 1 / 3]),
(sqrt(111 * 162), [1]),
(162, [1, 2, 1 / 2, 3, 1 / 3]),
(sqrt(213 * 162), [1]),
(213, [1, 2, 1 / 2]),
(sqrt(213 * 264), [1]),
(264, [1, 2, 1 / 2]),
(sqrt(264 * 315), [1]),
]
for _x, _ylst in config:
def main():
'''
1.读取默认参数
2.读取数据
3.创建网络
4.设置优化器、损失函数
5.进入train函数或者eval函数
'''
args = get_args()
# 读取数据
if args.dataset == 'VOC':
train_dataset = VOCDetection(root=args.dataset_root)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
num_workers=4,
collate_fn=detection_collate,
shuffle=True,
pin_memory=True)
# create batch iterator
batch_iterator = iter(train_loader)
# 创建网络
ssd_net = build_ssd('train', voc['min_dim'], voc['num_classes'])
if torch.cuda.is_available():
ssd_net = ssd_net.cuda()
cudnn.benchmark = True
# 优化器SGD
optimizer = optim.SGD(ssd_net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# 损失函数
criterion = MultiBoxLoss(num_classes=voc['num_classes'],
overlap_thresh=0.5,
neg_pos=3)
# 进入训练
ssd_net.train()
start_iter = 0 # 训练步数
step_index = 0 # 统计学习率
for iteration in tqdm(range(start_iter, voc['max_iter'])):
# 调整学习率
if iteration in voc['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args, step_index)
# 加载图像
try:
images, targets, _, _ = next(batch_iterator)
except:
batch_iterator = iter(
torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
num_workers=4,
collate_fn=detection_collate,
shuffle=True,
pin_memory=True))
images, targets, _, _ = next(batch_iterator)
images = images.cuda()
targets = [x.cuda() for x in targets]
# 前向
out = ssd_net(images)
# 后向
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
if iteration % 100 == 0:
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.item()))
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(),
'weights/ssd300_VOC_' + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(), 'weights/VOC_final.pth')
def train():
if args.dataset == 'COCO':
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
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))
elif args.dataset == 'CUSTOM':
if args.dataset_root == VOC_ROOT or args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')
cfg = custom
dataset = customDetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))
if args.visdom:
import visdom
viz = visdom.Visdom()
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:
pass
# resnet_weights = torch.load(args.save_folder + args.basenet)
# print('Loading base network...')
# ssd_net.resnet.load_state_dict(resnet_weights)
if args.cuda:
net = net.cuda()
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)
criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Epochj Size:', epoch_size)
print('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.visdom:
vis_title = 'SSD.PyTorch on ' + dataset.name
vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
batch_iterator = iter(data_loader)
for iteration in range(args.start_iter, cfg['max_iter']):
if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None,
'append', epoch_size)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
# images, targets = next(batch_iterator)
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
else:
images = Variable(images)
targets = [Variable(ann, volatile=True) for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
# loc_loss += loss_l.data[0]
# conf_loss += loss_c.data[0]
loc_loss += loss_l.item()
conf_loss += loss_c.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
# print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.item()), end=' ')
if args.visdom:
# update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
# iter_plot, epoch_plot, 'append')
update_vis_plot(iteration, loss_l.item(), loss_c.item(),
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(), args.save_folder + '/ssd300_COCO_' +
repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
# Puddle Detection Dataset
puddle_root = '/home/sean/Downloads/road_slip_hazards'
pjson_file = 'mine_and_car_puddles'
p_save = os.path.join(os.path.dirname(m_save),
'all_puddle_dataset_means.txt')
if not os.path.isfile(p_save):
p_mean = getMiningMean(puddle_root, pjson_file)
with open(p_save, 'w') as f:
f.write(str(p_mean))
else:
p_mean = np.array([107.86730173, 106.28081276, 107.2159824])
print('Testing Constructors')
voc_dataset = VOCDetection(voc_root, voc_sets,
SSDAugmentation(ssd_dim, voc_means),
AnnotationTransform())
mdataset = MiningDataset(mining_root,
transform=SSDMiningAugmentation(ssd_dim, m_mean),
target_transform=MiningAnnotationTransform(),
json_set=json_file)
pud_dataset = PuddleDataset(
puddle_root,
transform=SSDMiningAugmentation(ssd_dim, p_mean),
target_transform=MiningAnnotationTransform(class_to_ind={'puddle': 0}),
json_set=pjson_file,
dataset_name='PUDDLES')
data_iterators = [pud_dataset, voc_dataset, mdataset]