def train():
tb_writer, cfg_path, snapshot_dir, log_dir = setup_folder(args, cfg)
step_index = 0
train_loader = dataset_factory(phase='train', cfg=cfg)
val_loader = dataset_factory(phase='eval', cfg=cfg)
eval_solver = eval_solver_factory(val_loader, cfg)
ssd_net, priors, _ = model_factory(phase='train', cfg=cfg, tb_writer=tb_writer)
net = ssd_net # net is the parallel version of ssd_net
print(net)
print(cfg.TRAIN.OPTIMIZER)
# return
if args.cuda:
net = torch.nn.DataParallel(ssd_net)
priors = Variable(priors.cuda(), volatile=True)
else:
priors = Variable(priors)
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
checkpoint = torch.load(args.resume)
args.start_iter = checkpoint['iteration']
step_index = checkpoint['step_index']
ssd_net.load_state_dict(checkpoint['state_dict'])
else:
# pretained weights
pretrained_weights = torch.load(osp.join(cfg.GENERAL.WEIGHTS_ROOT, args.basenet))
print('Loading base network...')
try:
ssd_net.base.load_state_dict(pretrained_weights)
except:
model_dict = ssd_net.base.state_dict()
pretrained_weights = {k: v for k,
v in pretrained_weights.items() if k in model_dict}
model_dict.update(pretrained_weights)
ssd_net.base.load_state_dict(model_dict)
# initialize newly added layers' weights with xavier method
print('Initializing weights...')
ssd_net.extras.apply(weights_init)
ssd_net.loc.apply(weights_init)
ssd_net.conf.apply(weights_init)
if args.cuda:
net = net.cuda()
optimizer = optim.SGD(net.parameters(), lr=cfg.TRAIN.OPTIMIZER.LR,
momentum=cfg.TRAIN.OPTIMIZER.MOMENTUM,
weight_decay=cfg.TRAIN.OPTIMIZER.WEIGHT_DECAY)
criterion = MultiBoxLoss(cfg.MODEL.NUM_CLASSES, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
# continue training at 8w, 12w...
if args.start_iter not in cfg.TRAIN.LR_SCHEDULER.STEPS and step_index != 0:
adjust_learning_rate(optimizer, cfg.TRAIN.LR_SCHEDULER.GAMMA, step_index)
net.train()
epoch_size = len(train_loader.dataset) // cfg.DATASET.TRAIN_BATCH_SIZE
num_epochs = (cfg.TRAIN.MAX_ITER + epoch_size - 1) // epoch_size
print('Training SSD on:', train_loader.dataset.name)
print('Using the specified args:')
print(args)
# timer
t_ = {'network': Timer(), 'misc': Timer(), 'all': Timer(), 'eval': Timer()}
t_['all'].tic()
iteration = args.start_iter
for epoch in range(num_epochs):
tb_writer.cfg['epoch'] = epoch
for images, targets, _ in train_loader:
tb_writer.cfg['iteration'] = iteration
# t_['misc'].tic()
# if iteration in cfg.TRAIN.LR_SCHEDULER.STEPS:
# t_['misc'].tic()
# step_index += 1
# adjust_learning_rate(optimizer, cfg.TRAIN.LR_SCHEDULER.GAMMA, step_index)
#
# 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
# t_['network'].tic()
# out = net(images)
# out1 = [out[0], out[1], priors]
#
# # backward
# optimizer.zero_grad()
# loss_l, loss_c = criterion(out1, targets)
# loss = loss_l + loss_c
# loss.backward()
# optimizer.step()
# t_['network'].toc()
#
# # log
# if iteration % cfg.TRAIN.LOG_LOSS_ITER == 0:
# t_['misc'].toc()
# print('Iter ' + str(iteration) + ' || Loss: %.3f' % (loss.data[0]) +
# '|| conf_loss: %.3f' % (loss_c.data[0]) + ' || loc loss: %.3f ' % (loss_l.data[0]), end=' ')
# print('Timer: %.3f sec.' % t_['misc'].diff, ' Lr: %.6f' % optimizer.param_groups[0]['lr'])
# if args.tensorboard:
# phase = tb_writer.cfg['phase']
# tb_writer.writer.add_scalar('{}/loc_loss'.format(phase), loss_l.data[0], iteration)
# tb_writer.writer.add_scalar('{}/conf_loss'.format(phase), loss_c.data[0], iteration)
# tb_writer.writer.add_scalar('{}/all_loss'.format(phase), loss.data[0], iteration)
# tb_writer.writer.add_scalar('{}/time'.format(phase), t_['misc'].diff, iteration)
#
# # save model
# if iteration % cfg.TRAIN.SAVE_ITER == 0 and iteration != args.start_iter or \
# iteration == cfg.TRAIN.MAX_ITER:
# print('Saving state, iter:', iteration)
# save_checkpoint({'iteration': iteration,
# 'step_index': step_index,
# 'state_dict': ssd_net.state_dict()},
# snapshot_dir,
# args.cfg_name + '_' + repr(iteration) + '.pth')
# Eval
if (iteration % cfg.TRAIN.EVAL_ITER == 0 ) or \
iteration == cfg.TRAIN.MAX_ITER:
print('Start evaluation ......')
tb_writer.cfg['phase'] = 'eval'
t_['eval'].tic()
net.eval()
aps, mAPs = eval_solver.validate(net, priors, tb_writer=tb_writer)
net.train()
t_['eval'].toc()
print('Iteration ' + str(iteration) + ' || mAP: %.3f' % mAPs[0] + ' ||eval_time: %.4f/%.4f' %
(t_['eval'].diff, t_['eval'].average_time))
if cfg.DATASET.NAME == 'VOC0712':
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[0], iteration)
else:
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[0], iteration)
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[1], iteration)
tb_writer.cfg['phase'] = 'train'
return
if iteration == cfg.TRAIN.MAX_ITER:
break
iteration += 1
backup_jobs(cfg, cfg_path, log_dir)
def train():
tb_writer, cfg_path, snapshot_dir, log_dir = setup_folder(args, cfg)
step_index = 0
train_loader = dataset_factory(phase='train', cfg=cfg)
val_loader = dataset_factory(phase='eval', cfg=cfg)
eval_solver = eval_solver_factory(val_loader, cfg)
ssd_net, priors, _ = model_factory(phase='train', cfg=cfg)
net = ssd_net # net is the parallel version of ssd_net
print(net)
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
checkpoint = torch.load(args.resume)
args.start_iter = checkpoint['iteration'] + 1
step_index = checkpoint['step_index']
ssd_net.load_state_dict(checkpoint['state_dict'])
elif cfg.MODEL.PRETRAIN_MODEL != '':
# pretained weights
pretrain_weights = torch.load(cfg.MODEL.PRETRAIN_MODEL)
if 'reducedfc' not in cfg.MODEL.PRETRAIN_MODEL:
ssd_net.apply(weights_init)
try:
ssd_net.load_state_dict(pretrain_weights['state_dict'],
strict=False)
except RuntimeError: # another dataset
entries = [
i for i in pretrain_weights['state_dict'].keys()
if i.startswith('conf')
]
for key in entries:
del pretrain_weights['state_dict'][key]
ssd_net.load_state_dict(pretrain_weights['state_dict'],
strict=False)
else:
print('Loading base network...')
ssd_net.base.load_state_dict(pretrain_weights)
# initialize newly added layers' weights with xavier method
print('Initializing weights...')
ssd_net.extras.apply(weights_init)
ssd_net.loc.apply(weights_init)
ssd_net.conf.apply(weights_init)
else:
print('Initializing weights...')
ssd_net.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=cfg.TRAIN.OPTIMIZER.LR,
momentum=cfg.TRAIN.OPTIMIZER.MOMENTUM,
weight_decay=cfg.TRAIN.OPTIMIZER.WEIGHT_DECAY)
if args.cuda:
net = torch.nn.DataParallel(ssd_net, device_ids=cfg.GENERAL.NET_CPUS)
priors = Variable(priors.cuda(cfg.GENERAL.LOSS_GPU),
requires_grad=False)
net = net.cuda()
else:
priors = Variable(priors, requires_grad=False)
ssd_net.priors = priors
ssd_net.criterion = DetectLoss(cfg)
criterion_post = DetectLossPost(cfg)
net.train()
print('Using the specified args: \n', args)
epoch_size = len(train_loader.dataset) // cfg.DATASET.TRAIN_BATCH_SIZE
num_epochs = (cfg.TRAIN.MAX_ITER + epoch_size - 1) // epoch_size
iteration = args.start_iter
start_epoch = int(iteration * 1.0 / epoch_size)
# continue training at 8w, 12w...
if step_index > 0:
adjust_learning_rate(optimizer, cfg.TRAIN.OPTIMIZER.LR,
cfg.TRAIN.LR_SCHEDULER.GAMMA, 100, step_index,
None, None)
# timer
t_ = {'network': Timer(), 'misc': Timer(), 'eval': Timer()}
t_['misc'].tic()
iteration = args.start_iter
for epoch in range(start_epoch, num_epochs):
for images, targets, _ in train_loader:
if iteration in cfg.TRAIN.LR_SCHEDULER.STEPS or (
iteration <= cfg.TRAIN.WARMUP_EPOCH * epoch_size):
if iteration in cfg.TRAIN.LR_SCHEDULER.STEPS: step_index += 1
adjust_learning_rate(optimizer, cfg.TRAIN.OPTIMIZER.LR,
cfg.TRAIN.LR_SCHEDULER.GAMMA, epoch,
step_index, iteration, epoch_size,
cfg.TRAIN.WARMUP_EPOCH)
# save model
if iteration % cfg.TRAIN.SAVE_ITER == 0 and iteration != args.start_iter or \
iteration == cfg.TRAIN.MAX_ITER:
print('Saving state, iter:', iteration)
save_checkpoint(
{
'iteration': iteration,
'step_index': step_index,
'state_dict': ssd_net.state_dict()
}, snapshot_dir,
args.cfg_name + '_' + repr(iteration) + '.pth')
# Eval
if iteration % cfg.TRAIN.EVAL_ITER == 0 or iteration == cfg.TRAIN.MAX_ITER:
t_['eval'].tic()
net.eval()
aps, mAPs = eval_solver.validate(net,
priors,
tb_writer=tb_writer)
net.train()
t_['eval'].toc()
print('Iteration ' + str(iteration) +
' || mAP: %.3f' % mAPs[0] + ' ||eval_time: %.4f/%.4f' %
(t_['eval'].diff, t_['eval'].average_time))
if tb_writer is not None:
if cfg.DATASET.NAME == 'VOC0712' or 'FACE':
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[0],
iteration)
else:
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[0],
iteration)
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[1],
iteration)
if iteration == cfg.TRAIN.MAX_ITER:
break
if args.cuda:
images = Variable(images.cuda(), requires_grad=False)
targets = [
Variable(ann.cuda(cfg.GENERAL.LOSS_GPU), volatile=True)
for ann in targets
]
else:
images = Variable(images)
targets = [Variable(ann, volatile=True) for ann in targets]
# forward
t_['network'].tic()
match_result = matching(targets,
priors,
cfg.LOSS.OVERLAP_THRESHOLD,
cfg.MODEL.VARIANCE,
args.cuda,
cfg.GENERAL.LOSS_GPU,
cfg=cfg)
net_outputs = net(images,
match_result=match_result,
tb_writer=tb_writer)
loss, (loss_l, loss_c) = criterion_post(net_outputs)
loss_str = ' || Loss: %.3f' % (
loss.data[0]) + '|| conf_loss: %.3f' % (
loss_c) + ' || loc_loss: %.3f ' % (loss_l)
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
t_['network'].toc()
t_['misc'].toc()
# log
if iteration % cfg.TRAIN.LOG_LOSS_ITER == 0:
current_date = time.strftime("%Y-%m-%d %H:%M:%S",
time.localtime())
print('Iter ' + str(iteration) + loss_str, end=' ')
print(
'Timer: %.3f(%.3f) %.3f(%.3f) sec.' %
(t_['misc'].diff, t_['misc'].average_time,
t_['network'].diff, t_['network'].average_time),
'lr: %.6f' % optimizer.param_groups[0]['lr'], ' sys_time:',
current_date)
if tb_writer is not None:
phase = 'train'
tb_writer.writer.add_scalar('{}/loc_loss'.format(phase),
loss_l, iteration)
tb_writer.writer.add_scalar('{}/conf_loss'.format(phase),
loss_c, iteration)
tb_writer.writer.add_scalar('{}/all_loss'.format(phase),
loss.data[0], iteration)
tb_writer.writer.add_scalar('{}/time'.format(phase),
t_['misc'].diff, iteration)
iteration += 1
t_['misc'].tic()
import caffe
resize_h, resize_w = cfg.DATASET.IMAGE_SIZE
InputShape = [1, 3, resize_h, resize_w]
""" Convert! """
print('Converting...')
text_net, binary_weights = ConvertModel_caffe(net,
InputShape,
softmax=False)
model_path = os.path.join(log_dir, 'caffemodel')
model_name = args.trained_model.strip('.pth')
prototxtfile = os.path.join(model_path, model_name + '.prototxt')
caffemodelfile = os.path.join(model_path, model_name + '.caffemodel')
import google.protobuf.text_format
with open(prototxtfile, 'w') as f:
f.write(google.protobuf.text_format.MessageToString(text_net))
with open(caffemodelfile, 'w') as f:
f.write(binary_weights.SerializeToString())
exit()
print('test_type:', cfg.DATASET.TEST_SETS, 'test_model:',
args.trained_model, 'device_id:', cfg.GENERAL.CUDA_VISIBLE_DEVICES)
eval_solver = eval_solver_factory(loader, cfg)
res, mAPs = eval_solver.validate(net, priors, tb_writer=tb_writer)
print('final mAP', mAPs)
eval_solver.eval_pr_curve()
def train():
tb_writer, cfg_path, snapshot_dir, log_dir = setup_folder(args, cfg)
print(cfg_path)
step_index = 0
train_loader = dataset_factory(phase='train', cfg=cfg)
val_loader = dataset_factory(phase='eval', cfg=cfg)
eval_solver = eval_solver_factory(val_loader, cfg)
cls_net = clsmodel_factory(phase='train', cfg=cfg)
net = cls_net # net is the parallel version of cls_net
print(net)
if args.shownet:
return;
if args.cuda:
net = torch.nn.DataParallel(cls_net, device_ids=cfg.GENERAL.NET_CPUS)
cls_net.apply(weights_init)
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
checkpoint = torch.load(args.resume)
args.start_iter = checkpoint['iteration']
step_index = checkpoint['step_index']
cls_net.load_state_dict(checkpoint['state_dict'].state_dict())
elif args.pretrain:
# pretained weights
print('Loading pretrained model: {}'.format(cfg.MODEL.PRETRAIN_MODEL))
pretrain_weights = torch.load(cfg.MODEL.PRETRAIN_MODEL)
if 'reducedfc' not in cfg.MODEL.PRETRAIN_MODEL:
print('Loading whole network...')
cls_net.load_state_dict(pretrain_weights['state_dict'].state_dict(), strict=False)
# cls_net.apply(weights_init)
# try:
# cls_net.load_state_dict(pretrain_weights, strict=False)
# except RuntimeError: # another dataset
# entries = [i for i in pretrain_weights['state_dict'].keys() if i.startswith('conf')]
# for key in entries:
# del pretrain_weights['state_dict'][key]
# cls_net.load_state_dict(pretrain_weights['state_dict'], strict=False)
else:
print('Loading base network...')
cls_net.base.load_state_dict(pretrain_weights['state_dict'].state_dict(), strict=False)
else:
cls_net.apply(weights_init)
print('random init net weight with xavier')
if args.cuda:
net = net.cuda()
optimizer = optim.SGD(net.parameters(), lr=cfg.TRAIN.OPTIMIZER.LR,
momentum=cfg.TRAIN.OPTIMIZER.MOMENTUM,
weight_decay=cfg.TRAIN.OPTIMIZER.WEIGHT_DECAY)
# criterion = MultiBoxLoss(cfg, args.cuda)
# cls_net.criterion = DetectLoss(cfg)
# criterion_post = DetectLossPost(cfg)
# criterion_post = nn.BCEWithLogitsLoss()
criterion_post = FocalLoss_BCE(alpha=0.8,gamma=2,num_classes=1)
# criterion_post = FocalLoss_BCE(alpha=0.5,gamma=0,num_classes=1)
# continue training at 8w, 12w...
if args.start_iter not in cfg.TRAIN.LR_SCHEDULER.STEPS and step_index != 0:
adjust_learning_rate(optimizer, cfg.TRAIN.OPTIMIZER.LR, cfg.TRAIN.LR_SCHEDULER.GAMMA, 100,
step_index, None, None)
net.train()
epoch_size = len(train_loader.dataset) // cfg.DATASET.TRAIN_BATCH_SIZE
num_epochs = (cfg.TRAIN.MAX_ITER - args.start_iter + epoch_size - 1) // epoch_size
print('Training SSD on:', train_loader.dataset.name)
print('Using the specified args:')
print(args)
# timer
t_ = {'network': Timer(), 'forward': Timer(), 'misc': Timer(), 'all': Timer(), 'eval': Timer()}
t_['all'].tic()
iteration = args.start_iter
epoch_bias = int(iteration/epoch_size)
for epoch in range(num_epochs):
epoch += epoch_bias
tb_writer.cfg['epoch'] = epoch
for images,targets,_ in train_loader:
tb_writer.cfg['iteration'] = iteration
t_['misc'].tic()
if iteration in cfg.TRAIN.LR_SCHEDULER.STEPS or \
(epoch < cfg.TRAIN.WARMUP_EPOCH and not args.resume):
if epoch >= cfg.TRAIN.WARMUP_EPOCH:
step_index += 1
adjust_learning_rate(optimizer, cfg.TRAIN.OPTIMIZER.LR, cfg.TRAIN.LR_SCHEDULER.GAMMA, epoch,
step_index, iteration, epoch_size, cfg.TRAIN.WARMUP_EPOCH)
# save model
if iteration % cfg.TRAIN.SAVE_ITER == 0 and iteration != args.start_iter or \
iteration == cfg.TRAIN.MAX_ITER:
print('Saving state, iter:', iteration)
save_checkpoint({'iteration': iteration,
'step_index': step_index,
'state_dict': cls_net},
snapshot_dir,
args.cfg_name + '_' + repr(iteration) + '.pth')
# Eval
if iteration % cfg.TRAIN.EVAL_ITER == 0 or iteration == cfg.TRAIN.MAX_ITER:
tb_writer.cfg['phase'] = 'eval'
tb_writer.cfg['iter'] = iteration
t_['eval'].tic()
net.eval()
if torch.cuda.is_available():
net = nn.DataParallel(net, device_ids=[0])
aps, mAPs = eval_solver.validate(net, tb_writer=tb_writer)
net.train()
if torch.cuda.is_available():
net = torch.nn.DataParallel(cls_net, device_ids=cfg.GENERAL.NET_CPUS)
t_['eval'].toc()
print('Iteration ' + str(iteration) + ' || mAP: %.3f' % mAPs[0] + ' ||eval_time: %.4f/%.4f' %
(t_['eval'].diff, t_['eval'].average_time))
if cfg.DATASET.NAME == 'VOC0712':
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[0], iteration)
else:
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[0], iteration)
tb_writer.writer.add_scalar('mAP/[email protected]', mAPs[1], iteration)
tb_writer.cfg['phase'] = 'train'
if iteration == cfg.TRAIN.MAX_ITER:
break
targets = torch.stack(targets)
if args.cuda:
images = Variable(images.cuda(), requires_grad=False)
# targets = [Variable(ann.cuda(cfg.GENERAL.LOSS_GPU), volatile=True)
# for ann in targets]
targets = Variable(targets.cuda(cfg.GENERAL.LOSS_GPU), requires_grad=False)
else:
images = Variable(images)
# targets = [Variable(ann, volatile=True) for ann in targets]
targets = Variable(targets, requires_grad=False)
# forward
t_['network'].tic()
t_['forward'].tic()
net_outputs = net(images)
t_['forward'].toc()
# import pdb
# pdb.set_trace()
# net_outputs = net_outputs.view(1,-1)
# targets = targets.view(1,-1)
loss = criterion_post(net_outputs,targets)
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
t_['network'].toc()
# log
if iteration % cfg.TRAIN.LOG_LOSS_ITER == 0 or iteration == 1:
t_['misc'].toc()
now_time = datetime.datetime.now()
time_str = datetime.datetime.strftime(now_time,'%Y-%m-%d %H:%M:%S')
print(time_str+'\tIter ' + str(iteration) + ' || Loss: %.3f' % (loss.data[0]), end=' ')
print('Forward Timer: %.3f sec.' % t_['forward'].diff, ' Lr: %.6f' % optimizer.param_groups[0]['lr'])
if args.tensorboard:
phase = tb_writer.cfg['phase']
tb_writer.writer.add_scalar('{}/all_loss'.format(phase), loss.data[0], iteration)
tb_writer.writer.add_scalar('{}/time'.format(phase), t_['misc'].diff, iteration)
iteration += 1
backup_jobs(cfg, cfg_path, log_dir)