imcui/third_party/SGMNet/train/train.py

import torch
import torch.optim as optim
from tqdm import trange
import os
from tensorboardX import SummaryWriter
import numpy as np
import cv2
from loss import SGMLoss,SGLoss
from valid import valid,dump_train_vis

import sys
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, ROOT_DIR)


from utils import train_utils

def train_step(optimizer, model, match_loss, data,step,pre_avg_loss):
    data['step']=step
    result=model(data,test_mode=False)
    loss_res=match_loss.run(data,result)
    
    optimizer.zero_grad()
    loss_res['total_loss'].backward()
    #apply reduce on all record tensor
    for key in loss_res.keys():
        loss_res[key]=train_utils.reduce_tensor(loss_res[key],'mean')
  
    if loss_res['total_loss']<7*pre_avg_loss or step<200 or pre_avg_loss==0:
        optimizer.step()
        unusual_loss=False
    else:
        optimizer.zero_grad()
        unusual_loss=True
    return loss_res,unusual_loss


def train(model, train_loader, valid_loader, config,model_config):
    model.train()
    optimizer = optim.Adam(model.parameters(), lr=config.train_lr)
    
    if config.model_name=='SGM':
        match_loss = SGMLoss(config,model_config) 
    elif config.model_name=='SG':
        match_loss= SGLoss(config,model_config)
    else:
        raise NotImplementedError
    
    checkpoint_path = os.path.join(config.log_base, 'checkpoint.pth')
    config.resume = os.path.isfile(checkpoint_path)
    if config.resume:
        if config.local_rank==0:
            print('==> Resuming from checkpoint..')
        checkpoint = torch.load(checkpoint_path,map_location='cuda:{}'.format(config.local_rank))
        model.load_state_dict(checkpoint['state_dict'])
        best_acc = checkpoint['best_acc']
        start_step = checkpoint['step']
        optimizer.load_state_dict(checkpoint['optimizer'])
    else:
        best_acc = -1
        start_step = 0
    train_loader_iter = iter(train_loader)
    
    if config.local_rank==0:
        writer=SummaryWriter(os.path.join(config.log_base,'log_file'))

    train_loader.sampler.set_epoch(start_step*config.train_batch_size//len(train_loader.dataset))
    pre_avg_loss=0
    
    progress_bar=trange(start_step, config.train_iter,ncols=config.tqdm_width) if config.local_rank==0 else range(start_step, config.train_iter)
    for step in progress_bar:
        try:
            train_data = next(train_loader_iter)
        except StopIteration:
            if config.local_rank==0:
                print('epoch: ',step*config.train_batch_size//len(train_loader.dataset))
            train_loader.sampler.set_epoch(step*config.train_batch_size//len(train_loader.dataset))
            train_loader_iter = iter(train_loader)
            train_data = next(train_loader_iter)
    
        train_data = train_utils.tocuda(train_data)
        lr=min(config.train_lr*config.decay_rate**(step-config.decay_iter),config.train_lr)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # run training
        loss_res,unusual_loss = train_step(optimizer, model, match_loss, train_data,step-start_step,pre_avg_loss)
        if (step-start_step)<=200:
            pre_avg_loss=loss_res['total_loss'].data
        if (step-start_step)>200 and not unusual_loss:
            pre_avg_loss=pre_avg_loss.data*0.9+loss_res['total_loss'].data*0.1
        if unusual_loss and config.local_rank==0:
            print('unusual loss! pre_avg_loss: ',pre_avg_loss,'cur_loss: ',loss_res['total_loss'].data)
        #log
        if config.local_rank==0 and step%config.log_intv==0 and not unusual_loss:
            writer.add_scalar('TotalLoss',loss_res['total_loss'],step)
            writer.add_scalar('CorrLoss',loss_res['loss_corr'],step)
            writer.add_scalar('InCorrLoss', loss_res['loss_incorr'], step)
            writer.add_scalar('dustbin', model.module.dustbin, step)

            if config.model_name=='SGM':
                writer.add_scalar('SeedConfLoss', loss_res['loss_seed_conf'], step)
                writer.add_scalar('MidCorrLoss', loss_res['loss_corr_mid'].sum(), step)
                writer.add_scalar('MidInCorrLoss', loss_res['loss_incorr_mid'].sum(), step)
            

        # valid ans save
        b_save = ((step + 1) % config.save_intv) == 0
        b_validate = ((step + 1) % config.val_intv) == 0
        if b_validate:
            total_loss,acc_corr,acc_incorr,seed_precision_tower,seed_recall_tower,acc_mid=valid(valid_loader, model, match_loss, config,model_config)
            if config.local_rank==0:
                writer.add_scalar('ValidAcc', acc_corr, step)
                writer.add_scalar('ValidLoss', total_loss, step)
                
                if config.model_name=='SGM':
                    for i in range(len(seed_recall_tower)):
                        writer.add_scalar('seed_conf_pre_%d'%i,seed_precision_tower[i],step)
                        writer.add_scalar('seed_conf_recall_%d' % i, seed_precision_tower[i], step)
                    for i in range(len(acc_mid)):
                        writer.add_scalar('acc_mid%d'%i,acc_mid[i],step)
                    print('acc_corr: ',acc_corr.data,'acc_incorr: ',acc_incorr.data,'seed_conf_pre: ',seed_precision_tower.mean().data,
                     'seed_conf_recall: ',seed_recall_tower.mean().data,'acc_mid: ',acc_mid.mean().data)
                else:
                     print('acc_corr: ',acc_corr.data,'acc_incorr: ',acc_incorr.data)
                
                #saving best
                if acc_corr > best_acc:
                    print("Saving best model with va_res = {}".format(acc_corr))
                    best_acc = acc_corr
                    save_dict={'step': step + 1,
                    'state_dict': model.state_dict(),
                    'best_acc': best_acc,
                    'optimizer' : optimizer.state_dict()}
                    save_dict.update(save_dict)
                    torch.save(save_dict, os.path.join(config.log_base, 'model_best.pth'))

        if b_save:
            if config.local_rank==0:
                save_dict={'step': step + 1,
                'state_dict': model.state_dict(),
                'best_acc': best_acc,
                'optimizer' : optimizer.state_dict()}
                torch.save(save_dict, checkpoint_path)
            
            #draw match results
            model.eval()
            with torch.no_grad():
                if config.local_rank==0:
                    if not os.path.exists(os.path.join(config.train_vis_folder,'train_vis')):
                        os.mkdir(os.path.join(config.train_vis_folder,'train_vis'))
                    if not os.path.exists(os.path.join(config.train_vis_folder,'train_vis',config.log_base)):
                        os.mkdir(os.path.join(config.train_vis_folder,'train_vis',config.log_base))
                    os.mkdir(os.path.join(config.train_vis_folder,'train_vis',config.log_base,str(step)))
                res=model(train_data)
                dump_train_vis(res,train_data,step,config)
            model.train()
    
    if config.local_rank==0:
        writer.close()