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 == 'WHEAT':
# if args.dataset_root != WHEAT_ROOT:
# print('Must specify dataset if specifying dataset_root')
# cfg = wheat
# dataset = WHEATDetection(root=args.dataset_root,
# transform=SSDAugmentation(cfg['min_dim'],
# MEANS))
if args.dataset_root != WHEAT_ROOT:
print('Must specify dataset if specifying dataset_root')
cfg = wheat
dataset = WHEATDetection(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:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_weights)
if args.cuda:
#handbook
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# net = net.cuda()
net = net.to(device)
#handbook
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('Training SSD on:', dataset.name)
print('Using the specified args:')
print('epoch_size: ', epoch_size)
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)
print("epoch: ", epoch, " loc_loss: ", loc_loss, " conf_loss: ",
conf_loss)
# 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)
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if args.cuda:
# with torch.no_grad():
#handbook
# images = images.cuda()
# targets = [ann.cuda() for ann in targets]
device = 'cuda' if torch.cuda.is_available() else 'cpu'
images = images.to(device)
targets = [ann.to(device) for ann in targets]
#handbook
# 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.item()
conf_loss += loss_c.data.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data.item()),
end=' ')
if args.visdom:
update_vis_plot(iteration, loss_l.data.item(), loss_c.data.item(),
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 200 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(),
'weights/ssd300_WHEAT_' + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0
print('Loading Dataset...')
dataset = VOCDetection(args.voc_root, train_sets,
SSDAugmentation(ssd_dim, means),
AnnotationTransform())
epoch_size = len(dataset) // args.batch_size
print('Training SSD on', dataset.name)
step_index = 0
if args.visdom:
# initialize visdom loss plot
lot = viz.line(X=torch.zeros((1, )).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(xlabel='Iteration',
ylabel='Loss',
title='Current SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']))
epoch_lot = viz.line(X=torch.zeros((1, )).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Epoch',
ylabel='Loss',
title='Epoch SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']))
batch_iterator = None
data_loader = data.DataLoader(dataset,
batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate)
for iteration in range(args.start_iter, max_iter):
if (not batch_iterator) or (iteration % epoch_size == 0):
# create batch iterator
batch_iterator = iter(data_loader)
if iteration in stepvalues:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * epoch,
Y=torch.Tensor([loc_loss, conf_loss, loc_loss + conf_loss
]).unsqueeze(0).cpu() / epoch_size,
win=epoch_lot,
update='append')
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
# load train data
images, targets = next(batch_iterator)
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(anno.cuda()) for anno in targets]
else:
images = Variable(images)
targets = [Variable(anno) for anno 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]
if iteration % 10 == 0:
print('Timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data[0]),
end=' ')
if args.visdom and args.send_images_to_visdom:
random_batch_index = np.random.randint(images.size(0))
viz.image(images.data[random_batch_index].cpu().numpy())
if args.visdom:
viz.line(X=torch.ones((1, 3)).cpu() * iteration,
Y=torch.Tensor([
loss_l.data[0], loss_c.data[0],
loss_l.data[0] + loss_c.data[0]
]).unsqueeze(0).cpu(),
win=lot,
update='append')
# hacky fencepost solution for 0th epoch plot
if iteration == 0:
viz.line(X=torch.zeros((1, 3)).cpu(),
Y=torch.Tensor(
[loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu(),
win=epoch_lot,
update=True)
if iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(
ssd_net.state_dict(),
os.path.join(args.save_folder,
'ssd300_0712_iter_' + repr(iteration) + '.pth'))
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.version + '.pth')
def train(args, net, priors, optimizer, criterion, scheduler):
log_file = open(args.save_root+"training.log", "w", 1)
log_file.write(args.exp_name+'\n')
for arg in sorted(vars(args)):
print(arg, getattr(args, arg))
log_file.write(str(arg)+': '+str(getattr(args, arg))+'\n')
net.train()
# loss counters
batch_time = AverageMeter()
losses = AverageMeter()
loc_losses = AverageMeter()
cls_losses = AverageMeter()
print('Loading Dataset...')
train_dataset = ActionDetection(args, args.train_sets, SSDAugmentation(args.ssd_dim, args.means),
NormliseBoxes(), anno_transform=MatchPrior(priors, args.cfg['variance']))
log_file.write(train_dataset.print_str)
print(train_dataset.print_str)
val_dataset = ActionDetection(args, args.test_sets, BaseTransform(args.ssd_dim, args.means),
NormliseBoxes(), full_test=False)
log_file.write(val_dataset.print_str)
# print(val_dataset.print_str)
epoch_size = len(train_dataset) // args.batch_size
print('Training SSD on', train_dataset.name)
if args.visdom:
import visdom
viz = visdom.Visdom(env=args.exp_name, port=args.vis_port)
# initialize visdom loss plot
lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1, 6)).cpu(),
opts=dict(
xlabel='Iteration',
ylabel='Loss',
title='Current SSD Training Loss',
legend=['REG', 'CLS', 'AVG', 'S-REG', ' S-CLS', ' S-AVG']
)
)
# initialize visdom meanAP and class APs plot
legends = ['meanAP']
for cls in CLASSES[args.dataset]:
legends.append(cls)
val_lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1,args.num_classes)).cpu(),
opts=dict(
xlabel='Iteration',
ylabel='Mean AP',
title='Current SSD Validation mean AP',
legend=legends
)
)
batch_iterator = None
train_data_loader = data.DataLoader(train_dataset, args.batch_size, num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate, pin_memory=True)
val_data_loader = data.DataLoader(val_dataset, args.batch_size, num_workers=args.num_workers,
shuffle=False, collate_fn=detection_collate, pin_memory=True)
itr_count = 0
torch.cuda.synchronize()
t0 = time.perf_counter()
for iteration in range(args.max_iter + 1):
if (not batch_iterator) or (iteration % epoch_size == 0):
# create batch iterator
batch_iterator = iter(train_data_loader)
# load train data
images, _ , prior_gt_labels, prior_gt_locations, _, _ = next(batch_iterator)
# images, ground_truths, _ , _, num_mt, img_indexs
# pdb.set_trace()
images = [img.cuda(0, non_blocking=True) for img in images if not isinstance(img, list)]
prior_gt_labels = prior_gt_labels.cuda(0, non_blocking=True)
prior_gt_locations = prior_gt_locations.cuda(0, non_blocking=True)
# forward
cls_out, reg_out = net(images)
optimizer.zero_grad()
loss_l, loss_c = criterion(cls_out, reg_out, prior_gt_labels, prior_gt_locations)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
scheduler.step()
# pdb.set_trace()
loc_loss = loss_l.item()
conf_loss = loss_c.item()
# print('Loss data type ',type(loc_loss))
loc_losses.update(loc_loss)
cls_losses.update(conf_loss)
losses.update((loc_loss + conf_loss)/2.0)
if iteration == 103:
loc_losses.reset()
cls_losses.reset()
losses.reset()
batch_time.reset()
if iteration % args.print_step == 0:
if args.visdom and iteration>100:
losses_list = [loc_losses.val, cls_losses.val, losses.val, loc_losses.avg, cls_losses.avg, losses.avg]
viz.line(X=torch.ones((1, 6)).cpu() * iteration,
Y=torch.from_numpy(np.asarray(losses_list)).unsqueeze(0).cpu(),
win=lot,
update='append')
torch.cuda.synchronize()
t1 = time.perf_counter()
batch_time.update(t1 - t0)
print_line = 'Itration {:02d}/{:06d}/{:06d} loc-loss {:.3f}({:.3f}) cls-loss {:.3f}({:.3f}) ' \
'average-loss {:.3f}({:.3f}) Timer {:0.3f}({:0.3f})'.format(iteration//epoch_size,
iteration, args.max_iter, loc_losses.val, loc_losses.avg, cls_losses.val,
cls_losses.avg, losses.val, losses.avg, batch_time.val, batch_time.avg)
torch.cuda.synchronize()
t0 = time.perf_counter()
log_file.write(print_line+'\n')
print(print_line)
itr_count += 1
if itr_count % args.loss_reset_step == 0 and itr_count > 0:
loc_losses.reset()
cls_losses.reset()
losses.reset()
batch_time.reset()
print('Reset ', args.exp_name,' after', itr_count*args.print_step)
itr_count = 0
if (iteration % args.val_step == 0 or iteration in [1000, args.max_iter]) and iteration>0:
torch.cuda.synchronize()
tvs = time.perf_counter()
print('Saving state, iter:', iteration)
torch.save(net.state_dict(), args.save_root + 'AMTNet_' +
repr(iteration) + '.pth')
net.eval() # switch net to evaluation mode
with torch.no_grad():
mAP, ap_all, ap_strs = validate(args, net, priors, val_data_loader, val_dataset, iteration, iou_thresh=args.iou_thresh)
for ap_str in ap_strs:
print(ap_str)
log_file.write(ap_str+'\n')
ptr_str = '\nMEANAP:::=>'+str(mAP)+'\n'
print(ptr_str)
log_file.write(ptr_str)
if args.visdom:
aps = [mAP]
for ap in ap_all:
aps.append(ap)
viz.line(
X=torch.ones((1, args.num_classes)).cpu() * iteration,
Y=torch.from_numpy(np.asarray(aps)).unsqueeze(0).cpu(),
win=val_lot,
update='append'
)
net.train() # Switch net back to training mode
torch.cuda.synchronize()
t0 = time.perf_counter()
prt_str = '\nValidation TIME::: {:0.3f}\n\n'.format(t0-tvs)
print(prt_str)
log_file.write(ptr_str)
log_file.close()
if __name__ == '__main__':
if not args.cuda:
print("this file only supports cuda version now!")
# store pruning models
if not os.path.exists(args.prune_folder):
os.mkdir(args.prune_folder)
print(args)
# load model from previous pruning
model = torch.load(args.pruned_model).cuda()
print('Finished loading model!')
# data
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'], cfg['dataset_mean']))
testset = VOCDetection(root=args.dataset_root, image_sets=[('2007', 'test')],
transform=BaseTransform(cfg['min_dim'], cfg['testset_mean']))
data_loader = data.DataLoader(dataset, 32, num_workers=4,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
arm_criterion = RefineMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5, False, 0, args.cuda)
odm_criterion = RefineMultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5, False, 0.01, args.cuda)# 0.01 -> 0.99 negative confidence threshold
fine_tuner = FineTuner_refineDet(data_loader, testset, arm_criterion, odm_criterion, model)
# ------------------------ adjustable part
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
map = fine_tuner.train(optimizer = optimizer, epoches = args.epoch)
# ------------------------ adjustable part
def train(model, device):
global cfg, hr
# set GPU
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
use_focal = False
if args.use_focal == 1:
print("Let's use focal loss for objectness !!!")
use_focal = True
if args.multi_scale == 1:
print('Let us use the multi-scale trick.')
ms_inds = range(len(cfg['multi_scale']))
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation([608, 608], MEANS))
else:
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
from torch.utils.tensorboard import SummaryWriter
# c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
log_path = 'log/yolo_v2/voc2007/' # + c_time
if not os.path.exists(log_path):
os.mkdir(log_path)
writer = SummaryWriter(log_path)
print(
"----------------------------------------Object Detection--------------------------------------------"
)
print("Let's train OD network !")
net = model
net = net.to(device)
# optimizer = optim.Adam(net.parameters())
lr = args.lr
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# loss counters
print("----------------------------------------------------------")
print('Loading the dataset...')
print('Training on:', dataset.name)
print('The dataset size:', len(dataset))
print('The obj weight : ', args.obj)
print('The noobj weight : ', args.noobj)
print("----------------------------------------------------------")
input_size = cfg['min_dim']
step_index = 0
epoch_size = len(dataset) // args.batch_size
# each part of loss weight
obj_w = 1.0
cla_w = 1.0
box_w = 2.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
iteration = 0
# start training
for epoch in range(cfg['max_epoch']):
batch_iterator = iter(data_loader)
# No WarmUp strategy or WarmUp stage has finished.
if epoch in cfg['lr_epoch']:
step_index += 1
lr = adjust_learning_rate(optimizer, args.gamma, step_index)
for images, targets in batch_iterator:
# WarmUp strategy for learning rate
if args.warm_up == 'yes':
if epoch < args.wp_epoch:
lr = warmup_strategy(optimizer, epoch_size, iteration)
iteration += 1
# multi-scale trick
if iteration % 10 == 0 and args.multi_scale == 1:
ms_ind = random.sample(ms_inds, 1)[0]
input_size = cfg['multi_scale'][int(ms_ind)]
# multi scale
if args.multi_scale == 1:
images = torch.nn.functional.interpolate(images,
size=input_size,
mode='bilinear',
align_corners=True)
targets = [label.tolist() for label in targets]
if args.version == 'yolo_v2':
targets = tools.gt_creator(input_size, yolo_net.stride,
args.num_classes, targets)
elif args.version == 'yolo_v3':
targets = tools.multi_gt_creator(input_size, yolo_net.stride,
args.num_classes, targets)
targets = torch.tensor(targets).float().to(device)
# forward
t0 = time.time()
out = net(images.to(device))
optimizer.zero_grad()
obj_loss, class_loss, box_loss = tools.loss(
out,
targets,
num_classes=args.num_classes,
use_focal=use_focal,
obj=args.obj,
noobj=args.noobj)
# print(obj_loss.item(), class_loss.item(), box_loss.item())
total_loss = obj_w * obj_loss + cla_w * class_loss + box_w * box_loss
# viz loss
writer.add_scalar('object loss', obj_loss.item(), iteration)
writer.add_scalar('class loss', class_loss.item(), iteration)
writer.add_scalar('local loss', box_loss.item(), iteration)
# backprop
total_loss.backward()
optimizer.step()
t1 = time.time()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print(obj_loss.item(), class_loss.item(), box_loss.item())
print('Epoch[%d / %d]' % (epoch+1, cfg['max_epoch']) + ' || iter ' + repr(iteration) + \
' || Loss: %.4f ||' % (total_loss.item()) + ' || lr: %.8f ||' % (lr) + ' || input size: %d ||' % input_size[0], end=' ')
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(
yolo_net.state_dict(), args.save_folder + '/' + args.version +
'_' + repr(epoch + 1) + '.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_refinedet[args.input_size]
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
if args.visdom:
import visdom
viz = visdom.Visdom()
refinedet_net = build_refinedet('train', cfg['min_dim'],
cfg['num_classes'])
net = refinedet_net
print(net)
#input()
if args.cuda:
net = torch.nn.DataParallel(refinedet_net)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
refinedet_net.load_weights(args.resume)
else:
#vgg_weights = torch.load(args.save_folder + args.basenet)
vgg_weights = torch.load(args.basenet)
print('Loading base network...')
refinedet_net.vgg.load_state_dict(vgg_weights)
if args.cuda:
net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
refinedet_net.extras.apply(weights_init)
refinedet_net.arm_loc.apply(weights_init)
refinedet_net.arm_conf.apply(weights_init)
refinedet_net.odm_loc.apply(weights_init)
refinedet_net.odm_conf.apply(weights_init)
#refinedet_net.tcb.apply(weights_init)
refinedet_net.tcb0.apply(weights_init)
refinedet_net.tcb1.apply(weights_init)
refinedet_net.tcb2.apply(weights_init)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5, False,
args.cuda)
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
net.train()
# loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training RefineDet on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.visdom:
vis_title = 'RefineDet.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, arm_loc_loss, arm_conf_loss, epoch_plot,
None, 'append', epoch_size)
# reset epoch loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_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
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]
else:
images = images
targets = [ann for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
arm_loss_l, arm_loss_c = arm_criterion(out, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, targets)
#input()
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
loss.backward()
optimizer.step()
t1 = time.time()
arm_loc_loss += arm_loss_l.item()
arm_conf_loss += arm_loss_c.item()
odm_loc_loss += odm_loss_l.item()
odm_conf_loss += odm_loss_c.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || ARM_L Loss: %.4f ARM_C Loss: %.4f ODM_L Loss: %.4f ODM_C Loss: %.4f ||' \
% (arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item()), end=' ')
if args.visdom:
update_vis_plot(iteration, arm_loss_l.data[0], arm_loss_c.data[0],
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(
refinedet_net.state_dict(),
args.save_folder + '/RefineDet{}_{}_{}.pth'.format(
args.input_size, args.dataset, repr(iteration)))
torch.save(
refinedet_net.state_dict(), args.save_folder +
'/RefineDet{}_{}_final.pth'.format(args.input_size, args.dataset))
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0
print('Loading Dataset...')
dataset = Dataset(SSDAugmentation(ssd_dim, means))
epoch_size = len(dataset) // args.batch_size
print('Training SSD on', dataset.name)
step_index = 0
batch_iterator = None
data_loader = data.DataLoader(
dataset,
batch_size, #num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
scheduler = ReduceLROnPlateau(optimizer,
patience=1000,
min_lr=1e-6,
verbose=True)
for iteration in range(args.start_iter, max_iter):
if (not batch_iterator) or (iteration % epoch_size == 0):
# create batch iterator
batch_iterator = iter(data_loader)
if iteration in stepvalues:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
# load train data
images, targets = next(batch_iterator)
if args.cuda:
images = Variable(images.cuda())
targets = [
Variable(anno.cuda(), volatile=True) for anno in targets
]
else:
images = Variable(images)
targets = [Variable(anno, volatile=True) for anno 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]
print('Timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]),
end=' ')
scheduler.step(loss.data[0])
if iteration % 5000 == 0 and iteration != 0:
print('Saving state, iter:', iteration)
checkpoint(ssd_net, optimizer,
'weights/ssd300_0712_' + repr(iteration) + '.pth',
iteration)
checkpoint(ssd_net, optimizer,
args.save_folder + '' + args.version + '.pth', iteration)
def train():
best_map = 0
if args.dataset == 'VOC':
if args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')
cfg = voc_refinedet[args.input_size]
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
refinedet_net = build_refinedet('train', cfg['min_dim'],
cfg['num_classes'])
net = refinedet_net
# print(net)
#input()
if args.cuda:
# net = torch.nn.DataParallel(refinedet_net)
net = net.to(device)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
refinedet_net.load_weights(args.resume)
else:
#vgg_weights = torch.load(args.save_folder + args.basenet)
vgg_weights = torch.load(args.basenet)
print('Loading base network...')
refinedet_net.vgg.load_state_dict(vgg_weights)
# if args.cuda:
# net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
refinedet_net.extras.apply(weights_init)
refinedet_net.arm_loc.apply(weights_init)
refinedet_net.arm_conf.apply(weights_init)
refinedet_net.odm_loc.apply(weights_init)
refinedet_net.odm_conf.apply(weights_init)
#refinedet_net.tcb.apply(weights_init)
refinedet_net.tcb0.apply(weights_init)
refinedet_net.tcb1.apply(weights_init)
refinedet_net.tcb2.apply(weights_init)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5, False,
args.cuda)
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
# net.train()
# loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training RefineDet 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
iteration = 0
for epoch in range(args.epochs):
net.phase = 'train'
net.train()
for batch_i, (images, targets) in enumerate(data_loader):
iteration += 1
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
images = images.to(device)
targets = [ann.to(device) for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
arm_loss_l, arm_loss_c = arm_criterion(out, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, targets)
#input()
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
loss.backward()
optimizer.step()
t1 = time.time()
arm_loc_loss += arm_loss_l.item()
arm_conf_loss += arm_loss_c.item()
odm_loc_loss += odm_loss_l.item()
odm_conf_loss += odm_loss_c.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('epoch '+repr(epoch+1)+': iter ' + repr(iteration) + ' || ARM_L Loss: %.4f ARM_C Loss: %.4f ODM_L Loss: %.4f ODM_C Loss: %.4f ||' \
% (arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item()), end=' ')
if epoch % args.evaluation_interval == 0:
print("\n---- Evaluating Model ----")
map = evaluate(model=net,
save_folder=args.save_folders,
cuda=args.cuda,
top_k=5,
im_size=320,
thresh=0.05,
dataset_mean=((104, 117, 123)))
if map > best_map: #取最好的map保存
torch.save(net.state_dict(),
f"checkpoints/RefineDet320_%d.pth" % (epoch + 1))
best_map = map
def train():
if args.dataset == 'COCO':
cfg = coco
dataset = COCODetection(root=cfg['coco_root'],
transform=SSDAugmentation(cfg['min_dim'],
MEANS))
if args.dataset == 'VOC':
cfg = voc
dataset = VOCDetection(root=cfg['voc_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:
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()
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)
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
priors = priors.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
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
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(), requires_grad=False) for ann in targets]
else:
images = Variable(images)
targets = [Variable(ann, requires_grad=False) for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, 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 args.visdom:
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(), 'weights/ssd300_COCO_' +
repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
def train():
if args.dataset == 'COCO':
cfg = coco_refinedet[args.input_size]
root = '/home/soyeol/data/coco'
dataset = COCODetection(root=root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
print("len_dataset(COCO) :", len(dataset))
elif args.dataset == 'VOC':
cfg = voc_refinedet[args.input_size]
root = '/home/soyeol/data/VOCdevkit'
# root = "/content/drive/MyDrive/Colab Notebooks/dataset/VOCdevkit"
dataset = VOCDetection(root=root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
print("len_dataset(VOC) :", len(dataset))
val_dataset = VOCDetection(root=root,
transform=BaseTransform(args.input_size,
mean=(104, 117,
123)),
image_sets=[('2007', 'test')],
mode='test')
net = build_refinedet('train', cfg['min_dim'],
cfg['num_classes']) # , batch_size=args.batch_size)
print(net)
torch.multiprocessing.set_start_method('spawn')
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True,
drop_last=True)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5, False,
args.cuda)
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
if args.cuda:
# net = torch.nn.DataParallel(refinedet_net)
cudnn.benchmark = True
iteration = 0
num_epoch = int(cfg['max_iter'] * args.batch_size / len(data_loader)) + 1
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
state = torch.load(args.resume)
net.load_state_dict(state['model'])
optimizer.load_state_dict(state['optimizer'])
iteration = state['iter']
else:
#vgg_weights = torch.load(args.save_folder + args.basenet)
vgg_weights = torch.load(args.basenet)
print('Loading base network...')
net.vgg.load_state_dict(vgg_weights)
if args.cuda:
net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
net.extras.apply(weights_init)
# refinedet_net.arm_loc.apply(weights_init)
# refinedet_net.arm_conf.apply(weights_init)
# refinedet_net.odm_loc.apply(weights_init)
# refinedet_net.odm_conf.apply(weights_init)
#refinedet_net.tcb.apply(weights_init)
net.tcb0.apply(weights_init)
net.tcb1.apply(weights_init)
net.tcb2.apply(weights_init)
# if args.resume :
# optimizer =
net.train()
arm_iter = LossCal()
odm_iter = LossCal()
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training RefineDet on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.tensorboard:
writer = SummaryWriter(os.path.join(args.log_folder, args.log))
# create batch iterator
detect_time_sum = 0.
load_time_sum = 0.
print_count = 0
t2 = time.time()
# for iteration in range(args.start_iter, cfg['max_iter']):
for epoch in range(num_epoch):
for batch_data in data_loader:
t3 = time.time()
load_time_sum += (t3 - t2)
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
images, targets = batch_data
if args.cuda:
images = images.cuda()
targets = [ann.cuda() for ann in targets]
else:
images = images
targets = [ann for ann in targets]
t0 = time.time()
out = net(images)
# if mini_iter % args.iter_size == 1:
optimizer.zero_grad()
arm_loss_l, arm_loss_c = arm_criterion(out, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, targets)
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = (arm_loss + odm_loss)
loss.backward()
arm_iter.stack(arm_loss_l.item(), arm_loss_c.item())
odm_iter.stack(odm_loss_l.item(), odm_loss_c.item())
# if mini_iter % args.iter_size == 0:
optimizer.step()
t1 = time.time()
detect_time_sum += (t1 - t0)
iteration += 1
print_count += 1
if iteration % args.print_freq == 0:
al, ac = arm_iter.pop()
ol, oc = odm_iter.pop()
print('iter ' + repr(
iteration) + ' || Loss: ARM_L: %.4f ARM_C: %.4f ODM_L: %.4f ODM_C: %.4f ||' \
% (al, ac, ol, oc), end=' ')
print('detect_time: %.4f sec. || load_time: %.4f' %
(detect_time_sum / print_count,
load_time_sum / print_count))
print_count = 0
detect_time_sum = 0.
load_time_sum = 0.
if args.tensorboard:
writer.add_scalar('ARM_loss/loc', al, iteration)
writer.add_scalar('ARM_loss/conf', ac, iteration)
writer.add_scalar('ODM_loss/loc', ol, iteration)
writer.add_scalar('ODM_loss/conf', oc, iteration)
arm_iter.reset()
odm_iter.reset()
if iteration != 0 and iteration % args.save_freq == 0:
print('Saving state, iter:', iteration)
torch.save(
{
'model': net.state_dict(),
'optimizer': optimizer.state_dict(),
'iter': iteration,
'lr': optimizer.param_groups[0]['lr']
}, save_folder + '/RefineDet{}_{}_{}.pth'.format(
args.input_size, args.dataset, repr(iteration)))
if args.dataset == 'VOC' and iteration != 0 and iteration % args.valid_freq == 0:
net.phase = 'test'
map = _valid(net, args, iteration, val_dataset) # dataset[1])
writer.add_scalar('mAP', map, iteration)
net.phase = 'train'
if iteration == cfg['max_iter']:
break
t2 = time.time()
if iteration == cfg['max_iter']:
print('Saving state:', iteration)
torch.save(
{
'model': net.state_dict(),
'optimizer': optimizer.state_dict(),
'iter': iteration,
'lr': optimizer.param_groups[0]['lr']
}, save_folder + '/RefineDet{}_{}_{}.pth'.format(
args.input_size, args.dataset, repr(iteration)))
break
writer.close()
def train():
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
dataset = COCODetection(
image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
transform=SSDAugmentation(MEANS),
)
if args.validation_epoch > 0:
setup_eval()
val_dataset = COCODetection(
image_path=cfg.dataset.valid_images,
info_file=cfg.dataset.valid_info,
transform=BaseTransform(MEANS),
)
# Parallel wraps the underlying module, but when saving and loading we don't want that
yolact_net = Yolact()
net = yolact_net
net.train()
if args.log:
log = Log(
cfg.name,
args.log_folder,
dict(args._get_kwargs()),
overwrite=(args.resume is None),
log_gpu_stats=args.log_gpu,
)
# I don't use the timer during training (I use a different timing method).
# Apparently there's a race condition with multiple GPUs, so disable it just to be safe.
timer.disable_all()
# Both of these can set args.resume to None, so do them before the check
if args.resume == "interrupt":
args.resume = SavePath.get_interrupt(args.save_folder)
elif args.resume == "latest":
args.resume = SavePath.get_latest(args.save_folder, cfg.name)
if args.resume is not None:
print("Resuming training, loading {}...".format(args.resume))
yolact_net.load_weights(args.resume)
if args.start_iter == -1:
args.start_iter = SavePath.from_str(args.resume).iteration
else:
print("Initializing weights...")
yolact_net.init_weights(backbone_path=args.save_folder +
cfg.backbone.path)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.decay)
criterion = MultiBoxLoss(
num_classes=cfg.num_classes,
pos_threshold=cfg.positive_iou_threshold,
neg_threshold=cfg.negative_iou_threshold,
negpos_ratio=cfg.ohem_negpos_ratio,
)
if args.batch_alloc is not None:
args.batch_alloc = [int(x) for x in args.batch_alloc.split(",")]
if sum(args.batch_alloc) != args.batch_size:
print(
"Error: Batch allocation (%s) does not sum to batch size (%s)."
% (args.batch_alloc, args.batch_size))
exit(-1)
net = CustomDataParallel(NetLoss(net, criterion))
if args.cuda:
net = net.cuda()
# Initialize everything
if not cfg.freeze_bn:
yolact_net.freeze_bn() # Freeze bn so we don't kill our means
yolact_net(torch.zeros(1, 3, cfg.max_size, cfg.max_size).cuda())
if not cfg.freeze_bn:
yolact_net.freeze_bn(True)
# loss counters
loc_loss = 0
conf_loss = 0
iteration = max(args.start_iter, 0)
last_time = time.time()
epoch_size = len(dataset) // args.batch_size
num_epochs = math.ceil(cfg.max_iter / epoch_size)
# Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
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,
)
save_path = lambda epoch, iteration: SavePath(
cfg.name, epoch, iteration).get_path(root=args.save_folder)
time_avg = MovingAverage()
global loss_types # Forms the print order
loss_avgs = {k: MovingAverage(100) for k in loss_types}
print("Begin training!")
print()
# try-except so you can use ctrl+c to save early and stop training
try:
for epoch in range(num_epochs):
# Resume from start_iter
if (epoch + 1) * epoch_size < iteration:
continue
for datum in data_loader:
# Stop if we've reached an epoch if we're resuming from start_iter
if iteration == (epoch + 1) * epoch_size:
break
# Stop at the configured number of iterations even if mid-epoch
if iteration == cfg.max_iter:
break
# Change a config setting if we've reached the specified iteration
changed = False
for change in cfg.delayed_settings:
if iteration >= change[0]:
changed = True
cfg.replace(change[1])
# Reset the loss averages because things might have changed
for avg in loss_avgs:
avg.reset()
# If a config setting was changed, remove it from the list so we don't keep checking
if changed:
cfg.delayed_settings = [
x for x in cfg.delayed_settings if x[0] > iteration
]
# Warm up by linearly interpolating the learning rate from some smaller value
if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until:
set_lr(
optimizer,
(args.lr - cfg.lr_warmup_init) *
(iteration / cfg.lr_warmup_until) + cfg.lr_warmup_init,
)
# Adjust the learning rate at the given iterations, but also if we resume from past that iteration
while (step_index < len(cfg.lr_steps)
and iteration >= cfg.lr_steps[step_index]):
step_index += 1
set_lr(optimizer, args.lr * (args.gamma**step_index))
# Zero the grad to get ready to compute gradients
optimizer.zero_grad()
# Forward Pass + Compute loss at the same time (see CustomDataParallel and NetLoss)
losses = net(datum)
losses = {k: (v).mean()
for k, v in losses.items()
} # Mean here because Dataparallel
loss = sum([losses[k] for k in losses])
# no_inf_mean removes some components from the loss, so make sure to backward through all of it
# all_loss = sum([v.mean() for v in losses.values()])
# Backprop
loss.backward(
) # Do this to free up vram even if loss is not finite
if torch.isfinite(loss).item():
optimizer.step()
# Add the loss to the moving average for bookkeeping
for k in losses:
loss_avgs[k].add(losses[k].item())
cur_time = time.time()
elapsed = cur_time - last_time
last_time = cur_time
# Exclude graph setup from the timing information
if iteration != args.start_iter:
time_avg.add(elapsed)
if iteration % 10 == 0:
eta_str = str(
datetime.timedelta(seconds=(cfg.max_iter - iteration) *
time_avg.get_avg())).split(".")[0]
total = sum([loss_avgs[k].get_avg() for k in losses])
loss_labels = sum(
[[k, loss_avgs[k].get_avg()]
for k in loss_types if k in losses],
[],
)
print(
("[%3d] %7d ||" + (" %s: %.3f |" * len(losses)) +
" T: %.3f || ETA: %s || timer: %.3f") %
tuple([epoch, iteration] + loss_labels +
[total, eta_str, elapsed]),
flush=True,
)
if args.log:
precision = 5
loss_info = {
k: round(losses[k].item(), precision)
for k in losses
}
loss_info["T"] = round(loss.item(), precision)
if args.log_gpu:
log.log_gpu_stats = iteration % 10 == 0 # nvidia-smi is sloooow
log.log(
"train",
loss=loss_info,
epoch=epoch,
iter=iteration,
lr=round(cur_lr, 10),
elapsed=elapsed,
)
log.log_gpu_stats = args.log_gpu
iteration += 1
if iteration % args.save_interval == 0 and iteration != args.start_iter:
if args.keep_latest:
latest = SavePath.get_latest(args.save_folder,
cfg.name)
print("Saving state, iter:", iteration)
yolact_net.save_weights(save_path(epoch, iteration))
if args.keep_latest and latest is not None:
if (args.keep_latest_interval <= 0
or iteration % args.keep_latest_interval !=
args.save_interval):
print("Deleting old save...")
os.remove(latest)
# This is done per epoch
if args.validation_epoch > 0:
if epoch % args.validation_epoch == 0 and epoch > 0:
compute_validation_map(
epoch,
iteration,
yolact_net,
val_dataset,
log if args.log else None,
)
# Compute validation mAP after training is finished
compute_validation_map(epoch, iteration, yolact_net, val_dataset,
log if args.log else None)
except KeyboardInterrupt:
if args.interrupt:
print("Stopping early. Saving network...")
# Delete previous copy of the interrupted network so we don't spam the weights folder
SavePath.remove_interrupt(args.save_folder)
yolact_net.save_weights(
save_path(epoch,
repr(iteration) + "_interrupt"))
exit()
yolact_net.save_weights(save_path(epoch, iteration))
viz.line(X=torch.zeros((1, 3)).cpu(),
Y=torch.Tensor([loc, conf, loc + conf]).unsqueeze(0).cpu(),
win=window2,
update=True)
if __name__ == '__main__':
print("time before calling train(): ", datetime.now().time())
if args.dataset == 'VOC':
cfg = VOC_scat
dataset = VOCDetection(root=VOC_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=VOC_MEANS,
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'VOC_small':
cfg = VOC_scat_small
dataset = VOCDetection(root=VOC_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=VOC_MEANS,
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
def train():
print("start...")
dataset = DetectionDataset(root=args.dataset_root,
transform=SSDAugmentation(300, cfg['mean']),
phase='train')
print(len(dataset))
ssd_net = build_ssd('train', 300, cfg_polyp['num_classes'])
net = ssd_net
if args.cuda:
torch.cuda.set_device(cfg_polyp['cuda_device'])
# net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
ssd_net.load_weights(args.save_folder + args.resume)
else:
# vgg_weights = torch.load(args.save_folder + args.basenet)
# model_dict = ssd_net.vgg.state_dict()
# model_dict.pop('0.bias')
# model_dict.pop('0.weight')
# vgg_weights = bbbuild('train', 300, cfg['num_classes'])
vgg_weights = torch.load(args.save_folder + args.basenet)
model_dict = net.state_dict()
vgg_weights = {k: v for k, v in vgg_weights.items() if k in model_dict}
# model_dict = ssd_net.vgg.state_dict()
model_dict.update(vgg_weights)
print('Loading base network...')
net.load_state_dict(model_dict)
if args.cuda:
net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
# ssd_net.vgg.apply(weights_init)
# ssd_net.extras.apply(weights_init)
# ssd_net.loc.apply(weights_init)
# ssd_net.conf.apply(weights_init)
# ssd_net.simi.apply(weights_init)
# ssd_net.simiSE.apply(weights_init)
# 只训练simi
# for param in net.parameters():
# param.requires_grad = False
# for param in net.simi.parameters():
# param.requires_grad = True
# for param in net.simiSE.parameters():
# param.requires_grad = True
# for param in net.loc.parameters():
# param.requires_grad = False
# for param in net.simi.parameters():
# param.requires_grad = False
# for param in net.simiSE.parameters():
# param.requires_grad = True
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr=args.lr, momentum=args.momentum,
# weight_decay=args.weight_decay)
# optimizer = optim.Adam(filter(lambda p: p.requires_grad,net.parameters()), lr=args.lr)
criterion = MultiBoxLoss(num_classes=cfg_polyp['num_classes'],
overlap_thresh=0.5,
prior_for_matching=True,
bkg_label=0,
neg_mining=True,
neg_pos=3,
neg_overlap=0.5,
encode_target=False,
use_gpu=args.cuda)
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
simi_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training SSD on: Gastric Anatomy Images')
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)
def cycle(iterable):
while True:
for x in iterable:
yield x
# create batch iterator
batch_iterator = iter(cycle(data_loader))
for iteration in range(args.start_iter, cfg_polyp['max_iter']):
if iteration in cfg_polyp['lr_steps']:
step_index += 1
if iteration >= 30000:
adjust_learning_rate(optimizer, args.gamma, step_index)
else:
adjust_learning_rate(optimizer, args.gamma * 100, step_index)
# load train data
images, targets = next(batch_iterator)
if args.cuda:
images = Variable(images.cuda())
# targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
with torch.no_grad():
targets = [Variable(ann.cuda()) 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.item()
conf_loss += loss_c.item()
# simi_loss +=loss_s.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) +
' || Loss: %.4f || Loss_c: %.4f || Loss_l: %.4f ' %
(loss.item(), loss_c.item(), loss_l.item()))
writer.add_scalars('polyp_origins_CVC_clinic/train_origin_2time', {
'loss_c': float(loss_c.item()),
'loss_l': float(loss_l.item())
}, iteration)
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(
ssd_net.state_dict(),
'AnatomyDetection/origins_2time_CVC_clinic_ssd300_polyp_' +
repr(iteration) + '.pth')
torch.save(
ssd_net.state_dict(),
'AnatomyDetection/origins_2time_CVC_clinic_ssd300_polyp_' +
repr(iteration) + '.pth')
writer.close()
def train():
if args.visdom:
import visdom
viz = visdom.Visdom()
print('Loading the dataset...')
if args.dataset == 'COCO':
if args.dataset_root == VOC_ROOT:
if not os.path.exists(COCOroot):
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 = COCOroot
cfg = coco_refinedet[args.input_size]
train_sets = [('train2017', 'val2017')]
dataset = COCODetection(COCOroot, train_sets, 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_refinedet[args.input_size]
dataset = VOCDetection(root=VOC_ROOT, transform=SSDAugmentation(cfg['min_dim'], MEANS))
print('Training RefineDet on:', dataset.name)
print('Using the specified args:')
print(args)
refinedet_net = build_refinedet('train', int(args.input_size), cfg['num_classes'], backbone_dict)
net = refinedet_net
print(net)
device = torch.device('cuda:0' if args.cuda else 'cpu')
if args.ngpu > 1 and args.cuda:
net = torch.nn.DataParallel(refinedet_net, device_ids=list(range(args.ngpu)))
cudnn.benchmark = True
net = net.to(device)
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
state_dict = torch.load(args.resume)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
refinedet_net.load_state_dict(new_state_dict)
else:
print('Initializing weights...')
refinedet_net.init_weights(pretrained=pretrained)
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
arm_criterion = RefineDetMultiBoxLoss( 2, 0.5, True, 0, True, negpos_ratio, 0.5,
False, args.cuda)
if wo_refined_anchor:
odm_criterion = MultiBoxLoss( cfg['num_classes'], 0.5, True, 0, True, negpos_ratio, 0.5,
False, args.cuda)
else:
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, negpos_ratio, 0.5,
False, args.cuda, use_ARM=True)
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
priors = priors.to(device)
net.train()
# loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
epoch = 0 + args.resume_epoch
epoch_size = math.ceil(len(dataset) / args.batch_size)
max_iter = args.max_epoch * epoch_size
stepvalues = (args.max_epoch * 2 // 3 * epoch_size, args.max_epoch * 8 // 9 * epoch_size, args.max_epoch * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
for step in stepvalues:
if step < start_iter:
step_index += 1
else:
start_iter = 0
if args.visdom:
vis_title = 'RefineDet.PyTorch on ' + dataset.name
vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
iter_plot = create_vis_plot(viz, 'Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot(viz, '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)
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
if args.visdom and iteration != 0:
update_vis_plot(viz, epoch, arm_loc_loss, arm_conf_loss, epoch_plot, None,
'append', epoch_size)
# reset epoch loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
# create batch iterator
batch_iterator = iter(data_loader)
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 ==0 and epoch > 200):
torch.save(net.state_dict(), args.save_folder+'RefineDet'+ args.input_size +'_'+args.dataset + '_epoches_'+
repr(epoch) + '.pth')
epoch += 1
t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.to(device)
targets = [ann.to(device) for ann in targets]
# for an in targets:
# for instance in an:
# for cor in instance[:-1]:
# if cor < 0 or cor > 1:
# raise StopIteration
# forward
out = net(images)
# backprop
optimizer.zero_grad()
if wo_refined_anchor:
arm_loss_l, arm_loss_c = torch.zeros(1), torch.zeros(1)
odm_loss_l, odm_loss_c = odm_criterion(out, priors, targets)
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
else:
arm_loss_l, arm_loss_c = arm_criterion(out, priors, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, priors, targets)
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
loss.backward()
optimizer.step()
arm_loc_loss += arm_loss_l.item()
arm_conf_loss += arm_loss_c.item()
odm_loc_loss += odm_loss_l.item()
odm_conf_loss += odm_loss_c.item()
t1 = time.time()
batch_time = t1 - t0
eta = int(batch_time * (max_iter - iteration))
print('Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || ARM_L Loss: {:.4f} ARM_C Loss: {:.4f} ODM_L Loss: {:.4f} ODM_C Loss: {:.4f} loss: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'.\
format(epoch, args.max_epoch, (iteration % epoch_size) + 1, epoch_size, iteration + 1, max_iter, arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item(), loss.item(), lr, batch_time, str(datetime.timedelta(seconds=eta))))
if args.visdom:
update_vis_plot(viz, iteration, arm_loss_l.item(), arm_loss_c.item(),
iter_plot, epoch_plot, 'append')
torch.save(refinedet_net.state_dict(), args.save_folder + '/RefineDet{}_{}_final.pth'.format(args.input_size, args.dataset))
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
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')
args.dataset_root = VOC_ROOT
cfg = voc
dataset = VOCDetection(root=args.dataset_root,
image_sets=[('2007', 'trainval'),
('2012', 'trainval')],
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif args.dataset == 'HOLO':
# if args.dataset_root == COCO_ROOT:
# parser.error('Must specify dataset if specifying dataset_root')
args.dataset_root = HOLO_ROOT
cfg = holo
dataset = HOLODetection(root=args.dataset_root,
image_sets='train',
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif args.dataset == 'HOLOV2':
# if args.dataset_root == COCO_ROOT:
# parser.error('Must specify dataset if specifying dataset_root')
args.dataset_root = HOLOV2_ROOT
cfg = holov2
dataset = HOLOV2Detection(root=args.dataset_root,
image_sets='train',
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif args.dataset == 'ICDAR2015': ####
args.dataset_root = ICDAR2015_ROOT
cfg = icdar2015
dataset = ICDAR2015Detection(root=args.dataset_root,
image_sets='train',
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif args.dataset == 'ICDAR2013': ####
args.dataset_root = ICDAR2013_ROOT
cfg = icdar2013
dataset = ICDAR2013Detection(root=args.dataset_root,
image_sets='train',
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif args.dataset == 'SynthText': ####
args.dataset_root = SynthText_ROOT
cfg = synthtext
dataset = SynthTextDetection(root=args.dataset_root,
image_sets='train',
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
if args.visdom:
from tensorboardX import SummaryWriter
tb_writer = SummaryWriter(log_dir='runs/Pelee_test')
pelee_net = PeleeNet('train', cfg) ####
net = pelee_net
# ##########
# pelee_net = PeleeNet('train', cfg)####
# pelee_net = pelee_net.cpu()
# import pdb
# pdb.set_trace()
# from ptflops import get_model_complexity_info
# flops, params = get_model_complexity_info(pelee_net, (304, 304), as_strings=True, print_per_layer_stat=True)
# print('Flops: ' + flops)
# print('Params: ' + params)
# pdb.set_trace()
# from thop import profile
# flops, params = profile(pelee_net, input_size=(1, 3, 304,304))
# print('flops',flops)
# print('params',params)
# exit()
###########
if args.cuda:
net = torch.nn.DataParallel(pelee_net)
cudnn.benchmark = True
###################
####random init first
print('Initializing weights...')
pelee_net.apply(weights_init) ####
# # initialize newly added layers' weights with xavier method
# pelee_net.extras.apply(weights_init)
# pelee_net.loc.apply(weights_init)
# pelee_net.conf.apply(weights_init)
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
pelee_net.load_weights(args.resume)
else:
pelee_weights = torch.load(args.save_folder + args.basenet)
print('Loading network except conf layers...')
pelee_net.load_state_dict(pelee_weights, strict=False)
if args.cuda:
net = net.cuda()
#############
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('Training Pelee 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 iteration in range(args.start_iter, cfg['max_iter']):
if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
epoch += 1 ####already moved it before update_vis_plot
# update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None, ####
# 'append', epoch_size)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
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 = 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
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 ####origin is loc_loss += loss_l.data[0]
conf_loss += loss_c.data ####origin is conf_loss += loss_c.data[0]
if iteration % 100 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data),
end=' ') ####loss.data[0] to loss.data
print(' \n ') ########
if args.visdom:
# update_vis_plot(iteration, loss_l.data, loss_c.data, ####loss_l.data[0] loss_c.data[0]
# iter_plot, epoch_plot, 'append')
tb_writer.add_scalar('runs/Pelee_test',
float(loss_l.data + loss_c.data), iteration)
if iteration != 0 and iteration % args.snapshot == 0:
print('Saving state, iter:', iteration)
torch.save(pelee_net.state_dict(),
'weights/Pelee_' + str(iteration) + '.pth')
os.system(
'python holov2_eval.py --trained_model weights/Pelee_' +
str(iteration) + '.pth')
# t2 = time.time()
# print('t012',t0,t1,t2)####
torch.save(pelee_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
os.system('python holov2_eval.py --trained_model' + ' ' +
args.save_folder + args.dataset + '.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))
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('[DEBUG] 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()
if not args.resume:
print('[DEBUG] 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)
sche = lr_scheduler.CosineAnnealingLR(optimizer, T_max=250)
# initSummaty()
log_folder = './results/' + net.__class__.__name__ + '/' + 'mixupCOCO' + '/' + str(1002) + '/'
print("log_folder: ", log_folder)
writer = SummaryWriter(log_folder)
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
dataset_len = epoch_size
print("[DEBUG] dataset len: {}".format(len(dataset)))
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)
# stepvalues = (150 , 200 , 250)
stepvalues = (80, 100, 120, 140, 160)
for epoch in range(args.start_iter, 160):
batch_iterator = iter(data_loader)
loc_loss = 0
conf_loss = 0
t0 = time.time()
# ---------------------------------
step_index = 0
for stepvalue in stepvalues:
if epoch >= stepvalue:
step_index += 1
# ---------------------------------
# sche.step(epoch)
for iteration in range(epoch_size):
lr = adjust_learning_rate_wr(optimizer, 0.1, step_index, iteration + (epoch_size * epoch), epoch, epoch_size)
# load train data
images, targets = next(batch_iterator)
t1 = targets[0]
print(t1)
# -------------------------------------------------------------------
print("# --------------------------------------------------------------");
exit(0)
### mixing
alpha = 0.1
lam = np.random.beta(alpha, alpha)
index = np.arange(len(targets))
np.random.shuffle(index)
if args.cuda:
images = Variable(images.cuda())
[Variable(anno.cuda(), volatile=True) for anno in targets]
# else:
# images = Variable(images)
# targets = [Variable(ann, volatile=True) for ann in targets]
### mixing
mixed_x = lam * images + (1 - lam) * images[index, :]
y_a= targets
y_b = [targets[idx] for idx in index]
# forward
out = net(mixed_x)
# backprop
optimizer.zero_grad()
# loss_l, loss_c = criterion(out, targets)
loss_l1, loss_c1 = criterion(out, y_a)
loss_l2, loss_c2 = criterion(out, y_b)
# loss_l = torch.sum(lam * loss_l1.sum() + (1 - lam) * loss_l2.sum())
# loss_c = torch.sum(lam * loss_c1.sum() + (1 - lam) * loss_c2.sum())
loss_l = lam * loss_l1 + (1 - lam) * loss_l2
loss_c = lam * loss_c1 + (1 - lam) * loss_c2
loss = loss_l + loss_c
loss.backward()
optimizer.step()
loc_loss += loss_l.data.item()
conf_loss += loss_c.data.item()
lr_now = optimizer.param_groups[0]['lr']
print('==>Train: Epoch [{}/{}] '.format(epoch, 400) + 'iter ' + repr(iteration) + '|| loss_l:%.4f | loss_c:%.4f || ' % (loss_l.data.item(),loss_c.data.item()) + 'lr={}'.format(lr_now), end='\r')
t1 = time.time()
print('\nEpoch [{}/{}] '.format(epoch, 400) + 'timer: %.4f sec.' % (t1 - t0) , end='\n')
writer.add_scalar('loc_loss', loc_loss/epoch_size, epoch)
writer.add_scalar('conf_loss', conf_loss/epoch_size, epoch)
lr_now = optimizer.param_groups[0]['lr']
writer.add_scalar('learning rate', lr_now, epoch)
# reset epoch loss counters
# if args.visdom:
# update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
# iter_plot, epoch_plot, 'append')
if epoch % 10 == 0 or epoch > 100 :
print('Saving state, epoch:', epoch)
torch.save(ssd_net.state_dict(), log_folder + 'ssd300_' + args.dataset + '_' +
repr(epoch) + '.pth')
torch.save(ssd_net.state_dict(), log_folder + 'ssd300_' + args.dataset + '_' +
repr(epoch) + '.pth')
def train(args, cfg, option, DataSet):
if args.exp_name is not None:
args.save_folder = os.path.join(args.save_folder, args.exp_name)
args.log_folder = os.path.join(args.log_folder, args.exp_name)
if not os.path.exists(args.save_folder):
os.makedirs(args.save_folder, exist_ok=True)
if not os.path.exists(args.log_folder):
os.makedirs(args.log_folder, exist_ok=True)
if True:
dataset = DataSet(image_path=cfg.dataset.train_images,
mask_out_ch=cfg.gt_inst_ch,
info_file=cfg.dataset.train_info,
option = cfg.dataset,
transform=SSDAugmentation(cfg, MEANS),
running_mode='train')
else:
dataset = DataSet(image_path=cfg.dataset.valid_images,
mask_out_ch=cfg.gt_inst_ch,
info_file=cfg.dataset.valid_info,
option = cfg.dataset,
transform=SSDAugmentation(cfg, MEANS),
running_mode='train')
# Parallel wraps the underlying module, but when saving and loading we don't want that
dvis_net = DVIS(cfg)
net = dvis_net
net.train()
if args.log:
log = Log(cfg.name, args.log_folder, dict(args._get_kwargs()),
overwrite=(args.resume is None), log_gpu_stats=args.log_gpu)
# I don't use the timer during training (I use a different timing method).
# Apparently there's a race condition with multiple GPUs, so disable it just to be safe.
timer.disable_all()
# Both of these can set args.resume to None, so do them before the check
if args.resume == 'interrupt':
args.resume = SavePath.get_interrupt(args.save_folder)
elif args.resume == 'latest':
args.resume = SavePath.get_latest(args.save_folder, cfg.name)
if args.resume is not None:
print('Resuming training, loading {}...'.format(args.resume))
dvis_net.load_weights(args.resume,
load_firstLayer=option['model_1stLayer_en'],
load_lastLayer=option['model_lastLayer_en'])
if args.start_iter == -1:
args.start_iter = SavePath.from_str(args.resume).iteration
else:
print('Initializing weights...')
dvis_net.init_weights(backbone_path=args.save_folder + cfg.backbone.path)
#optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
# weight_decay=args.decay)
optimizer = optim.SGD([{'params': net.backbone.parameters(), 'lr':args.lr*option['bkb_lr_alpha']},
{'params': net.fpn.parameters(), 'lr':args.lr*option['fpn_lr_alpha']},
{'params': net.proto_net.parameters(), 'lr':args.lr*option['proto_net_lr_alpha']}],
lr=args.lr, momentum=args.momentum, weight_decay=args.decay)
criterion = LossEvaluate(option,
class_weights=cfg.dataset.sem_weights)
if args.batch_alloc is not None:
args.batch_alloc = [int(x) for x in args.batch_alloc.split(',')]
if sum(args.batch_alloc) != args.batch_size:
print('Error: Batch allocation (%s) does not sum to batch size (%s).' % (args.batch_alloc, args.batch_size))
exit(-1)
net = NetLoss(net, criterion)
net = CustomDataParallel(net)
if args.cuda:
net = net.cuda()
# Initialize everything
if not cfg.freeze_bn:
dvis_net.freeze_bn() # Freeze bn so we don't kill our means
# loss counters
loc_loss = 0
conf_loss = 0
iteration = max(args.start_iter, 0)
last_time = time.time()
epoch_size = len(dataset) // args.batch_size
num_epochs = math.ceil(cfg.max_iter / epoch_size)
# Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
step_index = 0
data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=False, collate_fn=detection_collate,
pin_memory=True)
writer = SummaryWriter(log_dir=args.log_folder)
save_path = lambda epoch, iteration: SavePath(cfg.name, epoch, iteration).get_path(root=args.save_folder)
time_avg = MovingAverage()
loss_keys = ['binary', 'pi', 'l1', 'regul', 'iou', 'classify', 'eval_prec', 'eval_rec', 'eval_acc']
vis_keys = ['preds', 'gts', 'rgb', 'wghts', 'grad']
loss_avgs = { k: MovingAverage(100) for k in loss_keys }
print('Begin training!')
# try-except so you can use ctrl+c to save early and stop training
try:
log_loss = dict()
for epoch in range(num_epochs):
# Resume from start_iter
if (epoch+1)*epoch_size < iteration:
continue
for datum in data_loader:
# Stop if we've reached an epoch if we're resuming from start_iter
if iteration == (epoch+1)*epoch_size:
break
# Stop at the configured number of iterations even if mid-epoch
if iteration == cfg.max_iter:
break
if iteration < 99:
iteration += 1
continue
# Change a config setting if we've reached the specified iteration
changed = False
for change in cfg.delayed_settings:
if iteration >= change[0]:
changed = True
cfg.replace(change[1])
# Reset the loss averages because things might have changed
for avg in loss_avgs:
avg.reset()
# If a config setting was changed, remove it from the list so we don't keep checking
if changed:
cfg.delayed_settings = [x for x in cfg.delayed_settings if x[0] > iteration]
# Warm up by linearly interpolating the learning rate from some smaller value
if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until:
set_lr(optimizer, (args.lr - cfg.lr_warmup_init) * (iteration / cfg.lr_warmup_until) + cfg.lr_warmup_init)
# Adjust the learning rate at the given iterations, but also if we resume from past that iteration
while step_index < len(cfg.lr_steps) and iteration >= cfg.lr_steps[step_index]:
step_index += 1
set_lr(optimizer, args.lr * (args.gamma ** step_index))
# Zero the grad to get ready to compute gradients
optimizer.zero_grad()
# Forward Pass + Compute loss at the same time (see CustomDataParallel and NetLoss0)
ret = net(datum)
# Mean here because Dataparallel and do Backprop
losses = { k: ret[k].mean() for k in loss_keys if k in ret}
det_loss_keys = [k for k in loss_keys if k in losses]
all_loss = sum([losses[k] for k in det_loss_keys])
for k in det_loss_keys:
loss_avgs[k].add(losses[k].item())
# backward and optimize
if args.show_gradients==True:
ret['preds_0'].retain_grad()
all_loss.backward(retain_graph=True)
ret['grad'] = ret['preds_0'].grad[:, 0, :, :]
else:
all_loss.backward() # Do this to free up vram even if loss is not finite
if torch.isfinite(all_loss).item():
optimizer.step()
_, ret['preds'] = ret['preds'].max(axis=1, keepdim=True)
#ret['preds'] = torch.nn.Softmax2d()(ret['preds'])[:, :1, :, :]
vis_imgs = {k:ret[k] for k in vis_keys if k in ret}
cur_time = time.time()
elapsed = cur_time - last_time
last_time = cur_time
# Exclude graph setup from the timing information
if iteration != args.start_iter:
time_avg.add(elapsed)
if iteration % 10 == 0:
eta_str = str(datetime.timedelta(seconds=(cfg.max_iter-iteration) * time_avg.get_avg())).split('.')[0]
total = sum([loss_avgs[k].get_avg() for k in det_loss_keys if 'eval' not in k])
loss_labels = sum([[k, loss_avgs[k].get_avg()] for k in loss_keys if k in det_loss_keys], [])
print(('[%3d] %7d ||' + (' %s: %.3f |' * len(det_loss_keys)) + ' T: %.3f || ETA: %s || timer: %.3f')
% tuple([epoch, iteration] + loss_labels + [total, eta_str, elapsed]), flush=True)
if args.log:
log_step = 50//args.batch_size
for k in det_loss_keys:
if k not in log_loss:
log_loss[k] = loss_avgs[k].get_avg()
else:
log_loss[k] += loss_avgs[k].get_avg()
if iteration%log_step== log_step-1:
for k in det_loss_keys:
writer.add_scalar(k+'_loss',
log_loss[k]/float(log_step),
iteration/log_step)
log_loss[k] = 0
log_fig_step = 100
if iteration%log_fig_step == log_fig_step-1:
if 'davis' in args.dataset:
vis_imgs['rgb'] = vis_imgs['rgb'][:, :3, :, :]
fig = plot_tfboard_figure(cfg, vis_imgs, show_grad=args.show_gradients)
writer.add_figure('prediction _ grad', fig,
global_step=iteration/log_fig_step)
iteration += 1
if iteration % args.save_interval == 0 and iteration != args.start_iter:
if args.keep_latest:
latest = SavePath.get_latest(args.save_folder, cfg.name)
print('Saving state, iter:', iteration)
dvis_net.save_weights(save_path(epoch, iteration))
if args.keep_latest and latest is not None:
if args.keep_latest_interval <= 0 or iteration % args.keep_latest_interval != args.save_interval:
print('Deleting old save...')
os.remove(latest)
del ret, vis_imgs, losses
# end of batch run
# end of epoch
except KeyboardInterrupt:
if args.interrupt:
print('Stopping early. Saving network...')
# Delete previous copy of the interrupted network so we don't spam the weights folder
SavePath.remove_interrupt(args.save_folder)
writer.close()
dvis_net.save_weights(save_path(epoch, repr(iteration) + '_interrupt'))
exit()
writer.close()
dvis_net.save_weights(save_path(epoch, iteration))
normalize=args.normalize,
norm_value=255),
target_transform=kittiVOCAnnotationTransform())
elif args.dataset == 'kitti_voc_small':
if args.config == '300x300':
cfg = kitti300x300
elif args.config == '1000x300':
cfg = kitti1000x300
else:
raise ValueError('The given configuration is not possible')
dataset = kitti_small_Detection(root=kitti_small_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=KITTI_MEANS,
eval=True))
else:
raise (TypeError("config was not possible. check for typos"))
#test_image = "/home/marius/data/kitti_voc/JPEGImages/000000001_000468.jpg"
#test_label = "/home/marius/data/kitti_voc/Annotations/000000001_000468.xml"
#test_image = "/home/marius/data/kitti_voc/JPEGImages/000000032_007469.jpg"
#test_label = "/home/marius/data/kitti_voc/Annotations/000000032_007469.xml"
#test_image_toy_data = "/home/marius/data/toy_data/JPEGImages/0010000.jpg"
#test_label_toy_data = "/home/marius/data/toy_data/Annotations/0010000.xml"
#show_label_annotations(test_image_toy_data, test_label_toy_data, save_file=False)
#cv2.imwrite('Images/test_annotation.png',pred)
def main():
global my_dict, keys, k_len, arr, xxx, args, log_file, best_prec1
relative_path = '/data4/lilin/my_code'
parser = argparse.ArgumentParser(
description='Single Shot MultiBox Detector Training')
parser.add_argument('--version',
default='v2',
help='conv11_2(v2) or pool6(v1) as last layer')
parser.add_argument('--basenet',
default='vgg16_reducedfc.pth',
help='pretrained base model')
parser.add_argument('--dataset',
default='ucf24',
help='pretrained base model')
parser.add_argument('--ssd_dim',
default=300,
type=int,
help='Input Size for SSD') # only support 300 now
parser.add_argument(
'--modality',
default='rgb',
type=str,
help='INput tyep default rgb options are [rgb,brox,fastOF]')
parser.add_argument('--jaccard_threshold',
default=0.5,
type=float,
help='Min Jaccard index for matching')
parser.add_argument('--batch_size',
default=32,
type=int,
help='Batch size for training')
parser.add_argument('--num_workers',
default=0,
type=int,
help='Number of workers used in dataloading')
parser.add_argument('--max_iter',
default=120000,
type=int,
help='Number of training iterations')
parser.add_argument('--man_seed',
default=123,
type=int,
help='manualseed for reproduction')
parser.add_argument('--cuda',
default=True,
type=str2bool,
help='Use cuda to train model')
parser.add_argument('--ngpu',
default=1,
type=str2bool,
help='Use cuda to train model')
parser.add_argument('--base_lr',
default=0.001,
type=float,
help='initial learning rate')
parser.add_argument('--lr',
default=0.001,
type=float,
help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--weight_decay',
default=5e-4,
type=float,
help='Weight decay for SGD')
parser.add_argument('--gamma',
default=0.2,
type=float,
help='Gamma update for SGD')
parser.add_argument('--log_iters',
default=True,
type=bool,
help='Print the loss at each iteration')
parser.add_argument('--visdom',
default=False,
type=str2bool,
help='Use visdom to for loss visualization')
parser.add_argument('--data_root',
default=relative_path + '/realtime/',
help='Location of VOC root directory')
parser.add_argument('--save_root',
default=relative_path + '/realtime/saveucf24/',
help='Location to save checkpoint models')
parser.add_argument('--iou_thresh',
default=0.5,
type=float,
help='Evaluation threshold')
parser.add_argument('--conf_thresh',
default=0.01,
type=float,
help='Confidence threshold for evaluation')
parser.add_argument('--nms_thresh',
default=0.45,
type=float,
help='NMS threshold')
parser.add_argument('--topk',
default=50,
type=int,
help='topk for evaluation')
parser.add_argument('--clip_gradient',
default=40,
type=float,
help='gradients clip')
parser.add_argument('--resume',
default=None,
type=str,
help='Resume from checkpoint')
parser.add_argument('--start_epoch',
default=0,
type=int,
help='start epoch')
parser.add_argument('--epochs',
default=35,
type=int,
metavar='N',
help='number of total epochs to run')
parser.add_argument('--eval_freq',
default=2,
type=int,
metavar='N',
help='evaluation frequency (default: 5)')
parser.add_argument('--snapshot_pref',
type=str,
default="ucf101_vgg16_ssd300_end2end")
parser.add_argument('--lr_milestones',
default=[-2, -4],
type=float,
help='initial learning rate')
parser.add_argument('--arch', type=str, default="VGG16")
parser.add_argument('--Finetune_SSD', default=False, type=str)
parser.add_argument('-e',
'--evaluate',
dest='evaluate',
action='store_true',
help='evaluate model on validation set')
parser.add_argument(
'--step',
type=int,
default=[18, 27],
nargs='+',
help='the epoch where optimizer reduce the learning rate')
parser.add_argument('--log_lr',
default=True,
type=str2bool,
help='Use cuda to train model')
parser.add_argument('--print-log',
type=str2bool,
default=True,
help='print logging or not')
parser.add_argument('--end2end',
type=str2bool,
default=False,
help='print logging or not')
## Parse arguments
args = parser.parse_args()
print(__file__)
file_name = (__file__).split('/')[-1]
file_name = file_name.split('.')[0]
print_log(args, file_name)
## set random seeds
np.random.seed(args.man_seed)
torch.manual_seed(args.man_seed)
if args.cuda:
torch.cuda.manual_seed_all(args.man_seed)
if args.cuda and torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
args.cfg = v2
args.train_sets = 'train'
args.means = (104, 117, 123)
num_classes = len(CLASSES) + 1
args.num_classes = num_classes
# args.step = [int(val) for val in args.step.split(',')]
args.loss_reset_step = 30
args.eval_step = 10000
args.print_step = 10
args.data_root += args.dataset + '/'
## Define the experiment Name will used to same directory
args.snapshot_pref = ('ucf101_CONV-SSD-{}-{}-bs-{}-{}-lr-{:05d}').format(
args.dataset, args.modality, args.batch_size, args.basenet[:-14],
int(args.lr * 100000)) # + '_' + file_name + '_' + day
print_log(args, args.snapshot_pref)
if not os.path.isdir(args.save_root):
os.makedirs(args.save_root)
net = build_ssd(300, args.num_classes)
if args.Finetune_SSD is True:
print_log(args, "load snapshot")
pretrained_weights = "/home2/lin_li/zjg_code/realtime/ucf24/rgb-ssd300_ucf24_120000.pth"
pretrained_dict = torch.load(pretrained_weights)
model_dict = net.state_dict() # 1. filter out unnecessary keys
pretrained_dict_2 = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
} # 2. overwrite entries in the existing state dict
# pretrained_dict_2['vgg.25.bias'] = pretrained_dict['vgg.24.bias']
# pretrained_dict_2['vgg.25.weight'] = pretrained_dict['vgg.24.weight']
# pretrained_dict_2['vgg.27.bias'] = pretrained_dict['vgg.26.bias']
# pretrained_dict_2['vgg.27.weight'] = pretrained_dict['vgg.26.weight']
# pretrained_dict_2['vgg.29.bias'] = pretrained_dict['vgg.28.bias']
# pretrained_dict_2['vgg.29.weight'] = pretrained_dict['vgg.28.weight']
# pretrained_dict_2['vgg.32.bias'] = pretrained_dict['vgg.31.bias']
# pretrained_dict_2['vgg.32.weight'] = pretrained_dict['vgg.31.weight']
# pretrained_dict_2['vgg.34.bias'] = pretrained_dict['vgg.33.bias']
# pretrained_dict_2['vgg.34.weight'] = pretrained_dict['vgg.33.weight']
model_dict.update(pretrained_dict_2) # 3. load the new state dict
elif args.resume is not None:
if os.path.isfile(args.resume):
print_log(args, ("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
if args.end2end is False:
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
net.load_state_dict(checkpoint['state_dict'])
print_log(args, ("=> loaded checkpoint '{}' (epoch {})".format(
args.evaluate, checkpoint['epoch'])))
else:
print_log(args,
("=> no checkpoint found at '{}'".format(args.resume)))
elif args.modality == 'fastOF':
print_log(
args,
'Download pretrained brox flow trained model weights and place them at:::=> '
+ args.data_root + 'ucf24/train_data/brox_wieghts.pth')
pretrained_weights = args.data_root + 'train_data/brox_wieghts.pth'
print_log(args, 'Loading base network...')
net.load_state_dict(torch.load(pretrained_weights))
else:
vgg_weights = torch.load(args.data_root + 'train_data/' + args.basenet)
print_log(args, 'Loading base network...')
net.vgg.load_state_dict(vgg_weights)
if args.cuda:
net = net.cuda()
def xavier(param):
init.xavier_uniform(param)
def weights_init(m):
if isinstance(m, nn.Conv2d):
xavier(m.weight.data)
m.bias.data.zero_()
print_log(args, 'Initializing weights for extra layers and HEADs...')
# initialize newly added layers' weights with xavier method
if args.Finetune_SSD is False and args.resume is None:
print_log(args, "init layers")
net.extras.apply(weights_init)
net.loc.apply(weights_init)
net.conf.apply(weights_init)
parameter_dict = dict(net.named_parameters(
)) # Get parmeter of network in dictionary format wtih name being key
params = []
#Set different learning rate to bias layers and set their weight_decay to 0
for name, param in parameter_dict.items():
if args.end2end is False and name.find('vgg') > -1 and int(
name.split('.')[1]) < 23: # :and name.find('cell') <= -1
param.requires_grad = False
print_log(args, name + 'layer parameters will be fixed')
else:
if name.find('bias') > -1:
print_log(
args,
name + 'layer parameters will be trained @ {}'.format(
args.lr * 2))
params += [{
'params': [param],
'lr': args.lr * 2,
'weight_decay': 0
}]
else:
print_log(
args, name +
'layer parameters will be trained @ {}'.format(args.lr))
params += [{
'params': [param],
'lr': args.lr,
'weight_decay': args.weight_decay
}]
optimizer = optim.SGD(params,
lr=args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion = MultiBoxLoss(args.num_classes, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
scheduler = None
# scheduler = MultiStepLR(optimizer, milestones=args.step, gamma=args.gamma)
print_log(args, 'Loading Dataset...')
train_dataset = UCF24Detection(args.data_root,
args.train_sets,
SSDAugmentation(args.ssd_dim, args.means),
AnnotationTransform(),
input_type=args.modality)
val_dataset = UCF24Detection(args.data_root,
'test',
BaseTransform(args.ssd_dim, args.means),
AnnotationTransform(),
input_type=args.modality,
full_test=False)
train_data_loader = data.DataLoader(train_dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
val_data_loader = data.DataLoader(val_dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=False,
collate_fn=detection_collate,
pin_memory=True)
print_log(args, "train epoch_size: " + str(len(train_data_loader)))
print_log(args, 'Training SSD on' + train_dataset.name)
my_dict = copy.deepcopy(train_data_loader.dataset.train_vid_frame)
keys = list(my_dict.keys())
k_len = len(keys)
arr = np.arange(k_len)
xxx = copy.deepcopy(train_data_loader.dataset.ids)
# log_file = open(args.save_root + args.snapshot_pref + "_training_" + day + ".log", "w", 1)
# log_file.write()
print_log(args, args.snapshot_pref)
for arg in vars(args):
print(arg, getattr(args, arg))
print_log(args, str(arg) + ': ' + str(getattr(args, arg)))
print_log(args, str(net))
torch.cuda.synchronize()
for epoch in range(args.start_epoch, args.epochs):
train(train_data_loader, net, criterion, optimizer, epoch, scheduler)
print_log(args, 'Saving state, epoch:' + str(epoch))
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': net.state_dict(),
'best_prec1': best_prec1,
},
epoch=epoch)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
torch.cuda.synchronize()
tvs = time.perf_counter()
mAP, ap_all, ap_strs = validate(args,
net,
val_data_loader,
val_dataset,
epoch,
iou_thresh=args.iou_thresh)
# remember best prec@1 and save checkpoint
is_best = mAP > best_prec1
best_prec1 = max(mAP, best_prec1)
print_log(args, 'Saving state, epoch:' + str(epoch))
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': net.state_dict(),
'best_prec1': best_prec1,
}, is_best, epoch)
for ap_str in ap_strs:
# print(ap_str)
print_log(args, ap_str)
ptr_str = '\nMEANAP:::=>' + str(mAP)
# print(ptr_str)
# log_file.write()
print_log(args, ptr_str)
torch.cuda.synchronize()
t0 = time.perf_counter()
prt_str = '\nValidation TIME::: {:0.3f}\n\n'.format(t0 - tvs)
# print(prt_str)
# log_file.write(ptr_str)
print_log(args, ptr_str)
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['dim_x'],
# cfg['dim_y'],
# MEANS))
if args.dataset == 'VOC':
cfg = voc
dataset = VOCDetection(root=VOC_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=VOC_MEANS,
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'VOC_small':
cfg = voc_small
dataset = VOCDetection(root=VOC_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=VOC_MEANS,
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'toy_data':
cfg = toy_data
dataset = toydataDetection(root=toy_data_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=(0, 0, 0),
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'rotation_data':
cfg = toy_data
dataset = rotationdataDetection(root=rotation_data_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=(0, 0, 0),
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'scale_data':
cfg = toy_data
dataset = scaledataDetection(root=scale_data_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=(0, 0, 0),
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'deformation_data':
cfg = toy_data
dataset = deformationdataDetection(root=deformation_data_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=(0, 0, 0),
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'translation_data':
cfg = toy_data
dataset = translationdataDetection(root=translation_data_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=(0, 0, 0),
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'toy_data_small':
cfg = toy_data_small
dataset = toydatasmallDetection(root=toy_data_small_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=(0, 0, 0),
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'kitti_voc':
if args.config == '300x300':
cfg = kitti300x300
elif args.config == '1000x300':
cfg = kitti1000x300
else:
raise ValueError('The given configuration is not possible')
dataset = kittiDetection(root=kitti_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=KITTI_MEANS,
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
elif args.dataset == 'kitti_voc_small':
if args.config == '300x300':
cfg = kitti300x300_small
elif args.config == '1000x300':
cfg = kitti1000x300_small
else:
raise ValueError('The given configuration is not possible')
dataset = kitti_small_Detection(root=kitti_small_ROOT,
transform=SSDAugmentation(
size_x=cfg['dim_x'],
size_y=cfg['dim_y'],
mean=KITTI_MEANS,
random=args.random,
normalize=args.normalize,
sub_mean=args.subtract_mean))
print("time after config: ", datetime.now().time())
print("config: ", cfg)
if args.visdom:
import visdom
viz = visdom.Visdom()
else:
ssd_net = build_ssd('train', cfg['dim_x'], cfg['dim_y'],
cfg['num_classes'], cfg, args.batch_norm)
print("ssd_net: ", ssd_net)
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)
elif args.pretrained_weights:
if args.batch_norm:
vgg_weights = torch.load(args.save_folder + args.basenet_bn)
else:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_weights)
#else: initizialize with xavier and do not use pretrained 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=cfg['lr'],
momentum=args.momentum, #replaced args.lr
weight_decay=args.weight_decay)
criterion = MultiBoxLoss(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
cfg=cfg,
use_gpu=args.cuda)
print("time after building: ", datetime.now().time())
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
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)
print("time after loading data: ", datetime.now().time())
# create batch iterator
batch_iterator = iter(data_loader)
for iteration in range(args.start_iter, cfg['max_iter'] + 100):
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
lr_ = cfg['lr_bn'] if args.batch_norm else cfg['lr']
adjust_learning_rate(optimizer, args.gamma, step_index, lr_in=lr_)
# 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 = 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.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 args.visdom:
update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % cfg['max_iter'] == 0:
print('Saving state, iter, file:', iteration, WEIGHTS_NAME)
torch.save(ssd_net.state_dict(),
'weights/' + WEIGHTS_NAME + repr(iteration) + '.pth')
torch.save(ssd_net.state_dict(),
args.save_folder + '' + args.dataset + '.pth')
def train():
# 选择不同的超参数配置和文件结构， 构建不同的dataset
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=COCO_ROOT,
transform=SSDAugmentation(
cfg['min_dim'], cfg['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=VOC_ROOT,
transform=SSDAugmentation(
cfg['min_dim'], cfg['means']))
elif args.dataset == 'VisDrone2018':
cfg = visdrone # 选择哪一个config
dataset = DroneDetection(root=DRONE_ROOT,
transform=SSDAugmentation(
cfg['min_dim'], cfg['means']))
# if args.visdom:
# import visdom
# viz = visdom.Visdom()
print('num_classes: ' + str(cfg['num_classes']))
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()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
ssd_net.vgg.apply(weights_init)
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) # L2 penalty
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('Training SSD on:', dataset.name)
print('Using the specified args:')
print(args)
# print(args.dataset)
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)
epoch_plot2 = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
if args.tensorboard:
logger = Logger('./logs')
# 初始化文件夹
with open('trainlogs.txt', 'w') as f:
f.write('Start training on {}'.format(args.dataset))
shutil.rmtree('args/')
shutil.rmtree('logs/')
os.mkdir('args/')
os.mkdir('logs/')
imgcnt = 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 size个图片
batch_iterator = iter(data_loader)
for iteration in range(args.start_iter, cfg['max_iter']):
# print('it: '+str(iteration))
if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot,
epoch_plot2, '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, 取数据
# 循环一次之后iter无法回到起点，需要重新赋值
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
# images, targets = next(batch_iterator)
# print('feed size')
# print(images.shape)
# for item in targets:
# print(item.shape)
if args.cuda:
images = Variable(images.cuda())
with torch.no_grad():
targets = [torch.Tensor(ann.cuda()) for ann in targets]
else:
images = Variable(images)
targets = [torch.Tensor(ann) for ann in targets]
# forward
t0 = time.time()
# optimizer.zero_grad()
# # output img
# with torch.no_grad():
# imgtensor=torch.Tensor(images)
# for img in imgtensor:
# imgnp=np.array(img.cpu().permute(1,2,0))
# rgbimg=imgnp[:, :, (2, 1, 0)]
# cv2.imwrite('trainimg/{}_{}.jpg'.format(args.dataset, imgcnt), rgbimg)
# imgcnt+=1
out = net(images)
# backprop
optimizer.zero_grad() # 写在计算新的梯度之前就可以backward之前
loss_l, loss_c = criterion(out, targets) # 对比network output和gt
loss = loss_l + loss_c
# print('loss: '+str(loss_l.data)+' '+str(loss_c.data))
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.data
conf_loss += loss_c.data
if iteration % 10 == 0:
# print('timer: %.4f sec.' % (t1 - t0))
# print('ok')
print('iter [{}/{}]'.format(iteration, cfg['max_iter'] -
args.start_iter) + ' || Loss: %.4f' %
(loss.data))
# # 打印参数
# with open('args/args_{}.txt'.format(iteration), 'a') as f:
# for item in net.named_parameters():
# f.write(' '+str(item[0])+': '+str(item[1]))
with open('trainlogs.txt', 'a') as f:
f.write('iter [{}/{}]'.format(
iteration, cfg['max_iter'] - args.start_iter) +
' || Loss: %.4f \n' % (loss.data))
if args.tensorboard:
info = {'loss': loss.data}
for tag, value in info.items():
logger.scalar_summary(tag, value, iteration)
for tag, value in net.named_parameters():
# print('tag: ' + str(tag))
# print('params: ' + str(value.data.cpu().numpy().shape)) # convert to cpu data and transform to numpy
tag = tag.replace('.', '/')
logger.histo_summary(tag,
value.data.cpu().numpy(), iteration)
logger.histo_summary(tag + '/grad',
value.grad.data.cpu().numpy(),
iteration)
info = {
'images':
images.view(-1, int(cfg['min_dim']),
int(cfg['min_dim'])).cpu().numpy()
}
for tag, img in info.items():
logger.image_summary(tag, img, iteration)
# print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data), end=' ')
# print(loss)
# print('isnan'+str(torch.isnan(loss)))
# print(torch.isnan(loss))
# # 检测loss爆炸
# if torch.isnan(loss).data!=0:
# print('Error')
# errorcnt=1
# with open('trainlogs.txt', 'a') as f:
# f.write('ERROR')
# # for img in images:
# # cv2.imwrite('./logs/'+str(errorcnt)+'.jpg', img)
# break
if args.visdom:
update_vis_plot(iteration, loss_l.data, loss_c.data, iter_plot,
epoch_plot, 'append')
if iteration != 0 and iteration % 1000 == 0:
print('Saving state, iter:', iteration)
torch.save(
ssd_net.state_dict(), 'weights/ssd' + str(cfg['min_dim']) +
'_' + str(args.dataset) + '_' + repr(iteration) + '.pth')
with open('trainlogs.txt', 'a') as f:
f.write('Saving state, iter:' + str(iteration))
torch.save(
ssd_net.state_dict(), args.save_folder + '' + 'SSD' +
str(cfg['min_dim']) + '_' + args.dataset + '.pth')
def main():
global my_dict, keys, k_len, arr, xxx, args, log_file, best_prec1
parser = argparse.ArgumentParser(description='Single Shot MultiBox Detector Training')
parser.add_argument('--version', default='v2', help='conv11_2(v2) or pool6(v1) as last layer')
parser.add_argument('--basenet', default='vgg16_reducedfc.pth', help='pretrained base model')
parser.add_argument('--dataset', default='ucf24', help='pretrained base model')
parser.add_argument('--ssd_dim', default=300, type=int, help='Input Size for SSD') # only support 300 now
parser.add_argument('--modality', default='rgb', type=str,
help='INput tyep default rgb options are [rgb,brox,fastOF]')
parser.add_argument('--jaccard_threshold', default=0.5, type=float, help='Min Jaccard index for matching')
parser.add_argument('--batch_size', default=32, type=int, help='Batch size for training')
parser.add_argument('--num_workers', default=0, type=int, help='Number of workers used in dataloading')
parser.add_argument('--max_iter', default=120000, type=int, help='Number of training iterations')
parser.add_argument('--man_seed', default=123, type=int, help='manualseed for reproduction')
parser.add_argument('--cuda', default=True, type=str2bool, help='Use cuda to train model')
parser.add_argument('--ngpu', default=1, type=str2bool, help='Use cuda to train model')
parser.add_argument('--lr', '--learning-rate', default=0.0005, type=float, help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--stepvalues', default='70000,90000', type=str,
help='iter number when learning rate to be dropped')
parser.add_argument('--weight_decay', default=5e-4, type=float, help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.2, type=float, help='Gamma update for SGD')
parser.add_argument('--log_iters', default=True, type=bool, help='Print the loss at each iteration')
parser.add_argument('--visdom', default=False, type=str2bool, help='Use visdom to for loss visualization')
parser.add_argument('--data_root', default='/data4/lilin/my_code/realtime/', help='Location of VOC root directory')
parser.add_argument('--save_root', default='/data4/lilin/my_code/realtime/realtime-lstm/save',
help='Location to save checkpoint models')
parser.add_argument('--iou_thresh', default=0.5, type=float, help='Evaluation threshold')
parser.add_argument('--conf_thresh', default=0.01, type=float, help='Confidence threshold for evaluation')
parser.add_argument('--nms_thresh', default=0.45, type=float, help='NMS threshold')
parser.add_argument('--topk', default=50, type=int, help='topk for evaluation')
parser.add_argument('--clip', default=40, type=float, help='gradients clip')
# parser.add_argument('--resume', default="/data4/lilin/my_code/realtime/realtime-lstm/saveucf24/cache/CONV-SSD-ucf24-rgb-bs-32-vgg16-lr-00050/ssd300_ucf24_30000.pth",
# type=str, help='Resume from checkpoint')
parser.add_argument('--resume', default=None,
type=str, help='Resume from checkpoint')
parser.add_argument('--start_epoch', default=0, type=int, help='start epoch')
parser.add_argument('--epochs', default=35, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--eval_freq', default=2, type=int, metavar='N', help='evaluation frequency (default: 5)')
parser.add_argument('--snapshot_pref', type=str, default="ucf101_vgg16_ssd300_")
print(__file__)
file_name = (__file__).split('/')[-1]
file_name = file_name.split('.')[0]
print(file_name)
## Parse arguments
args = parser.parse_args()
## set random seeds
np.random.seed(args.man_seed)
torch.manual_seed(args.man_seed)
if args.cuda:
torch.cuda.manual_seed_all(args.man_seed)
if args.cuda and torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
args.cfg = v2
args.train_sets = 'train'
args.means = (104, 117, 123)
num_classes = len(CLASSES) + 1
args.num_classes = num_classes
args.stepvalues = [int(val) for val in args.stepvalues.split(',')]
args.loss_reset_step = 30
args.eval_step = 10000
args.print_step = 10
## Define the experiment Name will used to same directory and ENV for visdom
args.exp_name = 'CONV-SSD-{}-{}-bs-{}-{}-lr-{:05d}'.format(args.dataset,
args.modality, args.batch_size, args.basenet[:-14], int(args.lr*100000))
args.save_root += args.dataset+'/'
args.save_root = args.save_root+'cache/'+args.exp_name+'/'
if not os.path.isdir(args.save_root):
os.makedirs(args.save_root)
net = build_ssd(300, args.num_classes)
# if args.has_snapshot is True:
# print ("load snapshot")
# pretrained_weights = "/data4/lilin/my_code/realtime/realtime-lstm/saveucf24/cache/CONV-SSD-ucf24-rgb-bs-32-vgg16-lr-00050/ssd300_ucf24_30000.pth"
# net.load_state_dict(torch.load(pretrained_weights))
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
net.load_state_dict(checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})"
.format(args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
elif args.modality == 'fastOF':
print('Download pretrained brox flow trained model weights and place them at:::=> ',args.data_root + 'ucf24/train_data/brox_wieghts.pth')
pretrained_weights = args.data_root + 'ucf24/train_data/brox_wieghts.pth'
print('Loading base network...')
net.load_state_dict(torch.load(pretrained_weights))
else:
vgg_weights = torch.load(args.data_root +'ucf24/train_data/' + args.basenet)
print('Loading base network...')
net.vgg.load_state_dict(vgg_weights)
args.data_root += args.dataset + '/'
if args.cuda:
net = net.cuda()
def xavier(param):
init.xavier_uniform(param)
def weights_init(m):
if isinstance(m, nn.Conv2d):
xavier(m.weight.data)
m.bias.data.zero_()
print('Initializing weights for extra layers and HEADs...')
# initialize newly added layers' weights with xavier method
if args.resume is None:
net.extras.apply(weights_init)
net.loc.apply(weights_init)
net.conf.apply(weights_init)
parameter_dict = dict(net.named_parameters()) # Get parmeter of network in dictionary format wtih name being key
params = []
#Set different learning rate to bias layers and set their weight_decay to 0
for name, param in parameter_dict.items():
if name.find('bias') > -1:
print(name, 'layer parameters will be trained @ {}'.format(args.lr*2))
params += [{'params': [param], 'lr': args.lr*2, 'weight_decay': 0}]
else:
print(name, 'layer parameters will be trained @ {}'.format(args.lr))
params += [{'params':[param], 'lr': args.lr, 'weight_decay':args.weight_decay}]
optimizer = optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
criterion = MultiBoxLoss(args.num_classes, 0.5, True, 0, True, 3, 0.5, False, args.cuda)
scheduler = None
scheduler = MultiStepLR(optimizer, milestones=args.stepvalues, gamma=args.gamma)
print('Loading Dataset...')
train_dataset = UCF24Detection(args.data_root, args.train_sets, SSDAugmentation(args.ssd_dim, args.means),
AnnotationTransform(), input_type=args.modality)
val_dataset = UCF24Detection(args.data_root, 'test', BaseTransform(args.ssd_dim, args.means),
AnnotationTransform(), input_type=args.modality,
full_test=False)
args.epoch_size = len(train_dataset) // args.batch_size
train_data_loader = data.DataLoader(train_dataset, args.batch_size, num_workers=args.num_workers,
shuffle=False, collate_fn=detection_collate, pin_memory=True)
val_data_loader = data.DataLoader(val_dataset, args.batch_size, num_workers=args.num_workers,
shuffle=False, collate_fn=detection_collate, pin_memory=True)
print ("epoch_size: ", args.epoch_size)
print('Training SSD on', train_dataset.name)
my_dict = copy.deepcopy(train_data_loader.dataset.train_vid_frame)
keys = list(my_dict.keys())
k_len = len(keys)
arr = np.arange(k_len)
xxx = copy.deepcopy(train_data_loader.dataset.ids)
log_file = open(args.save_root+"training.log", "w", 1)
log_file.write(args.exp_name+'\n')
for arg in vars(args):
print(arg, getattr(args, arg))
log_file.write(str(arg)+': '+str(getattr(args, arg))+'\n')
log_file.write(str(net))
torch.cuda.synchronize()
for epoch in range(args.start_epoch, args.epochs):
train(train_data_loader, net, criterion, optimizer, scheduler, epoch)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
torch.cuda.synchronize()
tvs = time.perf_counter()
mAP, ap_all, ap_strs = validate(args, net, val_data_loader, val_dataset, epoch, iou_thresh=args.iou_thresh)
# remember best prec@1 and save checkpoint
is_best = mAP > best_prec1
best_prec1 = max(mAP, best_prec1)
print('Saving state, epoch:', epoch)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': net.state_dict(),
'best_prec1': best_prec1,
}, is_best)
for ap_str in ap_strs:
print(ap_str)
log_file.write(ap_str+'\n')
ptr_str = '\nMEANAP:::=>'+str(mAP)+'\n'
print(ptr_str)
log_file.write(ptr_str)
torch.cuda.synchronize()
t0 = time.perf_counter()
prt_str = '\nValidation TIME::: {:0.3f}\n\n'.format(t0-tvs)
print(prt_str)
log_file.write(ptr_str)
log_file.close()
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 # coco位于config.py文件中
# COCODetection类 位于coco.py文件中
# SSDAugmentation类 位于utils/augmentations.py文件中
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))
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 # 这个flag可以让内置的cuDNN的auto-tuner自动寻找最适合当前配置的算法.
# resume 类型为 str, 值为checkpoint state_dict file
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.load_state_dict(vgg_weights)
# 将所有的参数都移送到GPU内存中
if args.cuda:
net = net.cuda()
# 用xavier方法初始化新添加层的权重
if not args.resume:
print('Initializing weights...')
ssd_net.extras.apply(weights_init)
ssd_net.loc.apply(weights_init)
ssd_net.conf.apply(weights_init)
# 随机梯度下降优化，计算梯度和误差并更新参数
# SGD在学习中增加了噪声，有正则化的效果
# 学习效率进行线性衰减可以保证SGD收敛
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# MultiBoxLoss类 位于layers/modules/multibox_loss.py文件中
criterion = MultiBoxLoss(cfg["num_classes"], 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
net.train()
# Loss计数器
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset)
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)
# 创建batch迭代器
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)
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)
# 加载训练数据
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_loc, loss_conf = criterion(out, targets)
loss = loss_loc + loss_conf
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_loc.data[0]
conf_loss += loss_conf.data[0]
# 每隔10次迭代就输出一次训练状态信息
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' ||Loss: %.4f ||' %
(loss.data[0]),
end=' ')
if args.visdom:
update_vis_plot(iteration, loss_loc.data[0], loss_conf.data[0],
iter_plot, epoch_plot, 'append')
# 迭代多少次保存一个模型以及保存的文件名
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(),
args.save_folder + '' + args.dataset + '.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 == 'VOC300':
if args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')
cfg = voc300
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
elif args.dataset == 'VOC512':
if args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')
cfg = voc512
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
if args.visdom:
import visdom
viz = visdom.Visdom()
finish_flag = True
while (finish_flag):
ssd_net = build_ssd_con('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)
# ssd_net.vgg_t.load_state_dict(vgg_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)
csd_criterion = CSDLoss(args.cuda)
isd_criterion = ISDLoss(args.cuda)
conf_consistency_criterion = torch.nn.KLDivLoss(size_average=False,
reduce=False).cuda()
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
supervised_flag = 1
print('Loading the dataset...')
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)
total_un_iter_num = 0
supervised_batch = args.batch_size
#unsupervised_batch = args.batch_size - supervised_batch
#data_shuffle = 0
if (args.start_iter == 0):
dataset = VOCDetection_con_init(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
else:
supervised_flag = 0
dataset = VOCDetection_con(
root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS)) #,shuffle_flag=data_shuffle)
#data_shuffle = 1
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True,
drop_last=True)
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)
try:
images, targets, semis = next(batch_iterator)
except StopIteration:
supervised_flag = 0
dataset = VOCDetection_con(
root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS)) #, shuffle_flag=data_shuffle)
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True,
drop_last=True)
batch_iterator = iter(data_loader)
images, targets, semis = 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()
images_flip = images.clone()
images_flip = flip(images_flip, 3)
images_shuffle = images_flip.clone()
images_shuffle[:int(args.batch_size / 2), :, :, :] = images_flip[
int(args.batch_size / 2):, :, :, :]
images_shuffle[
int(args.batch_size /
2):, :, :, :] = images_flip[:int(args.batch_size /
2), :, :, :]
lam = np.random.beta(5.0, 5.0)
images_mix = lam * images.clone() + (1 -
lam) * images_shuffle.clone()
out, conf, conf_flip, loc, loc_flip, conf_shuffle, conf_interpolation, loc_shuffle, loc_interpolation = net(
images, images_flip, images_mix)
sup_image_binary_index = np.zeros([len(semis), 1])
for super_image in range(len(semis)):
if (int(semis[super_image]) == 1):
sup_image_binary_index[super_image] = 1
else:
sup_image_binary_index[super_image] = 0
if (int(semis[len(semis) - 1 - super_image]) == 0):
del targets[len(semis) - 1 - super_image]
sup_image_index = np.where(sup_image_binary_index == 1)[0]
unsup_image_index = np.where(sup_image_binary_index == 0)[0]
loc_data, conf_data, priors = out
if (len(sup_image_index) != 0):
loc_data = loc_data[sup_image_index, :, :]
conf_data = conf_data[sup_image_index, :, :]
output = (loc_data, conf_data, priors)
# backprop
# loss = Variable(torch.cuda.FloatTensor([0]))
loss_l = Variable(torch.cuda.FloatTensor([0]))
loss_c = Variable(torch.cuda.FloatTensor([0]))
if (len(sup_image_index) != 0):
try:
loss_l, loss_c = criterion(output, targets)
except:
break
print('--------------')
consistency_loss = csd_criterion(args, conf, conf_flip, loc,
loc_flip,
conf_consistency_criterion)
interpolation_consistency_conf_loss, fixmatch_loss = isd_criterion(
args, lam, conf, conf_flip, loc, loc_flip, conf_shuffle,
conf_interpolation, loc_shuffle, loc_interpolation,
conf_consistency_criterion)
consistency_loss = consistency_loss.mean()
interpolation_loss = torch.mul(
interpolation_consistency_conf_loss.mean(),
0.1) + fixmatch_loss.mean()
ramp_weight = rampweight(iteration)
consistency_loss = torch.mul(consistency_loss, ramp_weight)
interpolation_loss = torch.mul(interpolation_loss, ramp_weight)
if (supervised_flag == 1):
loss = loss_l + loss_c + consistency_loss + interpolation_loss
else:
if (len(sup_image_index) == 0):
loss = consistency_loss + interpolation_loss
else:
loss = loss_l + loss_c + consistency_loss + interpolation_loss
if (loss.data > 0):
optimizer.zero_grad()
loss.backward()
optimizer.step()
t1 = time.time()
if (len(sup_image_index) == 0):
loss_l.data = Variable(torch.cuda.FloatTensor([0]))
loss_c.data = Variable(torch.cuda.FloatTensor([0]))
else:
loc_loss += loss_l.data # [0]
conf_loss += loss_c.data # [0]
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) +
' || Loss: %.4f || consistency_loss : %.4f ||' %
(loss.data, consistency_loss.data),
end=' ')
print(
'loss: %.4f , loss_c: %.4f , loss_l: %.4f , loss_con: %.4f, loss_interpolation: %.4f, lr : %.4f, super_len : %d\n'
% (loss.data, loss_c.data, loss_l.data,
consistency_loss.data, interpolation_loss.data,
float(optimizer.param_groups[0]['lr']),
len(sup_image_index)))
if (float(loss) > 100):
break
if args.visdom:
update_vis_plot(iteration, loss_l.data, loss_c.data, iter_plot,
epoch_plot, 'append')
if iteration != 0 and (iteration + 1) % 120000 == 0:
print('Saving state, iter:', iteration)
torch.save(
ssd_net.state_dict(),
'weights/ssd300_COCO_' + repr(iteration + 1) + '.pth')
# torch.save(ssd_net.state_dict(), args.save_folder + '' + args.dataset + '.pth')
print('-------------------------------\n')
print(loss.data)
print('-------------------------------')
if ((iteration + 1) == cfg['max_iter']):
finish_flag = False
def train(args):
if args.dataset == 'COCO':
if args.dataset_root == VOC_ROOT:
if not os.path.exists(COCO_ROOT):
# FIXME: change parser.error
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:
# FIXME: change parser.error
parser.error('Must specify dataset if specifying dataset_root')
cfg = voc
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))
viz = None
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:
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()
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('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(
viz, 'Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot(
viz, '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(viz, 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, args.lr)
# 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 = 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]
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' %
(loss.data[0]), end=' ')
if args.visdom:
update_vis_plot(viz, iteration, loss_l.data[0], loss_c.data[0],
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
model_name = 'ssd300_{}_{}.pth'.format(
args.dataset, repr(iteration))
model_path = os.path.join(args.save_folder, model_name)
torch.save(ssd_net.state_dict(), model_path)
finished_model_path = os.path.join(
args.save_folder, '{}.pth'.format(args.dataset))
torch.save(ssd_net.state_dict(), finished_model_path)
def train():
cfg = config.Damage
train_dataset = Damage_Dataset(name='train', label_root=args.train_label,
transform=SSDAugmentation())
val_dataset = Damage_Dataset(name='validation', label_root=args.val_label,
transform=BaseTransform())
ssd_net = build_ssd('train', cfg['min_dim'], config.num_classes)
net = ssd_net
best_loss = float('inf')
#cycle_cos_lr = cycle_lr(500, cfg['peak_lr'], cfg['T_init'], cfg['T_warmup'])
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('../../pretrained/vgg16_reducedfc.pth')
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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=config.lr, momentum=config.momentum,
weight_decay=config.weight_decay)
criterion = MultiBoxLoss(config.num_classes, overlap_thresh=0.5,
prior_for_matching=True, bkg_label=0,
neg_mining=True, neg_pos=3, neg_overlap=0.5,
encode_target=False, use_gpu=args.cuda)
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading the train dataset...')
epoch_size = len(train_dataset) // config.batch_size
print('Training SSD on:', train_dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.visdom:
vis_title = os.getcwd().split('/')[-1]
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)
iter_val_plot = create_vis_plot('Iteration', 'Val Loss', vis_title, vis_legend)
data_loader = data.DataLoader(train_dataset, config.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
val_data_loader = data.DataLoader(val_dataset, config.batch_size,
num_workers=args.num_workers,shuffle=True,
collate_fn=detection_collate, pin_memory=True)
# create batch iterator
batch_iterator = iter(data_loader)
val_batch_iterator = iter(val_data_loader)
for iteration in range(args.start_iter, config.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 config.lr_steps:
step_index += 1
adjust_learning_rate(optimizer, config.gamma, step_index)
# cycle lr
#for param_group in optimizer.param_groups:
# param_group['lr'] = cycle_cos_lr.get_lr(iteration)
# 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 = 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.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 args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * iteration,
Y=torch.Tensor([loss_l, loss_c, loss]).unsqueeze(0).cpu(),
win=iter_plot,
update='True' if iteration == 10 else 'append'
)
if iteration % 100 == 0 and iteration != 0:
val_loss_l, val_loss_c, val_loss = val(net, val_data_loader, criterion)
print('Val_Loss: %.4f ||' % (val_loss.item()), end=' ')
if val_loss < best_loss:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(), args.save_folder + 'model/'
+ 'best.pth')
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * iteration,
Y=torch.Tensor([val_loss_l, val_loss_c, val_loss]).unsqueeze(0).cpu(),
win=iter_val_plot,
update='True' if iteration == 100 else 'append'
)
def main():
global args
global minmum_loss
args.gpu = 0
args.world_size = 1
if args.distributed:
args.gpu = args.local_rank % torch.cuda.device_count()
torch.cuda.set_device(args.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.total_batch_size = args.world_size * args.batch_size
## DATA loading code
if args.dataset == 'COCO':
if not os.path.exists(cfg['coco_root']):
parser.error('Must specify dataset_root if specifying dataset')
print("WARNING: Using default COCO dataset_root because " +
"--dataset_root was not specified.")
cfg = coco
dataset = COCODetection(root=cfg['coco_root'],
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
if args.dataset == 'VOC':
cfg = voc
dataset = VOCDetection(root=cfg['voc_root'],
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
print('Training SSD on:', dataset.name)
print('Loading the dataset...')
train_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
print("Build ssd network")
model = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
if args.pretrained:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Loading base network...')
model.vgg.load_state_dict(vgg_weights)
model = model.cuda()
# optimizer and loss function
optimizer = optim.SGD(model.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, True)
## get the priorbox of ssd
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(
args.resume,
map_location=lambda storage, loc: storage.cuda(args.gpu))
args.start_epoch = checkpoint['epoch']
minmum_loss = checkpoint['minmum_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
model.extras.apply(weights_init)
model.loc.apply(weights_init)
model.conf.apply(weights_init)
print('Using the specified args:')
print(args)
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
end = time.time()
loss = train(train_loader, model, priors, criterion, optimizer, epoch)
# remember best prec@1 and save checkpoint
if args.local_rank == 0:
is_best = loss < minmum_loss
minmum_loss = min(loss, minmum_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': minmum_loss,
'optimizer': optimizer.state_dict(),
}, is_best, epoch)
epoch_time = time.time() - end
print('Epoch %s time cost %f' % (epoch, epoch_time))
def train():
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
dataset = COCODetection(image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
transform=SSDAugmentation(MEANS))
if args.validation_epoch > 0:
setup_eval()
val_dataset = COCODetection(image_path=cfg.dataset.valid_images,
info_file=cfg.dataset.valid_info,
transform=BaseTransform(MEANS))
# Parallel wraps the underlying module, but when saving and loading we don't want that
yolact_net = Yolact()
net = yolact_net
net.train()
# I don't use the timer during training (I use a different timing method).
# Apparently there's a race condition with multiple GPUs.
timer.disable_all()
# Both of these can set args.resume to None, so do them before the check
if args.resume == 'interrupt':
args.resume = SavePath.get_interrupt(args.save_folder)
elif args.resume == 'latest':
args.resume = SavePath.get_latest(args.save_folder, cfg.name)
args.resume = './weights/yolact_base_2_8000.pth'
if args.resume is not None:
print('Resuming training, loading {}...'.format(args.resume))
yolact_net.load_weights(args.resume)
if args.start_iter == -1:
args.start_iter = SavePath.from_str(args.resume).iteration
else:
print('Initializing weights...')
yolact_net.init_weights(backbone_path=args.save_folder +
cfg.backbone.path)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.decay)
criterion = MultiBoxLoss(num_classes=cfg.num_classes,
pos_threshold=cfg.positive_iou_threshold,
neg_threshold=cfg.negative_iou_threshold,
negpos_ratio=3)
if args.cuda:
cudnn.benchmark = True
net = nn.DataParallel(net).cuda()
criterion = nn.DataParallel(criterion).cuda()
# loss counters
loc_loss = 0
conf_loss = 0
iteration = max(args.start_iter, 0)
last_time = time.time()
epoch_size = len(dataset) // args.batch_size
num_epochs = math.ceil(cfg.max_iter / epoch_size)
# Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
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)
save_path = lambda epoch, iteration: SavePath(
cfg.name, epoch, iteration).get_path(root=args.save_folder)
time_avg = MovingAverage()
global loss_types # Forms the print order
loss_avgs = {k: MovingAverage(100) for k in loss_types}
print('Begin training!')
print()
# try-except so you can use ctrl+c to save early and stop training
try:
for epoch in range(num_epochs):
# Resume from start_iter
if (epoch + 1) * epoch_size < iteration:
continue
for datum in data_loader:
# Stop if we've reached an epoch if we're resuming from start_iter
if iteration == (epoch + 1) * epoch_size:
break
# Stop at the configured number of iterations even if mid-epoch
if iteration == cfg.max_iter:
break
# Change a config setting if we've reached the specified iteration
changed = False
for change in cfg.delayed_settings:
if iteration >= change[0]:
changed = True
cfg.replace(change[1])
# Reset the loss averages because things might have changed
for avg in loss_avgs:
avg.reset()
# If a config setting was changed, remove it from the list so we don't keep checking
if changed:
cfg.delayed_settings = [
x for x in cfg.delayed_settings if x[0] > iteration
]
# Warm up by linearly interpolating the learning rate from some smaller value
if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until:
set_lr(optimizer, (args.lr - cfg.lr_warmup_init) *
(iteration / cfg.lr_warmup_until) +
cfg.lr_warmup_init)
# Adjust the learning rate at the given iterations, but also if we resume from past that iteration
while step_index < len(
cfg.lr_steps
) and iteration >= cfg.lr_steps[step_index]:
step_index += 1
set_lr(optimizer, args.lr * (args.gamma**step_index))
# Load training data
# Note, for training on multiple gpus this will use the custom replicate and gather I wrote up there
images, targets, masks, num_crowds = prepare_data(datum)
# Forward Pass
out = net(images)
# Compute Loss
optimizer.zero_grad()
wrapper = ScatterWrapper(targets, masks, num_crowds)
losses = criterion(out, wrapper, wrapper.make_mask())
losses = {k: v.mean()
for k, v in losses.items()
} # Mean here because Dataparallel
loss = sum([losses[k] for k in losses])
# Backprop
loss.backward(
) # Do this to free up vram even if loss is not finite
if torch.isfinite(loss).item():
optimizer.step()
# Add the loss to the moving average for bookkeeping
for k in losses:
loss_avgs[k].add(losses[k].item())
cur_time = time.time()
elapsed = cur_time - last_time
last_time = cur_time
# Exclude graph setup from the timing information
if iteration != args.start_iter:
time_avg.add(elapsed)
if iteration % 50 == 0:
eta_str = str(
datetime.timedelta(seconds=(cfg.max_iter - iteration) *
time_avg.get_avg())).split('.')[0]
total = sum([loss_avgs[k].get_avg() for k in losses])
loss_labels = sum([[k, loss_avgs[k].get_avg()]
for k in loss_types if k in losses], [])
print(('[%3d] %7d ||' + (' %s: %.3f |' * len(losses)) +
' T: %.3f || ETA: %s || timer: %.3f') %
tuple([epoch, iteration] + loss_labels +
[total, eta_str, elapsed]),
flush=True)
iteration += 1
if iteration % args.save_interval == 0 and iteration != args.start_iter:
if args.keep_latest:
latest = SavePath.get_latest(args.save_folder,
cfg.name)
print('Saving state, iter:', iteration)
yolact_net.save_weights(save_path(epoch, iteration))
if args.keep_latest and latest is not None:
if args.keep_latest_interval <= 0 or iteration % args.keep_latest_interval != args.save_interval:
print('Deleting old save...')
os.remove(latest)
# This is done per epoch
if args.validation_epoch > 0:
if epoch % args.validation_epoch == 0 and epoch > 0:
compute_validation_map(yolact_net, val_dataset)
except KeyboardInterrupt:
print('Stopping early. Saving network...')
# Delete previous copy of the interrupted network so we don't spam the weights folder
SavePath.remove_interrupt(args.save_folder)
yolact_net.save_weights(
save_path(epoch,
repr(iteration) + '_interrupt'))
exit()
yolact_net.save_weights(save_path(epoch, iteration))
def train():
print(args)
# 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))
# el
cfg = barcode_cfg_dict
if args.dataset == 'barcode':
cfg = barcode_cfg_dict
dataset = BARCODEDetection(img_list_path=args.img_list_file_path,
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(f'Resuming training, loading {args.resume}...')
ssd_net.load_weights(args.resume)
else:
vgg_weights = torch.load(os.path.join(args.save_folder, args.basenet))
print('Loading base network...')
ssd_net.vgg.load_state_dict(vgg_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
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
if args.visdom:
pass
# 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
# begin epoch rewrite
for epoch_num in range(0, cfg['max_epoch'] + 1):
epoch_begin = time.time()
loc_loss = 0
conf_loss = 0
# if args.visdom:
# update_vis_plot(epoch_num, loc_loss, conf_loss, epoch_plot, None, 'append', epoch_size)
# reset epoch loss counters
t0 = time.time()
if epoch_num in (280, 350, 400):
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
print(f"epoch:{epoch_num} at time {time.time()}")
for iter_int, iteration in enumerate(iter(data_loader)):
images, targets = iteration
with torch.no_grad():
if args.cuda:
images = Variable(images.cuda(non_blocking=True))
targets = [
Variable(ann.cuda(non_blocking=True))
for ann in targets
]
else:
images = Variable(images)
targets = [Variable(ann) for ann in targets]
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
loc_loss += loss_l.data.item()
conf_loss += loss_c.data.item()
if iter_int % 10 == 0:
print(f'iter {repr(iter_int)} || Loss: {loss.data.item()} ||',
flush=True)
print(f'timer: {time.time() - t0} sec.', flush=True)
print(f'loc_loss {float(loc_loss)}. conf_loss {conf_loss}.',
flush=True)
t0 = time.time()
# if args.visdom:
# update_vis_plot(iter_int, loss_l.data.item(), loss_c.data.item(), iter_plot, epoch_plot, 'append')
if epoch_num % 5 == 0:
print(f'Saving state, iter: {epoch_num}', flush=True)
torch.save(
ssd_net.state_dict(),
f'{args.save_folder}/ssd300_{args.dataset}/{repr(epoch_num)}.pth'
)
torch.save(ssd_net.state_dict(),
os.path.join(args.save_folder, f'{args.dataset}.pth'))
if __name__ == '__main__':
devices = [0]
device, device_ids = prepare_device(devices)
T_file = 'model/KDGAN_1F1D/T_var/prune_' + repr(
compression_rate) + '/' + repr(Epoch) + 'epoch/lr_0508/20200329/'
if os.path.exists(Path(T_file)) is False:
os.makedirs(T_file)
T_model = build_ssd('train', 300, 21, device)
try:
dataset = VOCDetection(root='/home/hanchengye/data/VOCdevkit/',
image_sets=[('2007', 'train'),
('2012', 'train')],
transform=SSDAugmentation(300, MEANS))
except FileNotFoundError:
dataset = VOCDetection(
root='/remote-home/source/remote_desktop/share/Dataset/VOCdevkit/',
image_sets=[('2007', 'train'), ('2012', 'train')],
transform=SSDAugmentation(300, MEANS))
data_loader_rank = data.DataLoader(dataset,
batch_size=32,
num_workers=4,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
param_T = torch.load('weights/ssd300_COCO_395000.pth')
T_model.load_state_dict(param_T)
