def main(): import argparse parser = argparse.ArgumentParser( description="imsitu VSRL. Training, evaluation and prediction.") parser.add_argument("--gpuid", default=-1, help="put GPU id > -1 in GPU mode", type=int) parser.add_argument('--output_dir', type=str, default='./trained_models', help='Location to output the model') parser.add_argument('--resume_training', action='store_true', help='Resume training from the model [resume_model]') parser.add_argument('--resume_model', type=str, default='', help='The model we resume') parser.add_argument('--evaluate', action='store_true', help='Only use the testing mode') parser.add_argument('--evaluate_rare', action='store_true', help='Only use the testing mode') parser.add_argument('--test', action='store_true', help='Only use the testing mode') parser.add_argument('--dataset_folder', type=str, default='./imSitu', help='Location of annotations') parser.add_argument('--imgset_dir', type=str, default='./resized_256', help='Location of original images') parser.add_argument('--train_file', default="train_freq2000.json", type=str, help='trainfile name') parser.add_argument('--dev_file', default="dev_freq2000.json", type=str, help='dev file name') parser.add_argument('--test_file', default="test_freq2000.json", type=str, help='test file name') parser.add_argument('--model_saving_name', type=str, help='saving name of the outpul model') parser.add_argument('--epochs', type=int, default=500) parser.add_argument('--model', type=str, default='vgg_caq_joint') parser.add_argument('--batch_size', type=int, default=64) parser.add_argument('--seed', type=int, default=1111, help='random seed') parser.add_argument('--clip_norm', type=float, default=0.25) parser.add_argument('--num_workers', type=int, default=3) parser.add_argument('--vgg_verb_model', type=str, default='', help='Pretrained vgg verb model') parser.add_argument('--tda_verb_model', type=str, default='', help='Pretrained topdown verb model') parser.add_argument('--caq_model', type=str, default='', help='Pretrained CAQ model') args = parser.parse_args() n_epoch = args.epochs batch_size = args.batch_size clip_norm = args.clip_norm n_worker = args.num_workers dataset_folder = args.dataset_folder imgset_folder = args.imgset_dir train_set = json.load(open(dataset_folder + '/' + args.train_file)) encoder = imsitu_encoder.imsitu_encoder(train_set) train_set = imsitu_loader.imsitu_loader(imgset_folder, train_set, encoder, 'train', encoder.train_transform) constructor = 'build_vgg_verb_classifier' vgg_verb_model = getattr(vgg_verb_classifier, constructor)(len(encoder.verb_list)) constructor = 'build_top_down_baseline' role_module = getattr(top_down_baseline_addemb, constructor)(encoder.get_num_roles(), encoder.get_num_verbs(), encoder.get_num_labels(), encoder) constructor = 'build_top_down_verb' tda_verb_model = getattr(top_down_verb, constructor)(encoder.get_num_labels(), encoder.get_num_verbs(), role_module) constructor = 'build_top_down_baseline' tda_role_module = getattr(top_down_baseline, constructor)(encoder.get_num_roles(), encoder.get_num_verbs(), encoder.get_num_labels(), encoder) constructor = 'build_top_down_query_context' caq_model = getattr(top_down_query_context, constructor)(encoder.get_num_roles(), encoder.get_num_verbs(), encoder.get_num_labels(), encoder, tda_role_module) constructor = 'build_%s' % args.model model = getattr(revgg_caq_joint_eval, constructor)(vgg_verb_model, tda_verb_model, caq_model) train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=n_worker) dev_set = json.load(open(dataset_folder + '/' + args.dev_file)) dev_set = imsitu_loader.imsitu_loader(imgset_folder, dev_set, encoder, 'val', encoder.dev_transform) dev_loader = torch.utils.data.DataLoader(dev_set, batch_size=batch_size, shuffle=True, num_workers=n_worker) test_set = json.load(open(dataset_folder + '/' + args.test_file)) test_set = imsitu_loader.imsitu_loader(imgset_folder, test_set, encoder, 'test', encoder.dev_transform) test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=True, num_workers=n_worker) if not os.path.exists(args.output_dir): os.mkdir(args.output_dir) torch.manual_seed(args.seed) if args.gpuid >= 0: model.cuda() torch.cuda.manual_seed(args.seed) torch.backends.cudnn.benchmark = True #load models utils.load_net(args.vgg_verb_model, [model.vgg_model]) print('successfully loaded vgg_verb_model!') utils.load_net(args.tda_verb_model, [model.tda_model]) print('successfully loaded tda_verb_model!') utils.load_net(args.caq_model, [model.caq_model]) print('successfully loaded caq_model!') if args.evaluate: top1, top5, val_loss = eval(model, dev_loader, encoder, args.gpuid, write_to_file=True) top1_avg = top1.get_average_results() top5_avg = top5.get_average_results() avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \ top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"] avg_score /= 8 print('Dev average :{:.2f} {} {}'.format( avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'), utils.format_dict(top5_avg, '{:.2f}', '5-'))) elif args.test: top1, top5, val_loss = eval(model, test_loader, encoder, args.gpuid, write_to_file=True) top1_avg = top1.get_average_results() top5_avg = top5.get_average_results() avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \ top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"] avg_score /= 8 print('Test average :{:.2f} {} {}'.format( avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'), utils.format_dict(top5_avg, '{:.2f}', '5-')))
def train(model, train_loader, dev_loader, optimizer, scheduler, max_epoch, model_dir, encoder, gpu_mode, clip_norm, model_name, model_saving_name, eval_frequency=4000): model.train() train_loss = 0 total_steps = 0 print_freq = 400 dev_score_list = [] if gpu_mode >= 0: ngpus = 2 device_array = [i for i in range(0, ngpus)] pmodel = torch.nn.DataParallel(model, device_ids=device_array) else: pmodel = model top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3) top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3) for epoch in range(max_epoch): for i, (_, img, verb, labels) in enumerate(train_loader): total_steps += 1 if gpu_mode >= 0: img = torch.autograd.Variable(img.cuda()) verb = torch.autograd.Variable(verb.cuda()) labels = torch.autograd.Variable(labels.cuda()) else: img = torch.autograd.Variable(img) verb = torch.autograd.Variable(verb) labels = torch.autograd.Variable(labels) role_predict = pmodel(img, verb) loss = model.calculate_loss(verb, role_predict, labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm) optimizer.step() optimizer.zero_grad() train_loss += loss.item() top1.add_point_noun(verb, role_predict, labels) top5.add_point_noun(verb, role_predict, labels) if total_steps % print_freq == 0: top1_a = top1.get_average_results_nouns() top5_a = top5.get_average_results_nouns() print("{},{},{}, {} , {}, loss = {:.2f}, avg loss = {:.2f}". format(total_steps - 1, epoch, i, utils.format_dict(top1_a, "{:.2f}", "1-"), utils.format_dict(top5_a, "{:.2f}", "5-"), loss.item(), train_loss / ((total_steps - 1) % eval_frequency))) if total_steps % eval_frequency == 0: top1, top5, val_loss = eval(model, dev_loader, encoder, gpu_mode) model.train() top1_avg = top1.get_average_results_nouns() top5_avg = top5.get_average_results_nouns() avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \ top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"] avg_score /= 8 print('Dev {} average :{:.2f} {} {}'.format( total_steps - 1, avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'), utils.format_dict(top5_avg, '{:.2f}', '5-'))) dev_score_list.append(avg_score) max_score = max(dev_score_list) if max_score == dev_score_list[-1]: torch.save( model.state_dict(), model_dir + "/{}_{}.model".format(model_name, model_saving_name)) print('New best model saved! {0}'.format(max_score)) print('current train loss', train_loss) train_loss = 0 top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3) top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3) del role_predict, loss, img, verb, labels print('Epoch ', epoch, ' completed!') scheduler.step()
def printdict(d): return format_dict(d)
def main(): import argparse parser = argparse.ArgumentParser( description="imsitu VSRL. Training, evaluation and prediction.") parser.add_argument("--gpuid", default=-1, help="put GPU id > -1 in GPU mode", type=int) parser.add_argument('--output_dir', type=str, default='./trained_models', help='Location to output the model') parser.add_argument('--resume_training', action='store_true', help='Resume training from the model [resume_model]') parser.add_argument('--resume_model', type=str, default='', help='The model we resume') parser.add_argument('--evaluate', action='store_true', help='Only use the testing mode') parser.add_argument('--evaluate_visualize', action='store_true', help='Only use the testing mode to visualize ') parser.add_argument('--evaluate_rare', action='store_true', help='Only use the testing mode') parser.add_argument('--test', action='store_true', help='Only use the testing mode') parser.add_argument('--dataset_folder', type=str, default='./imSitu', help='Location of annotations') parser.add_argument('--imgset_dir', type=str, default='./resized_256', help='Location of original images') parser.add_argument('--train_file', default="train_freq2000.json", type=str, help='trainfile name') parser.add_argument('--dev_file', default="dev_freq2000.json", type=str, help='dev file name') parser.add_argument('--test_file', default="test_freq2000.json", type=str, help='test file name') parser.add_argument('--model_saving_name', type=str, help='saving name of the outpul model') parser.add_argument('--epochs', type=int, default=500) parser.add_argument('--model', type=str, default='top_down_baseline') parser.add_argument('--batch_size', type=int, default=64) parser.add_argument('--seed', type=int, default=1111, help='random seed') parser.add_argument('--clip_norm', type=float, default=0.25) parser.add_argument('--num_workers', type=int, default=3) args = parser.parse_args() n_epoch = args.epochs batch_size = args.batch_size clip_norm = args.clip_norm n_worker = args.num_workers dataset_folder = args.dataset_folder imgset_folder = args.imgset_dir train_set = json.load(open(dataset_folder + '/' + args.train_file)) encoder = imsitu_encoder.imsitu_encoder(train_set) train_set = imsitu_loader.imsitu_loader(imgset_folder, train_set, encoder, 'train', encoder.train_transform) constructor = 'build_%s' % args.model model = getattr(top_down_baseline, constructor)(encoder.get_num_roles(), encoder.get_num_verbs(), encoder.get_num_labels(), encoder) train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=n_worker) dev_set = json.load(open(dataset_folder + '/' + args.dev_file)) dev_set = imsitu_loader.imsitu_loader(imgset_folder, dev_set, encoder, 'val', encoder.dev_transform) dev_loader = torch.utils.data.DataLoader(dev_set, batch_size=batch_size, shuffle=True, num_workers=n_worker) test_set = json.load(open(dataset_folder + '/' + args.test_file)) test_set = imsitu_loader.imsitu_loader(imgset_folder, test_set, encoder, 'test', encoder.dev_transform) test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=True, num_workers=n_worker) if not os.path.exists(args.output_dir): os.mkdir(args.output_dir) torch.manual_seed(args.seed) if args.gpuid >= 0: model.cuda() torch.cuda.manual_seed(args.seed) torch.backends.cudnn.benchmark = True if args.resume_training: print('Resume training from: {}'.format(args.resume_model)) args.train_all = True if len(args.resume_model) == 0: raise Exception('[pretrained module] not specified') utils.load_net(args.resume_model, [model]) optimizer = torch.optim.Adamax(model.parameters(), lr=1e-3) model_name = 'resume_all' else: print('Training from the scratch.') model_name = 'train_full' utils.set_trainable(model, True) optimizer = torch.optim.Adamax( [{ 'params': model.convnet.parameters(), 'lr': 5e-5 }, { 'params': model.role_emb.parameters() }, { 'params': model.verb_emb.parameters() }, { 'params': model.query_composer.parameters() }, { 'params': model.v_att.parameters() }, { 'params': model.q_net.parameters() }, { 'params': model.v_net.parameters() }, { 'params': model.classifier.parameters() }], lr=1e-3) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9) if args.evaluate: top1, top5, val_loss = eval(model, dev_loader, encoder, args.gpuid) top1_avg = top1.get_average_results_nouns() top5_avg = top5.get_average_results_nouns() avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \ top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"] avg_score /= 8 print('Dev average :{:.2f} {} {}'.format( avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'), utils.format_dict(top5_avg, '{:.2f}', '5-'))) elif args.test: top1, top5, val_loss = eval(model, test_loader, encoder, args.gpuid) top1_avg = top1.get_average_results_nouns() top5_avg = top5.get_average_results_nouns() avg_score = top1_avg["verb"] + top1_avg["value"] + top1_avg["value-all"] + top5_avg["verb"] + \ top5_avg["value"] + top5_avg["value-all"] + top5_avg["value*"] + top5_avg["value-all*"] avg_score /= 8 print('Test average :{:.2f} {} {}'.format( avg_score * 100, utils.format_dict(top1_avg, '{:.2f}', '1-'), utils.format_dict(top5_avg, '{:.2f}', '5-'))) else: print('Model training started!') train( model, train_loader, dev_loader, optimizer, scheduler, n_epoch, args.output_dir, encoder, args.gpuid, clip_norm, model_name, args.model_saving_name, )
def eval(model, loader, encoder, logging=False): model.eval() verbloss = 0 nounsloss = 0 gtloss = 0 top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3) top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3) with torch.no_grad(): for __, (_, img, verb, nouns) in enumerate(loader): if torch.cuda.is_available(): img = img.cuda() verb = verb.cuda() nouns = nouns.cuda() with autocast(): # automix precision stuff pred_verb, pred_nouns, pred_gt_nouns = model(img, verb) top1.add_point_both(pred_verb, verb, pred_nouns, nouns, pred_gt_nouns) top5.add_point_both(pred_verb, verb, pred_nouns, nouns, pred_gt_nouns) if torch.cuda.is_available(): vl = model.module.verb_loss(pred_verb, verb) nl = model.module.nouns_loss(pred_nouns, nouns) gtl = model.module.nouns_loss(pred_gt_nouns, nouns) else: vl = model.verb_loss(pred_verb, verb) nl = model.nouns_loss(pred_nouns, nouns) gtl = model.nouns_loss(pred_gt_nouns, nouns) verbloss += vl.item() nounsloss += nl.item() gtloss += gtl.item() verbloss /= len(loader) nounsloss /= len(loader) gtloss /= len(loader) val_losses = { 'verb_loss': verbloss, 'nouns_loss': nounsloss, 'gt_loss': gtloss } #print scores avg_score = 0 if logging is True: top1_a = top1.get_average_results_both() top5_a = top5.get_average_results_both() avg_score = top1_a['verb'] + top1_a['value'] + top1_a['value-all'] + \ top5_a['verb'] + top5_a['value'] + top5_a['value-all'] + \ top1_a['gt-value'] + top1_a['gt-value-all'] avg_score /= 8 avg_score = avg_score * 100 print('val losses = [v: {:.2f}, n: {:.2f}, gt: {:.2f}]'.format( val_losses['verb_loss'], val_losses['nouns_loss'], val_losses['gt_loss'])) gt = {key: top1_a[key] for key in ['gt-value', 'gt-value-all']} one_val = {key: top1_a[key] for key in ['verb', 'value', 'value-all']} print('{}\n{}\n{}, mean = {:.2f}\n'.format( utils.format_dict(one_val, '{:.2f}', '1-'), utils.format_dict(top5_a, '{:.2f}', '5-'), utils.format_dict(gt, '{:.2f}', ''), avg_score)) return top1, top5, val_losses, avg_score
def train(model, train_loader, dev_loader, optimizer, max_epoch, encoder, model_saving_name, folder, checkpoint=None): model.train() avg_scores = [] verb_losses = [] nouns_losses = [] val_avg_scores = [] val_verb_losses = [] val_nouns_losses = [] epoch = 0 # if checkpoint resume stuffs if checkpoint is not None: epoch = checkpoint['epoch'] avg_scores = checkpoint['avg_scores'] verb_losses = checkpoint['verb_losses'] nouns_losses = checkpoint['nouns_losses'] val_avg_scores = checkpoint['val_avg_scores'] val_verb_losses = checkpoint['val_verb_losses'] val_nouns_losses = checkpoint['val_nouns_losses'] if torch.cuda.is_available(): model.module.load_state_dict(checkpoint['model_state_dict']) else: model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) #mix precision stuff scaler = GradScaler() for e in range(epoch, max_epoch): verb_loss_accum = 0 nouns_loss_accum = 0 gt_nouns_loss_accum = 0 print('Epoch-{}, lr: {:.4f}'.format(e, optimizer.param_groups[0]['lr'])) top1 = imsitu_scorer.imsitu_scorer(encoder, 1, 3) top5 = imsitu_scorer.imsitu_scorer(encoder, 5, 3) for i, (_, img, verb, nouns) in enumerate(train_loader): if torch.cuda.is_available(): img = img.cuda() verb = verb.cuda() nouns = nouns.cuda() optimizer.zero_grad() with autocast(): #mix precision stuff pred_verb, pred_nouns, pred_gt_nouns = model(img, verb) #predict and calculate lossess if torch.cuda.is_available(): verb_loss = model.module.verb_loss(pred_verb, verb) nouns_loss = model.module.nouns_loss(pred_nouns, nouns) gt_nouns_loss = model.module.nouns_loss( pred_gt_nouns, nouns) else: verb_loss = model.verb_loss(pred_verb, verb) nouns_loss = model.nouns_loss(pred_nouns, nouns) gt_nouns_loss = model.nouns_loss(pred_gt_nouns, nouns) loss = verb_loss + nouns_loss # backpropagate errors and stuffs scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1) scaler.step(optimizer) scaler.update() top1.add_point_both(pred_verb, verb, pred_nouns, nouns, pred_gt_nouns) top5.add_point_both(pred_verb, verb, pred_nouns, nouns, pred_gt_nouns) verb_loss_accum += verb_loss.item() nouns_loss_accum += nouns_loss.item() gt_nouns_loss_accum += gt_nouns_loss.item() #fin epoch #epoch accuracy mean top1_a = top1.get_average_results_both() top5_a = top5.get_average_results_both() avg_score = top1_a['verb'] + top1_a['value'] + top1_a['value-all'] + \ top5_a['verb'] + top5_a['value'] + top5_a['value-all'] + \ top1_a['gt-value'] + top1_a['gt-value-all'] avg_score /= 8 avg_score *= 100 avg_scores.append(avg_score) #epoch lossess verb_loss_mean = verb_loss_accum / len(train_loader) nouns_loss_mean = nouns_loss_accum / len(train_loader) gt_nouns_loss_mean = gt_nouns_loss_accum / len(train_loader) verb_losses.append(verb_loss_mean) nouns_losses.append(nouns_loss_mean) #print stuffs print('training losses = [v: {:.2f}, n: {:.2f}, gt: {:.2f}]'.format( verb_loss_mean, nouns_loss_mean, gt_nouns_loss_mean)) gt = {key: top1_a[key] for key in ['gt-value', 'gt-value-all']} one_val = {key: top1_a[key] for key in ['verb', 'value', 'value-all']} print('{}\n{}\n{}, mean = {:.2f}\n{}'.format( utils.format_dict(one_val, '{:.2f}', '1-'), utils.format_dict(top5_a, '{:.2f}', '5-'), utils.format_dict(gt, '{:.2f}', ''), avg_score, '-' * 50)) # evaluating top1, top5, val_losses, val_avg_score = eval(model, dev_loader, encoder, logging=True) model.train() #val mean scores val_avg_scores.append(val_avg_score) val_verb_losses.append(val_losses['verb_loss']) val_nouns_losses.append(val_losses['nouns_loss']) plt.plot(verb_losses, label='verb losses') plt.plot(nouns_losses, label='nouns losses') plt.plot(avg_scores, label='accuracy mean') plt.plot(val_verb_losses, '-.', label='val verb losses') plt.plot(val_nouns_losses, '-.', label='val nouns losses') plt.plot(val_avg_scores, '-.', label='val accuracy mean') plt.grid() plt.legend() plt.savefig(pjoin(folder, model_saving_name + '.png')) plt.clf() #always saving but no need if it's not the best score checkpoint = { 'epoch': e + 1, 'avg_scores': avg_scores, 'verb_losses': verb_losses, 'nouns_losses': nouns_losses, 'val_avg_scores': val_avg_scores, 'val_verb_losses': val_verb_losses, 'val_nouns_losses': val_nouns_losses, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict() } if torch.cuda.is_available(): checkpoint.update({'model_state_dict': model.module.state_dict()}) torch.save(checkpoint, pjoin(folder, model_saving_name))
def update_network(self): # initial configuration stop = False update_iters = 0 self.iter('S') self.iter('T') while not stop: # update learning rate self.update_lr(self.optimizer, self.base_lr) # set the status of network self.net.train() self.net.zero_grad() if self.adv_train: self.update_lr(self.optimizer_D, self.base_lr_D) self.net_D.train() self.net_D.zero_grad() loss = 0 for k in range(self.iter_size): sample_S = self.get_training_samples('S') data_S, gt_S = sample_S['Img'], sample_S['Label'] data_S, gt_S = gen_utils.to_cuda(data_S), gen_utils.to_cuda( gt_S) sample_T = self.get_training_samples('T') data_T, gt_T = sample_T['Img'], sample_T['Label'] data_T, gt_T = gen_utils.to_cuda(data_T), gen_utils.to_cuda( gt_T) loss_dict, out_dict = eval('self.%s' % cfg.TRAIN.METHOD)( data_S, gt_S, data_T, gt_T) loss = loss_dict['total'] / self.iter_size preds_S, preds_T = out_dict['preds_S'], out_dict['preds_T'] if self.adv_train: # G step: probs_S, probs_T = F.softmax(preds_S, dim=1), F.softmax(preds_T, dim=1) for param in self.net_D.parameters(): param.requires_grad = False loss_GD = self.G_step(probs_S, probs_T) / self.iter_size loss += cfg.TRAIN.ADV_W * loss_GD loss_dict['G_loss'] = loss_GD loss.backward() if self.adv_train: # D step: for param in self.net_D.parameters(): param.requires_grad = True loss_D = self.D_step(probs_S, probs_T) / self.iter_size loss_dict['D_loss'] = loss_D loss_D.backward() # update the network self.optimizer.step() if self.adv_train: # update the discriminator self.optimizer_D.step() if cfg.TRAIN.LOGGING and (update_iters+1) % \ (max(1, self.iters_per_epoch // cfg.TRAIN.NUM_LOGGING_PER_EPOCH)) == 0: preds = out_dict['preds_S'] accu = 100.0 * gen_utils.model_eval( torch.max(preds, dim=1).indices, gt_S, 'accuracy', preds.size(1), cfg.DATASET.IGNORE_LABEL).item() miou = 100.0 * gen_utils.model_eval( torch.max(preds, dim=1).indices, gt_S, 'mIoU', preds.size(1), cfg.DATASET.IGNORE_LABEL)[0].item() cur_loss = loss_dict eval_res = {'accu': accu, 'miou': miou} self.logging(cur_loss, eval_res) if cfg.TRAIN.TEST_INTERVAL > 0 and \ (self.iters+1) % int(cfg.TRAIN.TEST_INTERVAL * self.iters_per_epoch) == 0: with torch.no_grad(): accu, miou = self.test() print('Test at (epoch %d, iter %d) with %s.' % (self.epochs, self.iters, gen_utils.format_dict({ 'accu': accu, 'miou': miou }))) if not self.distributed or gen_utils.is_main_process(): if cfg.TRAIN.SAVE_CKPT_INTERVAL > 0 and \ (self.iters+1) % int(cfg.TRAIN.SAVE_CKPT_INTERVAL * self.iters_per_epoch) == 0: self.save_ckpt() update_iters += 1 self.iters += 1 # update stop condition if update_iters >= self.iters_per_epoch: stop = True else: stop = False
def logging(self, loss, res): print('[epoch: %d, iter: %d]: ' % (self.epochs, self.iters)) info_str = gen_utils.format_dict(loss) + '; ' + gen_utils.format_dict( res) print(info_str)