def siamdw_train(train_loader, model, optimizer, epoch, cur_lr, cfg, writer_dict, logger):
# unfix for FREEZE-OUT method
# model, optimizer = unfix_more(model, optimizer, epoch, cfg, cur_lr, logger) # you may try freeze-out
# prepare
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
end = time.time()
# switch to train mode
model.train()
model = model.cuda()
for iter, input in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# input and output/loss
label_cls = input[2].type(torch.FloatTensor) # BCE need float
template = input[0].cuda()
search = input[1].cuda()
label_cls = label_cls.cuda()
loss = model(template, search, label_cls)
loss = torch.mean(loss)
# compute gradient and do update step
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm(model.parameters(), 10) # gradient clip
if is_valid_number(loss.data[0]):
optimizer.step()
# record loss
loss = loss.data[0]
losses.update(loss, template.size(0))
batch_time.update(time.time() - end)
end = time.time()
if (iter + 1) % cfg.PRINT_FREQ == 0:
logger.info('Epoch: [{0}][{1}/{2}] lr: {lr:.7f}\t Batch Time: {batch_time.avg:.3f}s \t Data Time:{data_time.avg:.3f}s \t Loss:{loss.avg:.5f}'.format(
epoch, iter + 1, len(train_loader), lr=cur_lr, batch_time=batch_time, data_time=data_time, loss=losses))
print_speed((epoch - 1) * len(train_loader) + iter + 1, batch_time.avg, cfg.SIAMFC.TRAIN.END_EPOCH * len(train_loader), logger)
# write to tensorboard
writer = writer_dict['writer']
global_steps = writer_dict['train_global_steps']
writer.add_scalar('train_loss', loss, global_steps)
writer_dict['train_global_steps'] = global_steps + 1
return model, writer_dict
def ocean_train(train_loader, model, optimizer, epoch, cur_lr, cfg, writer_dict, logger, device):
# unfix for FREEZE-OUT method
# model, optimizer = unfix_more(model, optimizer, epoch, cfg, cur_lr, logger)
# prepare
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
cls_losses_align = AverageMeter()
cls_losses_ori = AverageMeter()
reg_losses = AverageMeter()
end = time.time()
# switch to train mode
model.train()
model = model.to(device)
for iter, input in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# input and output/loss
label_cls = input[2].type(torch.FloatTensor) # BCE need float
template = input[0].to(device)
search = input[1].to(device)
label_cls = label_cls.to(device)
reg_label = input[3].float().to(device)
reg_weight = input[4].float().to(device)
cls_loss_ori, cls_loss_align, reg_loss = model(template, search, label_cls, reg_target=reg_label, reg_weight=reg_weight)
cls_loss_ori = torch.mean(cls_loss_ori)
reg_loss = torch.mean(reg_loss)
if cls_loss_align is not None:
cls_loss_align = torch.mean(cls_loss_align)
loss = cls_loss_ori + cls_loss_align + 1.2 * reg_loss
else:
cls_loss_align = 0
loss = cls_loss_ori + 1.2 * reg_loss
loss = torch.mean(loss)
# compute gradient and do update step
optimizer.zero_grad()
loss.backward()
# torch.nn.utils.clip_grad_norm(model.parameters(), 10) # gradient clip
if is_valid_number(loss.item()):
optimizer.step()
# record loss
loss = loss.item()
losses.update(loss, template.size(0))
cls_loss_ori = cls_loss_ori.item()
cls_losses_ori.update(cls_loss_ori, template.size(0))
try:
cls_loss_align = cls_loss_align.item()
except:
cls_loss_align = 0
cls_losses_align.update(cls_loss_align, template.size(0))
reg_loss = reg_loss.item()
reg_losses.update(reg_loss, template.size(0))
batch_time.update(time.time() - end)
end = time.time()
if (iter + 1) % cfg.PRINT_FREQ == 0:
logger.info(
'Epoch: [{0}][{1}/{2}] lr: {lr:.7f}\t Batch Time: {batch_time.avg:.3f}s \t Data Time:{data_time.avg:.3f}s \t CLS_ORI Loss:{cls_loss_ori.avg:.5f} \t CLS_ALIGN Loss:{cls_loss_align.avg:.5f} \t REG Loss:{reg_loss.avg:.5f} \t Loss:{loss.avg:.5f}'.format(
epoch, iter + 1, len(train_loader), lr=cur_lr, batch_time=batch_time, data_time=data_time,
loss=losses, cls_loss_ori=cls_losses_ori, cls_loss_align=cls_losses_align, reg_loss=reg_losses))
print_speed((epoch - 1) * len(train_loader) + iter + 1, batch_time.avg,
cfg.OCEAN.TRAIN.END_EPOCH * len(train_loader), logger)
# write to tensorboard
writer = writer_dict['writer']
global_steps = writer_dict['train_global_steps']
writer.add_scalar('train_loss', loss, global_steps)
writer_dict['train_global_steps'] = global_steps + 1
return model, writer_dict
def oceanplus_train(train_loader, model, optimizer, epoch, cur_lr, cfg, writer_dict, logger, device):
# prepare
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
end = time.time()
# switch to train mode
print('====> fix again <=====')
model.train()
try:
model.module.features.features.eval()
except:
model.module.features.eval()
try:
model.module.neck.eval()
model.module.neck.apply(BNtoFixed)
except:
pass
try:
model.module.connect_model.eval()
model.module.connect_model.apply(BNtoFixed)
except:
pass
try:
model.module.bbox_tower.eval()
model.module.bbox_tower.apply(BNtoFixed)
except:
pass
try:
model.module.features.features.apply(BNtoFixed)
except:
model.module.features.apply(BNtoFixed)
model.module.mask_model.train()
model = model.to(device)
for iter, input in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# input and output/loss
label_cls = input[2].type(torch.FloatTensor) # BCE need float
template = input[0].to(device)
search = input[1].to(device)
label_cls = label_cls.to(device)
reg_label = input[3].float().to(device)
reg_weight = input[4].float().to(device)
mask = input[6].float().to(device)
template_mask = input[-1].float().to(device)
mask_weight = input[7].float().to(device)
_, _, loss = model(template, search, label_cls, reg_target=reg_label, reg_weight=reg_weight,
mask=mask, mask_weight=mask_weight, template_mask=template_mask)
loss = torch.mean(loss)
# compute gradient and do update step
optimizer.zero_grad()
loss.backward()
# torch.nn.utils.clip_grad_norm(model.parameters(), 10) # gradient clip
if is_valid_number(loss.item()):
optimizer.step()
torch.nn.utils.clip_grad_norm_(model.parameters(), 10) # gradient clip
# record loss
loss = loss.item()
losses.update(loss, template.size(0))
batch_time.update(time.time() - end)
end = time.time()
if (iter + 1) % cfg.PRINT_FREQ == 0:
logger.info(
'Epoch: [{0}][{1}/{2}] lr: {lr:.7f}\t Batch Time: {batch_time.avg:.3f}s \t Data Time:{data_time.avg:.3f}s \t MASK Loss:{mask_loss.avg:.5f}'.format(
epoch, iter + 1, len(train_loader), lr=cur_lr, batch_time=batch_time, data_time=data_time, mask_loss=losses))
print_speed((epoch - 1) * len(train_loader) + iter + 1, batch_time.avg,
cfg.FREEMASK.TRAIN.END_EPOCH * len(train_loader), logger)
# write to tensorboard
writer = writer_dict['writer']
global_steps = writer_dict['train_global_steps']
writer.add_scalar('train_loss', loss, global_steps)
writer_dict['train_global_steps'] = global_steps + 1
return model, writer_dict
def siamrpn_train(train_loader,
model,
optimizer,
epoch,
cur_lr,
cfg,
writer_dict,
logger,
cls_type='thicker'):
model, optimizer = unfix_more(model, optimizer, epoch, cfg, cur_lr, logger)
batch_time = AverageMeter()
data_time = AverageMeter()
cls_losses = AverageMeter()
reg_losses = AverageMeter()
losses = AverageMeter()
end = time.time()
# switch to train mode
model.train()
model = model.cuda()
for iter, input in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# input and output/loss
template = Variable(input[0]).cuda()
search = Variable(input[1]).cuda()
if cls_type == 'thinner':
label_cls = Variable(input[2].type(torch.FloatTensor)).cuda()
else:
label_cls = Variable(input[2]).cuda()
label_reg = Variable(input[3]).cuda()
label_reg_weight = Variable(input[4]).cuda()
sum_weight = Variable(input[5].type(torch.FloatTensor)).cuda()
cls_loss, loc_loss = model(template, search, label_cls, label_reg,
label_reg_weight, sum_weight)
cls_loss = torch.mean(cls_loss)
loc_loss = torch.mean(loc_loss)
loss = cfg.SIAMRPN.TRAIN.CLS_WEIGHT * cls_loss + cfg.SIAMRPN.TRAIN.REG_WEIGHT * loc_loss
# compute gradient and do update step
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 10) # gradient clip
if is_valid_number(loss.item()):
optimizer.step()
# record loss
cls_loss = cls_loss.item()
loc_loss = loc_loss.item()
loss = loss.item()
losses.update(loss, template.size(0))
cls_losses.update(cls_loss, template.size(0))
reg_losses.update(loc_loss, template.size(0))
batch_time.update(time.time() - end)
end = time.time()
if (iter + 1) % cfg.PRINT_FREQ == 0:
logger.info(
'Epoch: [{0}][{1}/{2}] lr : {lr:.7f} \t Batch Time: {batch_time.avg:.3f} \t Data Time:{data_time.avg:.3f} '
'\tCLS_Loss:{cls_loss.avg:.5f}\tREG_Loss:{reg_loss.avg:.5f} \t Loss:{loss.avg:.5f}'
.format(epoch,
iter + 1,
len(train_loader),
lr=cur_lr,
batch_time=batch_time,
data_time=data_time,
cls_loss=cls_losses,
reg_loss=reg_losses,
loss=losses))
print_speed(
(epoch - 1) * len(train_loader) + iter + 1, batch_time.avg,
cfg.SIAMRPN.TRAIN.END_EPOCH * len(train_loader), logger)
# write to tensorboard
writer = writer_dict['writer']
global_steps = writer_dict['train_global_steps']
writer.add_scalar('train_loss', loss, global_steps)
writer_dict['train_global_steps'] = global_steps + 1
return model, writer_dict
