Spaces:
Running
Running
File size: 7,168 Bytes
a80d6bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
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()
|