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