def train(gpu, config): dist.init_process_group(backend='nccl', init_method='env://', world_size=config['num_gpus'], rank=gpu) torch.cuda.set_device(gpu) """ @ build the dataset for training """ dataset = get_data(config) trainset = dataset(config, "train") testset = dataset(config, "test") sampler_train = DistributedSampler(trainset, num_replicas=config['num_gpus'], rank=gpu) sampler_val = DistributedSampler(testset, num_replicas=config['num_gpus'], rank=gpu) batch_size = config['batch_size'] loader_train = DataLoader(dataset=trainset, batch_size=batch_size, shuffle=False, num_workers=config['num_threads'], pin_memory=True, sampler=sampler_train, drop_last=True) loader_val = DataLoader(dataset=testset, batch_size=1, shuffle=False, num_workers=1, pin_memory=True, sampler=sampler_val, drop_last=True) model = UNet(config["in_channels"], config["out_channels"], post_processing=True) model.cuda(gpu) mask_sampling = masksamplingv2() """ @ init parameter """ save_folder = os.path.join( config['save_root'], 'batch_{}_lr_{}'.format(config['batch_size'], config['lr'])) best_epoch = 0 lowest_loss = 0. resume = 0 print('=>Save folder: {}\n'.format(save_folder)) if not os.path.exists(save_folder): os.makedirs(save_folder) optimizer = define_optim(config['optimizer'], model.parameters(), float(config['lr']), 0) criterion_1 = define_loss(config['loss_type']) criterion_2 = define_loss("Multimse") scheduler = define_scheduler(optimizer, config) """ @ justify the resume model """ if config['resume'] != 'None': checkpoint = torch.load(config['resume'], map_location=torch.device('cpu')) resume = checkpoint['epoch'] lowest_loss = checkpoint['loss'] best_epoch = checkpoint['best_epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) model, optimizer = amp.initialize(model, optimizer, opt_level='O0', verbosity=0) amp.load_state_dict(checkpoint['amp']) print("=> loaded checkpoint '{}' (epoch {})".format( resume, checkpoint['epoch'])) del checkpoint log_file = 'log_train_start_{}.txt'.format(resume) """ @ convert model to multi-gpus modes for training """ model = apex.parallel.convert_syncbn_model(model) if config['resume'] == 'None': model, optimizer = amp.initialize(model, optimizer, opt_level='O0', verbosity=0) model = DDP(model) if gpu == 0: sys.stdout = Logger(os.path.join(save_folder, log_file)) print("Number of parameters in model is {:.3f}M".format( sum(tensor.numel() for tensor in model.parameters()) / 1e6)) """ @ start to train """ for epoch in range(resume + 1, config['epoches'] + 1): print('=> Starch Epoch {}\n'.format(epoch)) print(datetime.now().strftime('%Y-%m-%d %H:%M:%S')) print('learning rate is set to {}.\n'.format( optimizer.param_groups[0]['lr'])) model.train() sampler_train.set_epoch(epoch) batch_time = AverageMeter() losses = AverageMeter() metric_train = Metrics() rmse_train = AverageMeter() mae_train = AverageMeter() time_snap = time.time() for i, inputs in tqdm(enumerate(loader_train)): gt, noise = inputs['gt'].cuda(gpu), inputs['noise'].cuda(gpu) optimizer.zero_grad() """ update the train inputs """ # patten = np.random.randint(0, 4, 1) patten = torch.randint(0, 8, (1, )) redinput, blueinput = mask_sampling(noise, patten) # redinput, blueinput = generator(noise, mask1, mask2) output = model(redinput) loss = criterion_1(output, blueinput) fulloutput = model(noise) redoutput, blueoutput = mask_sampling(fulloutput, patten) # redoutput, blueoutput = generator(fulloutput, mask1, mask2) loss2 = criterion_2(output, blueinput, redoutput, blueoutput) losssum = config["gamma"] * loss2 + loss with amp.scale_loss(losssum, optimizer) as scaled_loss: scaled_loss.backward() optimizer.step() "@ map-reduce tensor" rt = reduce_tensor(losssum.data) torch.cuda.synchronize() losses.update(rt.item(), loader_train.batch_size) metric_train.calculate(fulloutput.detach(), gt) rmse_train.update(metric_train.get_metric('mse'), metric_train.num) mae_train.update(metric_train.get_metric('mae'), metric_train.num) batch_time.update(time.time() - time_snap) time_snap = time.time() if (i + 1) % config['print_freq'] == 0: if gpu == 0: print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 'Loss {loss.val:.6f} ({loss.avg:.6f})\t' 'Metric {rmse_train.val:.6f} ({rmse_train.avg:.6f})'. format(epoch, i + 1, len(loader_train), batch_time=batch_time, loss=losses, rmse_train=rmse_train)) if (i + 1) % config['save_freq'] == 0: print('=> Start sub-selection validation set') rmse, mae = val(model, loader_val, epoch, gpu) model.train() if gpu == 0: print("===> Average RMSE score on selection set is {:.6f}". format(rmse)) print("===> Average MAE score on selection set is {:.6f}". format(mae)) print( "===> Last best score was RMSE of {:.6f} in epoch {}". format(lowest_loss, best_epoch)) if rmse > lowest_loss: lowest_loss = rmse best_epoch = epoch states = { 'epoch': epoch, 'best_epoch': best_epoch, 'loss': lowest_loss, 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'amp': amp.state_dict() } save_checkpoints(states, save_folder, epoch, gpu, True) # save checkpoints print('=> Start selection validation set') rmse, mae = val(model, loader_val, epoch, gpu) model.train() if gpu == 0: print("===> Average RMSE score on selection set is {:.6f}".format( rmse)) print("===> Average MAE score on selection set is {:.6f}".format( mae)) print("===> Last best score was RMSE of {:.6f} in epoch {}".format( lowest_loss, best_epoch)) if rmse > lowest_loss: best_epoch = epoch lowest_loss = rmse states = { 'epoch': epoch, 'best_epoch': best_epoch, 'loss': lowest_loss, 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'amp': amp.state_dict() } save_checkpoints(states, save_folder, epoch, gpu, True) if config['lr_policy'] == 'plateau': scheduler.step(rmse) else: scheduler.step() # if (epoch) % 10 == 0: # config["gamma"] += 0.5 print('=>> the model training finish!')
def main(): global args args = parser.parse_args() if args.num_samples == 0: args.num_samples = None if args.val_batch_size is None: args.val_batch_size = args.batch_size if args.seed: random.seed(args.seed) torch.manual_seed(args.seed) # torch.backends.cudnn.deterministic = True # warnings.warn('You have chosen to seed training. ' # 'This will turn on the CUDNN deterministic setting, ' # 'which can slow down your training considerably! ' # 'You may see unexpected behavior when restarting from checkpoints.') # For distributed training # init_distributed_mode(args) if not args.no_cuda and not torch.cuda.is_available(): raise Exception("No gpu available for usage") torch.backends.cudnn.benchmark = args.cudnn # Init model channels_in = 1 if args.input_type == 'depth' else 4 model = Models.define_model(mod=args.mod, in_channels=channels_in, thres=args.thres) define_init_weights(model, args.weight_init) # Load on gpu before passing params to optimizer if not args.no_cuda: if not args.multi: model = model.cuda() else: model = torch.nn.DataParallel(model).cuda() # model.cuda() # model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) # model = model.module save_id = '{}_{}_{}_{}_{}_batch{}_pretrain{}_wlid{}_wrgb{}_wguide{}_wpred{}_patience{}_num_samples{}_multi{}'.\ format(args.mod, args.optimizer, args.loss_criterion, args.learning_rate, args.input_type, args.batch_size, args.pretrained, args.wlid, args.wrgb, args.wguide, args.wpred, args.lr_decay_iters, args.num_samples, args.multi) # INIT optimizer/scheduler/loss criterion optimizer = define_optim(args.optimizer, model.parameters(), args.learning_rate, args.weight_decay) scheduler = define_scheduler(optimizer, args) # Optional to use different losses criterion_local = define_loss(args.loss_criterion) criterion_lidar = define_loss(args.loss_criterion) criterion_rgb = define_loss(args.loss_criterion) criterion_guide = define_loss(args.loss_criterion) # INIT dataset dataset = Datasets.define_dataset(args.dataset, args.data_path, args.input_type, args.side_selection) dataset.prepare_dataset() train_loader, valid_loader, valid_selection_loader = get_loader( args, dataset) # Resume training best_epoch = 0 lowest_loss = np.inf args.save_path = os.path.join(args.save_path, save_id) mkdir_if_missing(args.save_path) log_file_name = 'log_train_start_0.txt' args.resume = first_run(args.save_path) if args.resume and not args.test_mode and not args.evaluate: path = os.path.join( args.save_path, 'checkpoint_model_epoch_{}.pth.tar'.format(int(args.resume))) if os.path.isfile(path): log_file_name = 'log_train_start_{}.txt'.format(args.resume) # stdout sys.stdout = Logger(os.path.join(args.save_path, log_file_name)) print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(path) args.start_epoch = checkpoint['epoch'] lowest_loss = checkpoint['loss'] best_epoch = checkpoint['best epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: log_file_name = 'log_train_start_0.txt' # stdout sys.stdout = Logger(os.path.join(args.save_path, log_file_name)) print("=> no checkpoint found at '{}'".format(path)) # Only evaluate elif args.evaluate: print("Evaluate only") best_file_lst = glob.glob(os.path.join(args.save_path, 'model_best*')) if len(best_file_lst) != 0: best_file_name = best_file_lst[0] print(best_file_name) if os.path.isfile(best_file_name): sys.stdout = Logger( os.path.join(args.save_path, 'Evaluate.txt')) print("=> loading checkpoint '{}'".format(best_file_name)) checkpoint = torch.load(best_file_name) model.load_state_dict(checkpoint['state_dict']) else: print("=> no checkpoint found at '{}'".format(best_file_name)) else: print("=> no checkpoint found at due to empy list in folder {}". format(args.save_path)) validate(valid_selection_loader, model, criterion_lidar, criterion_rgb, criterion_local, criterion_guide) return # Start training from clean slate else: # Redirect stdout sys.stdout = Logger(os.path.join(args.save_path, log_file_name)) # INIT MODEL print(40 * "=" + "\nArgs:{}\n".format(args) + 40 * "=") print("Init model: '{}'".format(args.mod)) print("Number of parameters in model {} is {:.3f}M".format( args.mod.upper(), sum(tensor.numel() for tensor in model.parameters()) / 1e6)) # Load pretrained state for cityscapes in GLOBAL net if args.pretrained and not args.resume: if not args.load_external_mod: if not args.multi: target_state = model.depthnet.state_dict() else: target_state = model.module.depthnet.state_dict() check = torch.load('erfnet_pretrained.pth') for name, val in check.items(): # Exclude multi GPU prefix mono_name = name[7:] if mono_name not in target_state: continue try: target_state[mono_name].copy_(val) except RuntimeError: continue print('Successfully loaded pretrained model') else: check = torch.load('external_mod.pth.tar') lowest_loss_load = check['loss'] target_state = model.state_dict() for name, val in check['state_dict'].items(): if name not in target_state: continue try: target_state[name].copy_(val) except RuntimeError: continue print("=> loaded EXTERNAL checkpoint with best rmse {}".format( lowest_loss_load)) # Start training for epoch in range(args.start_epoch, args.nepochs): print("\n => Start EPOCH {}".format(epoch + 1)) print(datetime.now().strftime('%Y-%m-%d %H:%M:%S')) print(args.save_path) # Adjust learning rate if args.lr_policy is not None and args.lr_policy != 'plateau': scheduler.step() lr = optimizer.param_groups[0]['lr'] print('lr is set to {}'.format(lr)) # Define container objects batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() score_train = AverageMeter() score_train_1 = AverageMeter() metric_train = Metrics(max_depth=args.max_depth, disp=args.use_disp, normal=args.normal) # Train model for args.nepochs model.train() # compute timing end = time.time() # Load dataset for i, (input, gt) in tqdm(enumerate(train_loader)): # Time dataloader data_time.update(time.time() - end) # Put inputs on gpu if possible if not args.no_cuda: input, gt = input.cuda(), gt.cuda() prediction, lidar_out, precise, guide = model(input, epoch) loss = criterion_local(prediction, gt) loss_lidar = criterion_lidar(lidar_out, gt) loss_rgb = criterion_rgb(precise, gt) loss_guide = criterion_guide(guide, gt) loss = args.wpred * loss + args.wlid * loss_lidar + args.wrgb * loss_rgb + args.wguide * loss_guide losses.update(loss.item(), input.size(0)) metric_train.calculate(prediction[:, 0:1].detach(), gt.detach()) score_train.update(metric_train.get_metric(args.metric), metric_train.num) score_train_1.update(metric_train.get_metric(args.metric_1), metric_train.num) # Clip gradients (usefull for instabilities or mistakes in ground truth) if args.clip_grad_norm != 0: nn.utils.clip_grad_norm(model.parameters(), args.clip_grad_norm) # Setup backward pass optimizer.zero_grad() loss.backward() optimizer.step() # Time trainig iteration batch_time.update(time.time() - end) end = time.time() # Print info if (i + 1) % args.print_freq == 0: print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 'Metric {score.val:.4f} ({score.avg:.4f})'.format( epoch + 1, i + 1, len(train_loader), batch_time=batch_time, loss=losses, score=score_train)) print("===> Average RMSE score on training set is {:.4f}".format( score_train.avg)) print("===> Average MAE score on training set is {:.4f}".format( score_train_1.avg)) # Evaulate model on validation set print("=> Start validation set") score_valid, score_valid_1, losses_valid = validate( valid_loader, model, criterion_lidar, criterion_rgb, criterion_local, criterion_guide, epoch) print("===> Average RMSE score on validation set is {:.4f}".format( score_valid)) print("===> Average MAE score on validation set is {:.4f}".format( score_valid_1)) # Evaluate model on selected validation set if args.subset is None: print("=> Start selection validation set") score_selection, score_selection_1, losses_selection = validate( valid_selection_loader, model, criterion_lidar, criterion_rgb, criterion_local, criterion_guide, epoch) total_score = score_selection print("===> Average RMSE score on selection set is {:.4f}".format( score_selection)) print("===> Average MAE score on selection set is {:.4f}".format( score_selection_1)) else: total_score = score_valid print("===> Last best score was RMSE of {:.4f} in epoch {}".format( lowest_loss, best_epoch)) # Adjust lr if loss plateaued if args.lr_policy == 'plateau': scheduler.step(total_score) lr = optimizer.param_groups[0]['lr'] print('LR plateaued, hence is set to {}'.format(lr)) # File to keep latest epoch with open(os.path.join(args.save_path, 'first_run.txt'), 'w') as f: f.write(str(epoch)) # Save model to_save = False if total_score < lowest_loss: to_save = True best_epoch = epoch + 1 lowest_loss = total_score save_checkpoint( { 'epoch': epoch + 1, 'best epoch': best_epoch, 'arch': args.mod, 'state_dict': model.state_dict(), 'loss': lowest_loss, 'optimizer': optimizer.state_dict() }, to_save, epoch) if not args.no_tb: writer.close()
def main(): global args args = parser.parse_args() if args.num_samples == 0: # Use all lidar points args.num_samples = None else: args.data_path = "" # path to precomputed 500 samples assert args.num_samples == 500 print("changed path to samples500 dataset") if args.val_batch_size is None: args.val_batch_size = args.batch_size if args.seed: random.seed(args.seed) torch.manual_seed(args.seed) # init_distributed_mode(args) if not args.no_cuda and not torch.cuda.is_available(): raise Exception("No gpu available for usage") torch.backends.cudnn.benchmark = args.cudnn # Init model args.channels_in = 3 if args.input_type == 'rgb' else 4 model = Models.define_model(args.mod, args) # define_init_weights(model, args.weight_init) # Load on gpu before passing params to optimizer if not args.no_cuda: if not args.multi: model = model.cuda() else: model = torch.nn.DataParallel(model).cuda() # model.cuda() # model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) # model = model.module save_id = '{}_{}_{}_{}_{}_batch{}_pretrain{}_wrgb{}_drop{}_patience{}_num_samples{}_multi{}_submod{}'.\ format(args.mod, args.optimizer, args.loss_criterion, args.learning_rate, args.input_type, args.batch_size, args.pretrained, args.wrgb, args.drop, args.lr_decay_iters, args.num_samples, args.multi, args.submod) # INIT optimizer/scheduler/loss criterion optimizer = define_optim(args.optimizer, model.parameters(), args.learning_rate, args.weight_decay) scheduler = define_scheduler(optimizer, args) # Optional to use different losses criterion_local = define_loss(args.loss_criterion) criterion_rgb = define_loss(args.loss_criterion) # INIT dataset dataset = Datasets.define_dataset(args.dataset, args.data_path, args.input_type) dataset.prepare_dataset() train_loader, _, valid_loader, valid_selection_loader = get_loader( args, dataset) # Resume training best_epoch = 0 lowest_loss = np.inf args.save_path = os.path.join(args.save_path, save_id) mkdir_if_missing(args.save_path) log_file_name = 'log_train_start_0.txt' args.resume = first_run(args.save_path) if args.resume and not args.test_mode and not args.evaluate: path = os.path.join( args.save_path, 'checkpoint_model_epoch_{}.pth.tar'.format(int(args.resume))) if os.path.isfile(path): log_file_name = 'log_train_start_{}.txt'.format(args.resume) # stdout sys.stdout = Logger(os.path.join(args.save_path, log_file_name)) print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(path) args.start_epoch = checkpoint['epoch'] lowest_loss = checkpoint['loss'] best_epoch = checkpoint['best epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: log_file_name = 'log_train_start_0.txt' # stdout sys.stdout = Logger(os.path.join(args.save_path, log_file_name)) print("=> no checkpoint found at '{}'".format(path)) # Only evaluate elif args.evaluate: print("Evaluate only") best_file_lst = glob.glob(os.path.join(args.save_path, 'model_best*')) if len(best_file_lst) != 0: best_file_name = best_file_lst[0] print(best_file_name) if os.path.isfile(best_file_name): sys.stdout = Logger( os.path.join(args.save_path, 'Evaluate.txt')) print("=> loading checkpoint '{}'".format(best_file_name)) checkpoint = torch.load(best_file_name) model.load_state_dict(checkpoint['state_dict']) else: print("=> no checkpoint found at '{}'".format(best_file_name)) else: print("=> no checkpoint found at due to empy list in folder {}". format(args.save_path)) validate(valid_selection_loader, model, criterion_global, criterion_local) return # Start training from clean slate else: # Redirect stdout sys.stdout = Logger(os.path.join(args.save_path, log_file_name)) # INIT MODEL print(40 * "=" + "\nArgs:{}\n".format(args) + 40 * "=") print("Init model: '{}'".format(args.mod)) print("Number of parameters in model {} is {:.3f}M".format( args.mod.upper(), sum(tensor.numel() for tensor in model.parameters()) / 1e6)) # Load pretrained state for cityscapes in GLOBAL net # if args.pretrained and not args.resume: # if not args.multi: # target_state = model.depthnet.state_dict() # else: # target_state = model.module.depthnet.state_dict() # check = torch.load('erfnet_pretrained.pth') # for name, val in check.items(): # # Exclude multi GPU prefix # mono_name = name[7:] # if mono_name not in target_state: # continue # try: # target_state[mono_name].copy_(val) # except RuntimeError: # continue # print('Successfully loaded pretrained model') # Create summary writer log_path = os.path.join(args.save_path, "logs") if not os.path.exists(log_path): os.makedirs(log_path) logger = SummaryWriter(log_path) with open(os.path.join(args.save_path, 'commandline_args.txt'), 'w') as f: json.dump(args.__dict__, f, indent=2) # Start training for epoch in range(args.start_epoch, args.nepochs): print("\n => Start EPOCH {}".format(epoch + 1)) print(datetime.now().strftime('%Y-%m-%d %H:%M:%S')) print(args.save_path) # Adjust learning rate if args.lr_policy is not None and args.lr_policy != 'plateau': scheduler.step() lr = optimizer.param_groups[0]['lr'] print('lr is set to {}'.format(lr)) # Define container objects batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() score_train = AverageMeter() score_train_1 = AverageMeter() metric_train = Metrics(max_depth=args.max_depth, disp=args.use_disp) # Train model for args.nepochs model.train() # compute timing end = time.time() # Load dataset for i, (input, gt) in tqdm(enumerate(train_loader)): # Time dataloader data_time.update(time.time() - end) # Put inputs on gpu if possible if not args.no_cuda: input, gt = input.cuda(), gt.cuda() prediction, hidden = model(input, hidden=(None, None)) if 'mod' in args.mod or 'stacked' in args.mod: loss = criterion_local(prediction[0], gt) loss_rgb = criterion_rgb(prediction[1], gt) loss += args.wrgb * loss_rgb prediction = prediction[0] else: loss = criterion_local(prediction, gt) losses.update(loss.item(), input.size(0)) metric_train.calculate(prediction[:, 0:1].detach(), gt.detach()) score_train.update(metric_train.get_metric(args.metric), metric_train.num) score_train_1.update(metric_train.get_metric(args.metric_1), metric_train.num) # Clip gradients (usefull for instabilities or mistakes in ground truth) if args.clip_grad_norm != 0: nn.utils.clip_grad_norm(model.parameters(), args.clip_grad_norm) # Setup backward pass optimizer.zero_grad() loss.backward() optimizer.step() # Time trainig iteration batch_time.update(time.time() - end) end = time.time() # Print info if (i + 1) % args.print_freq == 0: print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 'Metric {score.val:.4f} ({score.avg:.4f})'.format( epoch + 1, i + 1, len(train_loader), batch_time=batch_time, loss=losses, score=score_train)) batch_num = len(train_loader) if ((i + 1) % 10 == 0) and (logger is not None): current_step = epoch * batch_num + i # Add scalar summaries logger.add_scalar('Train_loss/Loss', loss.item(), current_step) # utils.record_scalar_summary(result, average_meter, current_step, logger, "Train") # # Add system info # logger.add_scalar('System/gpu_time', average_meter.average().gpu_time, current_step) # logger.add_scalar('System/data_time', average_meter.average().data_time, current_step) if ((i + 1) % 200 == 0): # Add some image summary if args.modality == "rgb": input_images = input.cpu() else: input_images = input[:, :3, :, :].cpu() input_depth = torch.unsqueeze(input[:, 3, :, :], dim=1).cpu() rgb_grid = make_grid(input_images[0:6, :, :, :], nrow=3, normalize=True), target_grid = make_grid(gt.cpu()[0:6, :, :, :], nrow=3, normalize=True) pred_grid = make_grid(prediction.cpu()[0:6, :, :, :], nrow=3, normalize=True) logger.add_image('Train/RGB', rgb_grid[0].data.numpy(), current_step) logger.add_image('Train/Depth_gt', target_grid.data.numpy(), current_step) logger.add_image('Train/Depth_pred', pred_grid.data.numpy(), current_step) if args.modality == "rgbd": depth_grid = make_grid(input_depth[0:6, :, :, :], nrow=3, normalize=True) logger.add_image('Train/Depth_input', depth_grid.data.numpy(), current_step) print("===> Average RMSE score on training set is {:.4f}".format( score_train.avg)) print("===> Average MAE score on training set is {:.4f}".format( score_train_1.avg)) # Evaulate model on validation set print("=> Start validation set") score_valid, score_valid_1, losses_valid = validate( valid_loader, model, criterion_rgb, criterion_local, epoch) print("===> Average RMSE score on validation set is {:.4f}".format( score_valid)) print("===> Average MAE score on validation set is {:.4f}".format( score_valid_1)) # Evaluate model on selected validation set if args.subset is None: print("=> Start selection validation set") score_selection, score_selection_1, losses_selection = validate( valid_selection_loader, model, criterion_rgb, criterion_local, epoch) total_score = score_selection print("===> Average RMSE score on selection set is {:.4f}".format( score_selection)) print("===> Average MAE score on selection set is {:.4f}".format( score_selection_1)) else: total_score = score_valid print("===> Last best score was RMSE of {:.4f} in epoch {}".format( lowest_loss, best_epoch)) # Adjust lr if loss plateaued if args.lr_policy == 'plateau': scheduler.step(total_score) lr = optimizer.param_groups[0]['lr'] print('LR plateaued, hence is set to {}'.format(lr)) # File to keep latest epoch with open(os.path.join(args.save_path, 'first_run.txt'), 'w') as f: f.write(str(epoch)) # Save model to_save = False if total_score < lowest_loss: to_save = True best_epoch = epoch + 1 lowest_loss = total_score save_checkpoint( { 'epoch': epoch + 1, 'best epoch': best_epoch, 'arch': args.mod, 'state_dict': model.state_dict(), 'loss': lowest_loss, 'optimizer': optimizer.state_dict() }, to_save, epoch) if not args.no_tb: writer.close()