def run(self): self.model.eval() with torch.no_grad(): for batch_idx, data in enumerate(self.dataloader): bg, fgbg, mask, depth = data['bg'].to( self.model.device), data['fgbg'].to( self.model.device), data['mask'].to( self.model.device), data['depth'].to( self.model.device) mask_pred, depth_pred = self.model(bg, fgbg) # Calculate loss if self.criterion1 is not None: loss1 = self.criterion1(mask_pred, mask) m_ssim = torch.clamp((1 - ssim(mask_pred, mask)) * 0.5, 0, 1) loss1 = (0.84 * m_ssim) + (0.16 * loss1) if self.criterion2 is not None: loss2 = self.criterion2(depth_pred, depth) d_ssim = torch.clamp((1 - ssim(depth_pred, depth)) * 0.5, 0, 1) loss2 = (0.84 * d_ssim) + (0.16 * loss2) self.loss = 2 * loss1 + loss2 if batch_idx == 0: inp = fgbg.detach().cpu() orimp = mask.detach().cpu() mp = mask_pred.detach().cpu() oridp = depth.detach().cpu() dp = depth_pred.detach().cpu() print( "First batch in testing fgbg, (mask, predicted mask), (depth, predicted depth)" ) show(inp[:8, :, :, :], normalize=True) mdinp = torch.cat([ orimp[:8, :, :, :], mp[:8, :, :, :], oridp[:8, :, :, :], dp[:8, :, :, :] ], dim=0) show(mdinp) self.stats.add_batch_test_stats(self.loss.item(), 0, len(data)) if self.scheduler and isinstance( self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): #print("hello yes i am ") self.scheduler.step(self.loss)
def validate(loader, model, epoch, d1, d2, blind, noise_level): val_psnr = 0 val_ssim = 0 val_l1 = 0 model.train(False) k1 = model.weight[0].unsqueeze(0).expand(loader.batch_size, -1, -1, -1) k2 = model.weight[1].unsqueeze(0).expand(loader.batch_size, -1, -1, -1) d1 = d1.expand(loader.batch_size, -1, -1, -1) d2 = d2.expand(loader.batch_size, -1, -1, -1) # pre-create noise levels if blind: nls = np.linspace(0.5, noise_level, len(loader)) else: nls = noise_level * np.ones(len(loader)) with torch.no_grad(): for i, data in tqdm.tqdm(enumerate(loader)): x, y, k, d = data x = x.to(device) y = y.to(device) k = k.to(device) d = d.to(device) nl = nls[i] / 255 y += nl * torch.randn_like(y) y = y.clamp(0, 1) hat_x = model(y, k, d, k1, k2, d1, d2)[-1] hat_x.clamp_(0, 1) hat_x = utils.crop_valid(hat_x, k) x = utils.crop_valid(x, k) y = utils.crop_valid(y, k) val_psnr += loss.psnr(hat_x, x) val_ssim += loss.ssim(hat_x, x) val_l1 += F.l1_loss(hat_x, x).item() return val_psnr / len(loader), val_ssim / len(loader), val_l1 / len(loader)
def validate(loader, model, epoch, d1, d2, blind, noise_level): val_psnr = 0 val_ssim = 0 val_l1 = 0 model.train(False) k1 = model.weight[0].unsqueeze(0).expand(loader.batch_size, -1, -1, -1) k2 = model.weight[1].unsqueeze(0).expand(loader.batch_size, -1, -1, -1) d1 = d1.expand(loader.batch_size, -1, -1, -1) d2 = d2.expand(loader.batch_size, -1, -1, -1) # pre-create noise levels if blind: nls = np.linspace(0.5, noise_level, len(loader)) else: nls = noise_level * np.ones(len(loader)) with torch.no_grad(): for i, data in tqdm.tqdm(enumerate(loader)): x, y, mag, ori = data x = x.to(device) y = y.to(device) mag = mag.to(device) ori = ori.to(device) ori = (90 - ori).add(360).fmod(180) labels = utils.get_labels(mag, ori) ori = ori * np.pi / 180 nl = nls[i] / 255 y += nl * torch.randn_like(y) y = y.clamp(0, 1) hat_x = model(y, mag, ori, labels, k1, k2, d1, d2)[-1] hat_x.clamp_(0, 1) val_psnr += loss.psnr(hat_x, x) val_ssim += loss.ssim(hat_x, x) val_l1 += F.l1_loss(hat_x, x).item() return val_psnr / len(loader), val_ssim / len(loader), val_l1 / len(loader)
def eval(self): self.model.eval() psnr_losses = [] ssim_losses = [] with tqdm(total=len(self.valloader)) as t: # show batch evaluate process t.set_description('evaluating...') for LR, HR in self.valloader: LR = LR.to(self.device) HR = HR.to(self.device) output = self.model(LR) # # only calculation on y channel or gray # if self.r_mode == 'RGB' and self.img_channels == 3: # output = rgb2y_tensor(output) # HR = rgb2y_tensor(HR) ssim_loss = ssim(output, HR) psnr_loss = psnr(output, HR) # save losses ssim_losses.append(ssim_loss) psnr_losses.append(psnr_loss) t.update() avg_ssim = torch.stack(ssim_losses, dim=0).mean().item() avg_psnr = torch.stack(psnr_losses, dim=0).mean().item() t.set_postfix(avg_psnr=f'{avg_psnr:.010f}', avg_ssim=f'{avg_ssim:.010f}') t.set_description('evaluate') self.model.train() return avg_ssim, avg_psnr
def LogProgress(model, writer, test_loader, global_step): with torch.no_grad(): model.eval() sequential = test_loader _input, _label = next(iter(sequential)) image = torch.autograd.Variable(_input.float().cuda()) depth = torch.autograd.Variable(_label.float().cuda(non_blocking=True)) # Normalize depth # depth_n = DepthNorm(depth) # print() # if epoch == 0: # writer.add_image('Train.1.Image', vutils.make_grid(image.data, nrow=6, normalize=True), epoch) # if epoch == 0: output = model(image) # Predict # output = model(image) # edge = sobel_(output) # pred_edge = sobel_(depth) # Compute the loss # l_sobel = nn.L1Loss()(edge, pred_edge) l_depth = nn.L1Loss()(output, depth) l_ssim = torch.clamp((1 - ssim(output, depth, val_range=10.0)) * 0.5, 0, 1) # if torch.isnan(l_sobel): # print(1) # print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim)) writer.add_scalar('Test/L1', l_depth.item(), global_step) writer.add_scalar('Test/SSIM', l_ssim.item(), global_step) # writer.add_scalar('Test/EDGE', l_sobel.item(), global_step) # normal = Norm(output) # normal = (normal + 1.0) / 2. # normal = normal[3, :, :, :] # normal = normal.detach().cpu().numpy().astype(np.uint8) # normal = np.transpose(normal, (1, 2, 0)) # print() # plt.imshow(normal) # plt.show() # output = DepthNorm(output) # writer.add_image('Train.3.Normal', vutils.make_grid(normal.data, nrow=6, normalize=False), epoch) _fig = colorize(image.data, depth.data, output.data) # _ = vutils.make_grid(_, nrow=1, padding=0, normalize=False) # res1 = torch.cat([image[0:1, 3:4].data, depth[0:1].data, output[0:1].data], dim=0) # xx = colorize(vutils.make_grid(res1, nrow=6, padding=0, normalize=False)) # rgb = image[0:1, :3, :, :].data.cpu().numpy() * 255. # rgb = rgb.astype(np.uint8).squeeze() # out = np.concatenate([rgb, xx], axis=2) # print() writer.add_image('Train.2.test', _fig.transpose((2, 0, 1)), global_step) # writer.add_image('Train.2.Raw', colorize(vutils.make_grid(image[:,3:4,:,:].data, nrow=6, normalize=False)), # global_step) # writer.add_image('Train.2.Depth', colorize(vutils.make_grid(depth.data, nrow=6, normalize=False)), global_step) # writer.add_image('Test.2.Ours', colorize(vutils.make_grid(output.data, nrow=6, normalize=False)), global_step) # writer.add_image('Test.3.Diff', colorize(vutils.make_grid(torch.abs(output - depth).data, nrow=6, normalize=False)), # global_step) del image del depth del output
def run_epoch(self, epoch, dataloader, logimage=False, isTrain=True): # For details see training. psnr_value = 0 ssim_value = 0 loss_value = 0 if not isTrain: valid_images = [] for index, all_data in enumerate(dataloader, 0): self.optimizer.zero_grad() ( Ft_p, I0, IFrame, I1, g_I0_F_t_0, g_I1_F_t_1, FlowBackWarp_I0_F_1_0, FlowBackWarp_I1_F_0_1, F_1_0, F_0_1, ) = self.slomo(all_data, pred_only=False, isTrain=isTrain) if (not isTrain) and logimage: if index % self.args.logimagefreq == 0: valid_images.append( 255.0 * I0.cpu()[0] .resize_(1, 1, self.args.data_h, self.args.data_w) .repeat(1, 3, 1, 1) ) valid_images.append( 255.0 * IFrame.cpu()[0] .resize_(1, 1, self.args.data_h, self.args.data_w) .repeat(1, 3, 1, 1) ) valid_images.append( 255.0 * I1.cpu()[0] .resize_(1, 1, self.args.data_h, self.args.data_w) .repeat(1, 3, 1, 1) ) valid_images.append( 255.0 * Ft_p.cpu()[0] .resize_(1, 1, self.args.data_h, self.args.data_w) .repeat(1, 3, 1, 1) ) # loss loss = self.supervisedloss( Ft_p, IFrame, I0, I1, g_I0_F_t_0, g_I1_F_t_1, FlowBackWarp_I0_F_1_0, FlowBackWarp_I1_F_0_1, F_1_0, F_0_1, ) if isTrain: loss.backward() self.optimizer.step() self.scheduler.step() loss_value += loss.item() # metrics psnr_value += psnr(Ft_p, IFrame, outputTensor=False) ssim_value += ssim(Ft_p, IFrame, outputTensor=False) name_loss = "TrainLoss" if isTrain else "ValLoss" itr = int(index + epoch * (len(dataloader))) if self.comet_exp is not None: self.comet_exp.log_metric( "PSNR", psnr_value / len(dataloader), step=itr, epoch=epoch ) self.comet_exp.log_metric( "SSIM", ssim_value / len(dataloader), step=itr, epoch=epoch ) self.comet_exp.log_metric( name_loss, loss_value / len(dataloader), step=itr, epoch=epoch ) if logimage: upload_images( valid_images, epoch, exp=self.comet_exp, im_per_row=4, rows_per_log=int(len(valid_images) / 4), ) print( " Loss: %0.6f Iterations: %4d/%4d ValPSNR: %0.4f ValSSIM: %0.4f " % ( loss_value / len(dataloader), index, len(dataloader), psnr_value / len(dataloader), ssim_value / len(dataloader), ) ) return ( (psnr_value / len(dataloader)), (ssim_value / len(dataloader)), (loss_value / len(dataloader)), )
def main(): # Arguments parser = argparse.ArgumentParser(description='High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=20, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate') parser.add_argument('--bs', default=4, type=int, help='batch size') args = parser.parse_args() SAVE_DIR = 'models/191107_mod15' ifcrop = True if ifcrop: HEIGHT = 256 HEIGHT_WITH_RATIO = 240 WIDTH = 320 else: HEIGHT = 480 WIDTH = 640 with torch.cuda.device(0): # Create model # model = Model().cuda() # print('Model created.') # ============================================================================= # load saved model # model = Model_rgbd().cuda() # model.load_state_dict(torch.load(os.path.join(SAVE_DIR, 'model_overtraining.pth'))) # model.eval() # print('model loaded for evaluation.') # ============================================================================= # Create RGB-D model model = Model_rgbd().cuda() print('Model created.') # ============================================================================= # Training parameters optimizer = torch.optim.Adam( model.parameters(), args.lr, amsgrad=True ) batch_size = args.bs prefix = 'densenet_' + str(batch_size) # Load data train_loader, test_loader = getTrainingTestingData(batch_size=1, crop_halfsize=ifcrop) train_loader_l, test_loader_l = getTranslucentData(batch_size=1, crop_halfsize=ifcrop) # Test batch is manually enlarged! See getTranslucentData's return. # Logging writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(prefix, args.lr, args.epochs, args.bs), flush_secs=30) # Loss l1_criterion = nn.L1Loss() l1_criterion_masked = MaskedL1() grad_l1_criterion_masked = MaskedL1Grad() # Hand-craft loss weight of main task interval1 = 1 interval2 = 2 weight_t1loss = [1] * (10*interval1) + [0] * interval2 weight_txloss = [.0317] * interval1 + [.1] * interval1 + \ [.316] * interval1 + [1] * interval1 + \ [3.16] * interval1 + [10] * interval1 + \ [10] * interval1 + [5.62] * interval1 + \ [3.16] * interval1 + [1.78] * interval1 + \ [0] * interval2 weight_t2loss = [.001] * interval1 + [.00316] * interval1 + \ [.01] * interval1 + [.0316] * interval1 + \ [.1] * interval1 + [.316] * interval1 + \ [1.0] * interval1 + [3.16] * interval1 + \ [10.0] * interval1 + [31.6] * interval1 + \ [100.0] * interval2 if not os.path.exists('%s/img' % SAVE_DIR): os.makedirs('%s/img' % SAVE_DIR) # Start training... for epoch in range(0, 10*interval1 + interval2): batch_time = AverageMeter() losses_nyu = AverageMeter() losses_lucent = AverageMeter() losses_hole = AverageMeter() losses = AverageMeter() N = len(train_loader) # Switch to train mode model.train() end = time.time() # decide #(iter) tot_len = min(len(train_loader), len(train_loader_l)) # print(tot_len) trainiter = iter(train_loader) trainiter_l = iter(train_loader_l) for i in range(tot_len): # print("Iteration "+str(i)+". loop start:") try: sample_batched = next(trainiter) sample_batched_l = next(trainiter_l) except StopIteration: print(' (almost) end of iteration.') continue # print('in loop.') # Prepare sample and target image_nyu = torch.autograd.Variable(sample_batched['image'].cuda()) depth_nyu = torch.autograd.Variable(sample_batched['depth'].cuda(non_blocking=True)) image_raw = torch.autograd.Variable(sample_batched_l['image'].cuda()) mask_raw = torch.autograd.Variable(sample_batched_l['mask'].cuda()) depth_raw = torch.autograd.Variable(sample_batched_l['depth_raw'].cuda()) depth_gt = torch.autograd.Variable(sample_batched_l['depth_truth'].cuda(non_blocking=True)) # if i < 10: # print('========-=-=') # print(image_nyu.shape) # print(depth_nyu.shape) # print(image_raw.shape) # print(" " + str(torch.max(image_raw)) + " " + str(torch.min(image_raw))) # print(mask_raw.shape) # print(" " + str(torch.max(mask_raw)) + " " + str(torch.min(mask_raw))) # print(depth_raw.shape) # print(" " + str(torch.max(depth_raw)) + " " + str(torch.min(depth_raw))) # print(depth_gt.shape) # print(" " + str(torch.max(depth_gt)) + " " + str(torch.min(depth_gt))) N1 = image_nyu.shape[0] N2 = image_raw.shape[0] ########################### # (1) Pretext task: depth completion # if weight_t1loss[epoch] > 0: # Normalize depth depth_nyu_n = DepthNorm(depth_nyu) # Apply random mask to it rand_index = [random.randint(0, N2-1) for k in range(N1)] mask_new = mask_raw[rand_index, :, :, :] depth_nyu_masked = resize2d(depth_nyu_n, (HEIGHT, WIDTH)) * mask_new # if i < 1: # print('========') # print(image_nyu.shape) # print(" " + str(torch.max(image_raw)) + " " + str(torch.min(image_raw))) # print(depth_nyu_masked.shape) # print(" " + str(torch.max(depth_nyu_masked)) + " " + str(torch.min(depth_nyu_masked))) # Predict (output_t1, _) = model(image_nyu, depth_nyu_masked) # print(" (1): " + str(output_task1.shape)) # Calculate Loss and backprop l_depth_t1 = l1_criterion(output_t1, depth_nyu_n) l_grad_t1 = l1_criterion(grad_x(output_t1), grad_x(depth_nyu_n)) + l1_criterion(grad_y(output_t1), grad_y(depth_nyu_n)) l_ssim_t1 = torch.clamp((1 - ssim(output_t1, depth_nyu_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss_nyu = (0.1 * l_depth_t1) + (1.0 * l_grad_t1) + (1.0 * l_ssim_t1) # loss_nyu_weighted = weight_t1loss[epoch] * loss_nyu # https://discuss.pytorch.org/t/freeze-the-learnable-parameters-of-resnet-and-attach-it-to-a-new-network/949 # freeze_weight(model, e_stay=False, e=False, d1_stay=False, d1=True) # optimizer.zero_grad() # moved to its new position # loss_nyu_weighted.backward(retain_graph=True) # optimizer.step() if i % 150 == 0 or i < 2: vutils.save_image(DepthNorm(depth_nyu_masked), '%s/img/A_masked_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True) vutils.save_image(DepthNorm(output_t1), '%s/img/A_out_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True) save_error_image(DepthNorm(output_t1) - depth_nyu, '%s/img/A_diff_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True, range=(-500, 500)) torch.cuda.empty_cache() ########################### # (x) Transfer task: /*Fill*/ reconstruct sudo-translucent object depth_gt_n = DepthNorm(depth_gt) depth_raw_n = DepthNorm(depth_raw) # if weight_txloss[epoch] > 0: # Normalize depth depth_gt_large = resize2d(depth_gt, (HEIGHT, WIDTH)) object_mask = thresh_mask(depth_gt_large, depth_raw) # depth_holed = depth_raw * object_mask depth_holed = blend_depth(depth_raw, depth_gt_large, object_mask) # print('========') # print(object_mask.shape) # print(" " + str(torch.max(object_mask)) + " " + str(torch.min(object_mask))) # print(depth_holed.shape) # print(" " + str(torch.max(depth_holed)) + " " + str(torch.min(depth_holed))) # print(image_raw.shape) # print(" " + str(torch.max(image_raw)) + " " + str(torch.min(image_raw))) (output_tx, _) = model(image_raw, DepthNorm(depth_holed)) output_tx_n = DepthNorm(output_tx) # Calculate Loss and backprop mask_post = resize2dmask(mask_raw, (int(HEIGHT/2), int(WIDTH/2))) l_depth_tx = l1_criterion_masked(output_tx, depth_gt_n, mask_post) l_grad_tx = grad_l1_criterion_masked(output_tx, depth_gt_n, mask_post) # l_ssim_tx = torch.clamp((1 - ssim(output_tx, depth_nyu_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss_hole = (0.1 * l_depth_tx) + (1.0 * l_grad_tx) #+ (0 * l_ssim_tx) #### # loss_hole_weighted = weight_txloss[epoch] * loss_hole # for param in model.decoder1.parameters(): # param.requires_grad = False # for param in model.decoder2.parameters(): # param.requires_grad = True # freeze_weight(model, d1_stay=False, d1=False, d2_stay=False, d2=True) # optimizer.zero_grad() # loss_hole_weighted.backward(retain_graph=True) # https://pytorch.org/docs/stable/autograd.html # optimizer.step() if i % 150 == 0 or i < 2: vutils.save_image(DepthNorm(depth_holed), '%s/img/C_in_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True, range=(0, 500)) vutils.save_image(object_mask, '%s/img/C_mask_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True, range=(0, 1.5)) vutils.save_image(output_tx_n, '%s/img/C_out_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True, range=(0, 500)) save_error_image(output_tx_n - depth_gt, '%s/img/C_zdiff_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True, range=(-500, 500)) torch.cuda.empty_cache() ########################### # (2) Main task: Undistort translucent object # Predict (_, output_t2) = model(image_raw, depth_raw_n) output_t2_n = DepthNorm(output_t2) # print(" (2): " + str(output.shape)) # Calculate Loss and backprop l_depth_t2 = l1_criterion_masked(output_t2, depth_gt_n, mask_post) l_grad_t2 = grad_l1_criterion_masked(output_t2, depth_gt_n, mask_post) # l_ssim_t2 = torch.clamp((1 - ssim(output_t2, depth_gt_n, val_range=1000.0/10.0)) * 0.5, 0, 1) loss_lucent = (0.1 * l_depth_t2) + (1.0 * l_grad_t2) # + (0 * l_ssim_t2) # loss_lucent_weighted = weight_t2loss[epoch] * loss_lucent # optimizer.zero_grad() # moved to its new position # loss_lucent_weighted.backward(retain_graph=True) # optimizer.step() if i % 150 == 0 or i < 2: vutils.save_image(depth_raw, '%s/img/B_ln_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(0, 500)) vutils.save_image(depth_gt, '%s/img/B_gt_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(0, 500)) vutils.save_image(output_t2_n, '%s/img/B_out_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(0, 500)) save_error_image(output_t2_n-depth_gt, '%s/img/B_zdiff_%06d.png' % (SAVE_DIR, epoch*10000+i), normalize=True, range=(-500, 500)) if i % 150 == 0 : o2 = output_t2.cpu().detach().numpy() o3 = output_t2_n.cpu().detach().numpy() og = depth_gt.cpu().numpy() nm = DepthNorm(depth_nyu_masked).cpu().numpy() ng = depth_nyu.cpu().numpy() print('> ========') print("> Output_t2:" + str(np.min(o2)) + "~" + str(np.max(o2)) + " (" + str(np.mean(o2)) + ") // Converted to depth: " + str(np.min(o3)) + "~" + str(np.max(o3)) + " (" + str(np.mean(o3)) + ")") print("> GT depth : " + str(np.min(og)) + "~" + str(np.max(og)) + " // NYU GT depth from 0.0~" + str(np.max(nm)) + " to " + str(np.min(ng)) + "~" + str(np.max(ng)) + " (" + str(np.mean(ng)) + ")") ########################### # (3) Update the network parameters if i % 150 == 0 or i < 1: vutils.save_image(mask_post, '%s/img/_mask_%06d.png' % (SAVE_DIR, epoch * 10000 + i), normalize=True) loss = (weight_t1loss[epoch] * loss_nyu) + (weight_t2loss[epoch] * loss_lucent) + (weight_txloss[epoch] * loss_hole) #### # freeze_weight(model, e_stay=False, e=True, d2_stay=False, d2=False) optimizer.zero_grad() loss.backward() optimizer.step() # loss = loss_nyu + loss_lucent + loss_hole # Log losses losses_nyu.update(loss_nyu.detach().item(), image_nyu.size(0)) losses_lucent.update(loss_lucent.detach().item(), image_raw.size(0)) losses_hole.update(loss_hole.detach().item(), image_raw.size(0)) losses.update(loss.detach().item(), image_nyu.size(0) + image_raw.size(0)) # Measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i)))) # Log progress niter = epoch*N+i if i % 15 == 0: # Print to console print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' 'ETA {eta}\t' 'Loss {loss.val:.4f} ({loss.avg:.4f}) ||\t' 'NYU {l1.val:.4f} ({l1.avg:.4f}) [{l1d:.4f} | {l1g:.4f} | {l1s:.4f}]\t' 'LUC {l2.val:.4f} ({l2.avg:.4f}) [{l2d:.4f} | {l2g:.4f}]\t' 'TX {lx.val:.4f} ({lx.avg:.4f}) [{lxd:.4f} | {lxg:.4f}]' .format(epoch, i, N, batch_time=batch_time, loss=losses, l1=losses_nyu, l1d=l_depth_t1, l1g=l_grad_t1, l1s=l_ssim_t1, l2=losses_lucent, l2d=l_depth_t2, l2g=l_grad_t2, lx=losses_hole, lxd=l_depth_tx, lxg=l_grad_tx, eta=eta)) # Note that the numbers displayed are pre-weighted. # Log to tensorboard writer.add_scalar('Train/Loss', losses.val, niter) if i % 750 == 0: LogProgress(model, writer, test_loader, test_loader_l, niter, epoch*10000+i, SAVE_DIR, HEIGHT, WIDTH) path = os.path.join(SAVE_DIR, 'model_overtraining.pth') torch.save(model.cpu().state_dict(), path) # saving model model.cuda() # moving model to GPU for further training del image_nyu, depth_nyu_masked, output_t1, image_raw, depth_raw_n, output_t2 torch.cuda.empty_cache() # Record epoch's intermediate results LogProgress(model, writer, test_loader, test_loader_l, niter, epoch*10000+i, SAVE_DIR, HEIGHT, WIDTH) writer.add_scalar('Train/Loss.avg', losses.avg, epoch) # all the saves come from https://discuss.pytorch.org/t/how-to-save-a-model-from-a-previous-epoch/20252 torch.save(model.state_dict(), os.path.join(SAVE_DIR, 'epoch-{}.pth'.format(epoch))) print('Program terminated.')
def main(): # Arguments parser = argparse.ArgumentParser( description= 'High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=20, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate') parser.add_argument('--bs', default=4, type=int, help='batch size') args = parser.parse_args() # Create a new directory to save logs runs = sorted(glob.glob(os.path.join('logs', 'exp-*'))) prev_run_id = int(runs[-1].split('-')[-1]) if runs else 0 MODEL_LOG_DIR = os.path.join('logs', 'exp-{:03d}'.format(prev_run_id + 1)) CHECKPOINT_DIR = os.path.join(MODEL_LOG_DIR, 'checkpoints') os.makedirs(CHECKPOINT_DIR) print('Saving logs to folder: ' + colored('"{}"'.format(MODEL_LOG_DIR), 'blue')) # Create model model = Model() print('Model created.') # Enable Multi-GPU training print("Let's use", torch.cuda.device_count(), "GPUs!") if torch.cuda.device_count() > 1: print( 'Multiple GPUs being used, can\'t save model graph to Tensorboard') # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model = nn.DataParallel(model) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs prefix = 'densenet_' + str(batch_size) # Load data train_loader, test_loader = getTrainingTestingData(batch_size=batch_size) # Logging # writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(prefix, args.lr, args.epochs, args.bs), flush_secs=30) writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph') # Loss l1_criterion = nn.L1Loss() # Start training... for epoch in range(args.epochs): batch_time = AverageMeter() losses = AverageMeter() N = len(train_loader) # Switch to train mode model.train() end = time.time() for i, sample_batched in enumerate(train_loader): optimizer.zero_grad() # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict output = model(image) # Compute the loss l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss = (1.0 * l_ssim) + (0.1 * l_depth) # Update step losses.update(loss.data.item(), image.size(0)) loss.backward() optimizer.step() # Measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.val * (N - i)))) # Log progress niter = epoch * N + i if i % 5 == 0: # Print to console print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' 'ETA {eta}\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})'.format( epoch, i, N, batch_time=batch_time, loss=losses, eta=eta)) # Log to tensorboard writer.add_scalar('Train/Loss', losses.val, niter) if i % 300 == 0: LogProgress(model, writer, test_loader, niter) # Record epoch's intermediate results LogProgress(model, writer, test_loader, niter) writer.add_scalar('Train/Loss.avg', losses.avg, epoch) # Save the model checkpoint every N epochs saveModelInterval = 2 if (epoch % saveModelInterval) == 0: filename = os.path.join( CHECKPOINT_DIR, 'checkpoint-epoch-{:04d}.pth'.format(epoch)) if torch.cuda.device_count() > 1: model_params = model.module.state_dict( ) # Saving nn.DataParallel model else: model_params = model.state_dict() torch.save( { 'model_state_dict': model_params, 'optimizer_state_dict': optimizer.state_dict(), 'epoch': epoch }, filename)
def main(): # Arguments parser = argparse.ArgumentParser( description= 'High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=config.TRAIN_EPOCH, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float, help='initial learning rate') parser.add_argument('--bs', default=4, type=int, help='batch size') args = parser.parse_args() # Create model model = Model().cuda() print('Model created.') # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs prefix = 'densenet_' + str(batch_size) # Load data train_loader, test_loader = getTrainingTestingData(batch_size=batch_size) # Logging writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format( prefix, args.lr, args.epochs, args.bs), flush_secs=30) # Loss l1_criterion = nn.L1Loss() niter = 0 if config.load_model and config.model_name != '': model.load_state_dict(torch.load('Checkpoints\\%s' % config.model_name)) niter = int(config.model_name.split('_')[-3]) print('load success -> %s' % config.model_name) # Start training... for epoch in range(0, args.epochs): if niter % config.save_interval == 0: if not os.path.exists('Checkpoints\\%s' % config.save_name): os.makedirs('Checkpoints\\%s' % config.save_name) save_name = '%s\\net_params_%s_%s.pkl' % ( config.save_name, niter, datetime.datetime.now().strftime("%Y-%m-%d_%H-%M")) torch.save(model.state_dict(), 'Checkpoints\\%s' % save_name) print('save success -> %s' % save_name) # batch_time = AverageMeter() loss_l1 = AverageMeter() loss_ssim = AverageMeter() loss_edge = AverageMeter() N = len(train_loader) # Switch to train mode model.train() # end = time.time() for i, sample_batched in enumerate(train_loader): optimizer.zero_grad() # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict output = model(image) edge = sobel_(output) pred_edge = sobel_(depth_n) # Compute the loss l_sobel = nn.L1Loss()(edge, pred_edge) l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) # print('nan', torch.sum(torch.isnan(image)).item()) # print('inf', torch.sum(torch.isinf(image)).item()) loss = (1.0 * l_ssim) + (0.1 * l_depth) + l_sobel if torch.isnan(l_sobel) or torch.isnan(l_depth) or torch.isnan( l_ssim) or torch.isinf(l_sobel) or torch.isinf( l_depth) or torch.isinf(l_ssim): print(0) print('nan', torch.sum(torch.isnan(image)).item()) print('inf', torch.sum(torch.isinf(image)).item()) print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim)) print(torch.isinf(l_sobel), torch.isinf(l_depth), torch.isinf(l_ssim)) return # Update step loss_l1.update(l_depth.data.item(), image.size(0)) loss_ssim.update(l_ssim.data.item(), image.size(0)) loss_edge.update(l_sobel.data.item(), image.size(0)) loss.backward() optimizer.step() # Measure elapsed time # batch_time.update(time.time() - end) # end = time.time() # eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i)))) # Log progress niter += 1 if i % 5 == 0: # Print to console # print('Epoch: [{0}][{1}/{2}]\t' # 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' # 'ETA {eta}\t' # 'Loss {loss.val:.4f} ({loss.avg:.4f})' # .format(epoch, i, N, batch_time=batch_time, loss=losses, eta=eta)) # Log to tensorboard writer.add_scalar('Train/L1', loss_l1.val, niter) writer.add_scalar('Train/SSIM', loss_ssim.val, niter) writer.add_scalar('Train/EDGE', loss_edge.val, niter) if i % 300 == 0: LogProgress(model, writer, test_loader, niter) # Record epoch's intermediate results LogProgress(model, writer, test_loader, niter)
def main(): # Arguments parser = argparse.ArgumentParser( description= 'High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=20, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate') parser.add_argument('--bs', default=4, type=int, help='batch size') parser.add_argument('--o', dest='optimizer', help='training optimizer', default="adam", type=str) parser.add_argument('--lr_decay_step', dest='lr_decay_step', help='step to do learning rate decay, unit is epoch', default=5, type=int) parser.add_argument('--lr_decay_gamma', dest='lr_decay_gamma', help='learning rate decay ratio', default=0.1, type=float) args = parser.parse_args() # Create model # model = Model().cuda() model = UNet(n_channels=3, n_classes=1, bilinear=True).cuda() writer_train = SummaryWriter('runs/train_0') # model.load_state_dict(torch.load(r'models\06-05-2020_15-04-20-n15000-e10-bs4-lr0.0001\weights.epoch9_model.pth')) print('Model created.') logger = Logger('./logs') # hyperparams lr = args.lr bs = args.bs lr_decay_step = args.lr_decay_step lr_decay_gamma = args.lr_decay_gamma DOUBLE_BIAS = True WEIGHT_DECAY = 0.0001 # params params = [] for key, value in dict(model.named_parameters()).items(): if value.requires_grad: if 'bias' in key: params += [{'params': [value], 'lr': lr * (DOUBLE_BIAS + 1), \ 'weight_decay': 4e-5 and WEIGHT_DECAY or 0}] else: params += [{'params': [value], 'lr': lr, 'weight_decay': 4e-5}] # optimizer if args.optimizer == "adam": optimizer = torch.optim.Adam(params, lr=lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=4e-5) # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs prefix = 'unet_' + str(batch_size) # Load data train_loader, test_loader = getTrainingTestingData(batch_size=batch_size, trainCount=68730) # # Logging # writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(prefix, args.lr, args.epochs, args.bs), flush_secs=30) # Loss # l1_criterion = nn.L1Loss() # rmse = RMSE() l1_criterion = nn.L1Loss() grad_criterion = GradLoss() normal_criterion = NormalLoss() # eval_metric = RMSE_log() now = datetime.datetime.now() # current date and time runID = now.strftime("%m-%d-%Y_%H-%M-%S") + '-n' + str( len(train_loader)) + '-e' + str( args.epochs) + '-bs' + str(batch_size) + '-lr' + str(args.lr) outputPath = './models/' runPath = outputPath + runID pathlib.Path(runPath).mkdir(parents=True, exist_ok=True) # constants # grad_factor = 10. # normal_factor = 10. # Start training... for epoch in range(args.epochs): batch_time = AverageMeter() losses = AverageMeter() N = len(train_loader) # Switch to train mode model.train() end = time.time() for i, sample_batched in enumerate(train_loader): optimizer.zero_grad() # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict output = model(image) # Compute the loss l_depth = l1_criterion(output, depth_n) # l_ssim = torch.clamp((1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=5.0)) * 0.5, 0, 1) # sbasak01 # sbasak01 loss modification grad_real, grad_fake = imgrad_yx(depth_n), imgrad_yx(output) grad_loss = grad_criterion( grad_fake, grad_real) # * grad_factor # * (epoch > 3) normal_loss = normal_criterion( grad_fake, grad_real) # * normal_factor # * (epoch > 7) # loss = (1.0 * l_ssim) + (0.1 * l_depth) loss = (l_ssim) + (l_depth) + (grad_loss) + (normal_loss) # Update step losses.update(loss.data.item(), image.size(0)) loss.backward() optimizer.step() # ADD TRAINING LOSS TO SummaryWriter writer_train.add_scalar('LOSS', loss.data.item(), i) # Measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.val * (N - i)))) # Log progress niter = epoch * N + i if i % 5 == 0: # Print to console print( 'Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' 'ETA {eta}\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 'L1_Loss: {l1_loss:.4f} SSIM_Loss: {ssim_loss:.4f} grad_loss: {gradloss:.4f} normal_loss: {' 'normalloss:.4f} '.format( epoch, i, N, batch_time=batch_time, loss=losses, eta=eta, l1_loss=l_depth, ssim_loss=l_ssim, gradloss=grad_loss, # sbasak01 loss modification normalloss=normal_loss # sbasak01 loss modification )) # # 1. Log scalar values (scalar summary) # info = {'loss': loss.item()} # # for tag, value in info.items(): # logger.scalar_summary(tag, value, i + 1) # # # 2. Log values and gradients of the parameters (histogram summary) # for tag, value in model.named_parameters(): # tag = tag.replace('.', '/') # logger.histo_summary(tag, value.data.cpu().numpy(), i + 1) # logger.histo_summary(tag + '/grad', value.grad.data.cpu().numpy(), i + 1) if i % 100 == 0: sys.stdout.write("Iterations: %d \r" % (i)) sys.stdout.flush() # save Model intermediate path = runPath + '/weights.epoch{0}_model.pth'.format(epoch) torch.save(model.cpu().state_dict(), path) # saving model model.cuda()
def main(): # Arguments parser = argparse.ArgumentParser( description= 'High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=50, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate') parser.add_argument('--bs', default=6, type=int, help='batch size') parser.add_argument('--o', dest='optimizer', help='training optimizer', default="adam", type=str) parser.add_argument('--lr_decay_step', dest='lr_decay_step', help='step to do learning rate decay, unit is epoch', default=5, type=int) parser.add_argument('--lr_decay_gamma', dest='lr_decay_gamma', help='learning rate decay ratio', default=0.1, type=float) args = parser.parse_args() # Create model model = Model().cuda() macs, params = get_model_complexity_info(model, (3, 480, 640), as_strings=True, print_per_layer_stat=True, verbose=True) #print(summary(model, (3, 480, 640))) print('{:<30} {:<8}'.format('Computational complexity: ', macs)) print('{:<30} {:<8}'.format('Number of parameters: ', params)) # hyperparams lr = args.lr bs = args.bs lr_decay_step = args.lr_decay_step lr_decay_gamma = args.lr_decay_gamma DOUBLE_BIAS = True WEIGHT_DECAY = 0.0001 # params params = [] for key, value in dict(model.named_parameters()).items(): if value.requires_grad: if 'bias' in key: params += [{'params': [value], 'lr': lr * (DOUBLE_BIAS + 1), \ 'weight_decay': 4e-5 and WEIGHT_DECAY or 0}] else: params += [{'params': [value], 'lr': lr, 'weight_decay': 4e-5}] # optimizer if args.optimizer == "adam": optimizer = torch.optim.Adam(params, lr=lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=4e-5) # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs prefix = 'mobilenet_' + str(batch_size) # Load data train_loader, test_loader = getTrainingTestingData(batch_size=batch_size, trainCount=159569) # Logging writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format( prefix, args.lr, args.epochs, args.bs), flush_secs=30) # Loss # l1_criterion = nn.L1Loss() # rmse = RMSE() l1_criterion = nn.L1Loss() grad_criterion = GradLoss() normal_criterion = NormalLoss() # eval_metric = RMSE_log() now = datetime.datetime.now() # current date and time runID = now.strftime("%m-%d-%Y_%H-%M-%S") + '-n' + str( len(train_loader)) + '-e' + str( args.epochs) + '-bs' + str(batch_size) + '-lr' + str(args.lr) outputPath = './models_fk_mobile_sum_weights/' runPath = outputPath + runID pathlib.Path(runPath).mkdir(parents=True, exist_ok=True) # constants grad_factor = 10. normal_factor = 10. # Start training... for epoch in range(args.epochs): batch_time = AverageMeter() losses = AverageMeter() N = len(train_loader) # Switch to train mode model.train() end = time.time() for i, sample_batched in enumerate(train_loader): optimizer.zero_grad() # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict output = model(image) # Compute the loss l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=5.0)) * 0.5, 0, 1) # sbasak01 # FK01 loss modification grad_real, grad_fake = imgrad_yx(depth_n), imgrad_yx(output) grad_loss = grad_criterion( grad_fake, grad_real) * grad_factor # * (epoch > 3) normal_loss = normal_criterion( grad_fake, grad_real) * normal_factor # * (epoch > 7) loss = (0.28 * l_ssim) + (0.22 * l_depth) + (0.30 * grad_loss) + ( 0.20 * normal_loss) # Fk01 loss modification # Update step losses.update(loss.data.item(), image.size(0)) loss.backward() optimizer.step() # Measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.val * (N - i)))) # Log progress niter = epoch * N + i if i % 5 == 0: # Print to console print( 'Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' 'ETA {eta}\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 'L1_Loss: {l1_loss:.4f} SSIM_Loss: {ssim_loss:.4f} grad_loss: {gradloss:.4f} normal_loss: {' 'normalloss:.4f} '.format( epoch, i, N, batch_time=batch_time, loss=losses, eta=eta, l1_loss=l_depth, ssim_loss=l_ssim, gradloss=grad_loss, # Fk1 loss modification normalloss=normal_loss # Fk01 loss modification )) # Log to tensorboard writer.add_scalar('Train/Loss', losses.val, niter) writer.add_scalar('Train/Loss.avg', losses.avg, epoch) # save Model intermediate path = runPath + '/weights.epoch{0}_model.pth'.format(epoch) torch.save(model.cpu().state_dict(), path) # saving model model.cuda()
def trainAndVal(loader, model, l1_criterion, optimizer=None): batch_time = AverageMeter() losses = AverageMeter() rsmes = AverageMeter() if (optimizer): # switch to train mode model.train() print('Train', flush=True) else: # switch to evaluate mode model.eval() print('Val', flush=True) N = len(loader) end = time.time() start = end if (optimizer is None): predictions = [] testSetDepths = [] device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # process epoch for i, sample_batched in enumerate(loader): # Prepare sample and target image = sample_batched['image'].to(device) depth = sample_batched['depth'].to(device) # Normalize depth depth_n = DepthNorm(depth) # Predict output = model(image) # Compute the loss l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss = (1.0 * l_ssim) + (0.1 * l_depth) # measure accuracy and record loss losses.update(loss.data, image.size(0)) rmse = (depth_n.data.cpu() - output.data.cpu())**2 rmse = np.sqrt(rmse.mean()) rsmes.update(rmse, image.size(0)) if (optimizer): # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.avg * (N - i)))) total = str( datetime.timedelta(seconds=int((time.time() - start) + batch_time.avg * (N - i)))) minDepth = 10 maxDepth = 1000 if (optimizer is None): predictions.append(output.squeeze().data.cpu()) testSetDepths.append(depth_n.squeeze().data.cpu()) if i % 5 == 0: p = 100 * i / N bar = "[%-10s] %d%%" % ('=' * int(p * 10 / 100) + '.' * (10 - int(p * 10 / 100)), p) print('[{0}/{1}] {2} - ' 'Batch Time: {batch_time.val:.2f} ({batch_time.avg:.2f}) ' 'ETA: {eta}/{total} ' 'Loss: {loss.val:.3f} ({loss.avg:.3f}) ' 'RSME: {rsme.val:.3f} ({rsme.avg:.3f})'.format( i, N, bar, batch_time=batch_time, eta=eta, total=total, loss=losses, rsme=rsmes), flush=True) break if (optimizer is None): predictions = np.vstack(predictions) testSetDepths = np.vstack(testSetDepths) e = compute_errors(predictions, testSetDepths) print("{:>10}, {:>10}, {:>10}, {:>10}, {:>10}, {:>10}".format( 'a1', 'a2', 'a3', 'rel', 'rms', 'log_10')) print("{:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}, {:10.4f}". format(e[0], e[1], e[2], e[3], e[4], e[5])) return losses.avg
def main(): # Arguments parser = argparse.ArgumentParser( description='Finetuning for depth estimation') parser.add_argument('--epochs', default=60, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate') parser.add_argument('--model_type', default='densedepth', help='type of the depth estimation network') parser.add_argument('--layers', default=161, type=int, help='number of layers of encoder') parser.add_argument('--bs', default=8, type=int, help='batch size') args = parser.parse_args() # Create model if args.model_type == 'densedepth': model = Model(layers=args.layers) else: model = ResNet(layers=args.layers) model = model.cuda() model = nn.DataParallel(model) print('Model created.') # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs # Load data train_loader, test_loader = getTrainingTestingData(batch_size=batch_size) # Loss l1_criterion = nn.L1Loss() masked_criterion = MaskedL1Loss() for epoch in range(args.epochs): model.train() for i, sample_batched in enumerate(tqdm.tqdm(train_loader)): optimizer.zero_grad() # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict output = model(image) if args.model_type == 'densedepth': l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss = (1.0 * l_ssim) + (0.1 * l_depth) else: loss = masked_criterion(output, depth_n) # Update step loss.backward() optimizer.step() torch.save(model.module.state_dict(), 'checkpoints/%s_%d.pth' % (args.model_type, args.layers)) print('Epoch:%d Model Saved!' % (epoch + 1))
def main(): # Arguments parser = argparse.ArgumentParser( description= 'High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=config.TRAIN_EPOCH, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float, help='initial learning rate') parser.add_argument('--bs', default=4, type=int, help='batch size') args = parser.parse_args() # Create model model = Model().cuda() print('Model created.') # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs prefix = 'densenet_' + str(batch_size) nyu_list = open('nyu_list.txt').readlines() scannet_list = open('scannet_list.txt').readlines() real_list = nyu_list + scannet_list irs_list = open('irs_list.txt').readlines() mask_list = open('mask_list.txt').readlines() test_list = open('test_list.txt').readlines() train_dataset = TripleDataset(real_list, irs_list, mask_list, None) test_dataset = NyuV2Dataset(test_list) train_loader = DataLoader(dataset=train_dataset, batch_size=config.bs, num_workers=8, shuffle=True) test_loader = DataLoader(dataset=test_dataset, batch_size=4, num_workers=8, shuffle=True) # Load data # train_loader, test_loader = getTrainingTestingData(batch_size=batch_size) # Logging writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format( prefix, args.lr, args.epochs, args.bs), flush_secs=30) # Loss l1_criterion = nn.L1Loss().cuda() niter = 0 if config.load_model and config.model_name != '': checkpoint = torch.load('Checkpoints\\%s' % config.model_name) model.load_state_dict(checkpoint['net']) optimizer.load_state_dict(checkpoint['optimizer']) niter = int(config.model_name.split('_')[-3]) print('load success -> %s' % config.model_name) # Start training... for epoch in range(0, args.epochs): # batch_time = AverageMeter() loss_l1 = AverageMeter() loss_ssim = AverageMeter() loss_edge = AverageMeter() N = len(train_loader) # Switch to train mode model.train() # end = time.time() idx = 0 for syn_input, real_input, syn_label, real_label in (train_loader): if niter % config.save_interval == 0: if not os.path.exists('Checkpoints\\%s' % config.save_name): os.makedirs('Checkpoints\\%s' % config.save_name) save_name = '%s\\net_params_%s_%s.pkl' % ( config.save_name, niter, datetime.datetime.now().strftime("%Y-%m-%d_%H-%M")) state = { 'net': model.state_dict(), 'optimizer': optimizer.state_dict() } torch.save(state, 'Checkpoints\\%s' % save_name) print('save success -> %s' % save_name) optimizer.zero_grad() _input = torch.cat([syn_input, real_input], dim=0) _label = torch.cat([syn_label, real_label], dim=0) # Prepare sample and target image = torch.autograd.Variable(_input.float().cuda()) depth = torch.autograd.Variable( _label.float().cuda(non_blocking=True)) # Normalize depth # depth_n = DepthNorm(depth) # Predict output = model(image) # pred_edge = sobel_(output) # edge = sobel_(depth) # pred_edge = torch.where(pred_edge>0.2, pred_edge, torch.full_like(pred_edge, 0)) # edge = torch.where(edge > 0.2, edge, torch.full_like(pred_edge, 0)) # Compute the loss # l_sobel = nn.L1Loss()(edge, pred_edge) l_depth = l1_criterion(output, depth) l_ssim = torch.clamp( (1 - ssim(output, depth, val_range=10.0)) * 0.5, 0, 1) # print('nan', torch.sum(torch.isnan(image)).item()) # print('inf', torch.sum(torch.isinf(image)).item()) loss = (1.0 * l_ssim) + (0.1 * l_depth) # if torch.isnan(l_sobel) or torch.isnan(l_depth) or torch.isnan(l_ssim) or torch.isinf(l_sobel) or torch.isinf(l_depth) or torch.isinf(l_ssim): # print(0) # print('nan', torch.sum(torch.isnan(image)).item()) # print('inf', torch.sum(torch.isinf(image)).item()) # print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim)) # print(torch.isinf(l_sobel), torch.isinf(l_depth), torch.isinf(l_ssim)) # return # Update step loss_l1.update(l_depth.data.item(), image.size(0)) loss_ssim.update(l_ssim.data.item(), image.size(0)) # loss_edge.update(l_sobel.data.item(), image.size(0)) loss.backward() optimizer.step() # Measure elapsed time # batch_time.update(time.time() - end) # end = time.time() # eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i)))) # Log progress niter += 1 if idx % 50 == 0: # Print to console # print('Epoch: [{0}][{1}/{2}]\t' # 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' # 'ETA {eta}\t' # 'Loss {loss.val:.4f} ({loss.avg:.4f})' # .format(epoch, i, N, batch_time=batch_time, loss=losses, eta=eta)) # Log to tensorboard writer.add_scalar('Train/L1', loss_l1.val, niter) writer.add_scalar('Train/SSIM', loss_ssim.val, niter) # writer.add_scalar('Train/EDGE', loss_edge.val, niter) if idx % 1000 == 0: LogProgress(model, writer, test_loader, niter) idx += 1 # Record epoch's intermediate results LogProgress(model, writer, test_loader, niter)
def main(): # Arguments parser = argparse.ArgumentParser( description= 'High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=20, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate') parser.add_argument('--data-folder', type=str, help='location of dataset folder') parser.add_argument('--pretrained-weights', type=str, default=None, help='location of pretrained weights') parser.add_argument('--label-file', type=str, help='name of label txt file') parser.add_argument('--bs', default=4, type=int, help='batch size') parser.add_argument('--num-instances', default=None, type=int, help='number of instances of same image in the batch') parser.add_argument('--lamb', default=0., type=float, help='lambda: multiplier for added additional loss') parser.add_argument('--evaluate', type=str, default=None, help='path of model file for testing') args = parser.parse_args() print(args) # Create model model = PTResModel(pretrained_weights=args.pretrained_weights) evaluating = (args.evaluate != None) if (evaluating): print('Evaluating ', args.evaluate) checkpoint = torch.load(args.evaluate) model.load_state_dict(checkpoint) model = model.cuda() print('Model created.') # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs # Load data data_test, nyu2_test = loadToMem(args.data_folder, txtfile='/nyu2_test_updated.csv') transformed_testing = depthDatasetMemory( data_test, nyu2_test, transform=getNoTransform(is_test=True)) test_loader = DataLoader(transformed_testing, 1, shuffle=False) if (evaluating): evaluate_model(model, test_loader, args) return data, nyu2_train = loadToMem(args.data_folder, txtfile=args.label_file) transformed_training = depthDatasetMemory( data, nyu2_train, transform=getDefaultTrainTransform()) if (args.num_instances): train_loader = DataLoader(transformed_training, batch_size, sampler=RandomIdentitySampler( transformed_training, num_instances=args.num_instances), num_workers=4, drop_last=True) else: train_loader = DataLoader(transformed_training, batch_size, shuffle=True, num_workers=4, drop_last=True) # Logging writer = SummaryWriter('logs', flush_secs=30) # Loss l1_criterion = nn.L1Loss() l1_criterion = l1_criterion.cuda() ps = AllPositivePairSelector(balance=False) criterion2 = OnlineContrastiveLoss(1., ps) criterion2 = criterion2.cuda() # Start training... for epoch in range(args.epochs): batch_time = AverageMeter() losses = AverageMeter() N = len(train_loader) # Switch to train mode model.train() transformed_training.reset_seeds() end = time.time() for i, sample_batched in enumerate(train_loader): optimizer.zero_grad() # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) image_id = torch.autograd.Variable( sample_batched['image_id'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict output, feats = model(image) # Compute the loss l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss = (1.0 * l_ssim) + (0.1 * l_depth) if (args.lamb > 0): l_mse = criterion2(feats, image_id, neg_loss=False) loss += args.lamb * l_mse # Update step losses.update(loss.data.item(), image.size(0)) loss.backward() optimizer.step() # Measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.val * (N - i)))) # Log progress niter = epoch * N + i if i % 5 == 0: # Print to console print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' 'ETA {eta}\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})'.format( epoch, i, N, batch_time=batch_time, loss=losses, eta=eta)) # Log to tensorboard writer.add_scalar('Train/Loss', losses.val, niter) if i % 300 == 0: LogProgress(model, writer, test_loader, niter) # Record epoch's intermediate results LogProgress(model, writer, test_loader, niter) writer.add_scalar('Train/Loss.avg', losses.avg, epoch) torch.save(model.state_dict(), "models_" + str(epoch) + '.pth.tar') evaluate_model(model, test_loader, args)
def main(): # Arguments parser = argparse.ArgumentParser( description= 'High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('--epochs', default=20, type=int, help='number of total epochs to run') parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, help='initial learning rate') parser.add_argument('--bs', default=4, type=int, help='batch size') args = parser.parse_args() # Create model model = Model().cuda() print('Model created.') # Training parameters optimizer = torch.optim.Adam(model.parameters(), args.lr) batch_size = args.bs prefix = 'densenet_' + str(batch_size) # Load data train_loader, test_loader = getTrainingTestingData(batch_size=batch_size) # Logging writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format( prefix, args.lr, args.epochs, args.bs), flush_secs=30) # Loss l1_criterion = nn.L1Loss() # Start training... for epoch in range(args.epochs): batch_time = AverageMeter() losses = AverageMeter() N = len(train_loader) # Switch to train mode model.train() end = time.time() for i, sample_batched in enumerate(train_loader): optimizer.zero_grad() # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict output = model(image) # Compute the loss l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss = (1.0 * l_ssim) + (0.1 * l_depth) # Update step losses.update(loss.data.item(), image.size(0)) loss.backward() optimizer.step() # Measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.val * (N - i)))) # Log progress niter = epoch * N + i if i % 5 == 0: # Print to console print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' 'ETA {eta}\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})'.format( epoch, i, N, batch_time=batch_time, loss=losses, eta=eta)) # Log to tensorboard writer.add_scalar('Train/Loss', losses.val, niter) if i % 300 == 0: LogProgress(model, writer, test_loader, niter) # Record epoch's intermediate results LogProgress(model, writer, test_loader, niter) writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
# Forward pass of the mini-batch # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # Predict outputs = model(image) # Compute the loss l_depth = l1_criterion(outputs, depth_n) l_ssim = torch.clamp( (1 - ssim(outputs, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) loss = (1.0 * l_ssim) + (0.1 * l_depth) running_loss += loss.item() # Save output images, one at a time, to results inputs_tensor = image.detach().cpu() output_tensor = outputs.detach().cpu() label_tensor = depth_n.detach().cpu() depth_metric = depth * (config.train.max_depth / 1000.0) outputs_tmp = DepthNorm(outputs) outputs_metric = outputs_tmp * (config.train.max_depth / 1000.0) metrics = compute_errors(depth_metric, outputs_metric) # print(metrics)
def main(): # Arguments parser = argparse.ArgumentParser(description='High Quality Monocular Depth Estimation via Transfer Learning') parser.add_argument('-c', '--configFile', required=True, help='Path to config yaml file', metavar='path/to/config') args = parser.parse_args() CONFIG_FILE_PATH = args.configFile with open(CONFIG_FILE_PATH) as fd: config_yaml = oyaml.load(fd) # Returns an ordered dict. Used for printing config = AttrDict(config_yaml) print(colored('Config being used for training:\n{}\n\n'.format(oyaml.dump(config_yaml)), 'green')) # Create a new directory to save logs runs = sorted(glob.glob(os.path.join(config.train.logsDir, 'exp-*'))) prev_run_id = int(runs[-1].split('-')[-1]) if runs else 0 MODEL_LOG_DIR = os.path.join(config.train.logsDir, 'exp-{:03d}'.format(prev_run_id + 1)) CHECKPOINT_DIR = os.path.join(MODEL_LOG_DIR, 'checkpoints') os.makedirs(CHECKPOINT_DIR) print('Saving logs to folder: ' + colored('"{}"'.format(MODEL_LOG_DIR), 'blue')) # Save a copy of config file in the logs shutil.copy(CONFIG_FILE_PATH, os.path.join(MODEL_LOG_DIR, 'config.yaml')) # Create a tensorboard object and Write config to tensorboard writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph') string_out = io.StringIO() oyaml.dump(config_yaml, string_out, default_flow_style=False) config_str = string_out.getvalue().split('\n') string = '' for line in config_str: string = string + ' ' + line + '\n\r' writer.add_text('Config', string, global_step=None) # Create model model = Model() print('Model created.') # to continue training from a checkpoint if config.train.continueTraining: print('Transfer Learning enabled. Model State to be loaded from a prev checkpoint...') if not os.path.isfile(config.train.pathPrevCheckpoint): raise ValueError('Invalid path to the given weights file for transfer learning.\ The file {} does not exist'.format(config.train.pathPrevCheckpoint)) CHECKPOINT = torch.load(config.train.pathPrevCheckpoint, map_location='cpu') if 'model_state_dict' in CHECKPOINT: # Newer weights file with various dicts print(colored('Continuing training from checkpoint...Loaded data from checkpoint:', 'green')) print('Config Used to train Checkpoint:\n', oyaml.dump(CHECKPOINT['config']), '\n') print('From Checkpoint: Last Epoch Loss:', CHECKPOINT['epoch_loss'], '\n\n') model.load_state_dict(CHECKPOINT['model_state_dict']) elif 'state_dict' in CHECKPOINT: # reading original authors checkpoints if config.train.model != 'rednet': # original author deeplab checkpoint CHECKPOINT['state_dict'].pop('decoder.last_conv.8.weight') CHECKPOINT['state_dict'].pop('decoder.last_conv.8.bias') else: # rednet checkpoint # print(CHECKPOINT['state_dict'].keys()) CHECKPOINT['state_dict'].pop('final_deconv.weight') CHECKPOINT['state_dict'].pop('final_deconv.bias') CHECKPOINT['state_dict'].pop('out5_conv.weight') CHECKPOINT['state_dict'].pop('out5_conv.bias') CHECKPOINT['state_dict'].pop('out4_conv.weight') CHECKPOINT['state_dict'].pop('out4_conv.bias') CHECKPOINT['state_dict'].pop('out3_conv.weight') CHECKPOINT['state_dict'].pop('out3_conv.bias') CHECKPOINT['state_dict'].pop('out2_conv.weight') CHECKPOINT['state_dict'].pop('out2_conv.bias') model.load_state_dict(CHECKPOINT['state_dict'], strict=False) else: # Old checkpoint containing only model's state_dict() model.load_state_dict(CHECKPOINT) # Enable Multi-GPU training print("Let's use", torch.cuda.device_count(), "GPUs!") if torch.cuda.device_count() > 1: print('Multiple GPUs being used, can\'t save model graph to Tensorboard') # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model = nn.DataParallel(model) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) # Training parameters optimizer = torch.optim.Adam( model.parameters(), config.train.optimAdam.learningRate ) batch_size = config.train.batchSize prefix = 'densenet_' + str(batch_size) # Load data train_loader_list = [] test_loader_list = [] for dataset in config.train.datasetsTrain: train_data = getTrainingTestingData('rgb', 'train', dataset.images, dataset.labels) train_loader_list.append(train_data) for dataset in config.train.datasetsVal: print(dataset.images) test_data = getTrainingTestingData('rgb', 'eval', dataset.images, dataset.labels) test_loader_list.append(test_data) train_loader = DataLoader(torch.utils.data.ConcatDataset(train_loader_list), batch_size, num_workers=config.train.numWorkers, shuffle=True, drop_last=True, pin_memory=True) test_loader = DataLoader(torch.utils.data.ConcatDataset(test_loader_list), batch_size, num_workers=config.train.numWorkers, shuffle=False, drop_last=True, pin_memory=True) print(len(torch.utils.data.ConcatDataset(train_loader_list))) print(len(train_loader)) print(len(test_loader)) # Create a tensorboard object and Write config to tensorboard writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph') # Loss l1_criterion = nn.L1Loss() total_iter_num = 0 # Start training... for epoch in range(config.train.numEpochs): batch_time = AverageMeter() losses = AverageMeter() N = len(train_loader) # Log the current Epoch Number writer.add_scalar('data/Epoch Number', epoch, total_iter_num) # Switch to train mode model.train() end = time.time() running_loss = 0.0 for i, sample_batched in enumerate(train_loader): optimizer.zero_grad() total_iter_num += 1 # Prepare sample and target image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable(sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm( depth ) # Predict output = model(image) # Compute the loss l_depth = l1_criterion(output, depth_n) l_ssim = torch.clamp((1 - ssim(output, depth_n, val_range = 1000.0 / 10.0)) * 0.5, 0, 1) loss = (1.0 * l_ssim) + (0.1 * l_depth) # Update step losses.update(loss.data.item(), image.size(0)) loss.backward() optimizer.step() # statistics running_loss += loss.item() # Measure elapsed time batch_time.update(time.time() - end) end = time.time() eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i)))) # Log progress niter = epoch*N+i if i % 5 == 0: # Print to console print('Epoch: [{0}][{1}/{2}]\t' 'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t' 'ETA {eta}\t' 'Loss {loss.val:.4f} ({loss.avg:.4f})' .format(epoch, i, N, batch_time=batch_time, loss=losses, eta=eta)) # Log to tensorboard writer.add_scalar('Train/Loss', losses.val, niter) if i % 50 == 0: LogProgress(model, writer, test_loader, niter) # Log Epoch Loss epoch_loss = running_loss / (len(train_loader)) writer.add_scalar('data/Train Epoch Loss', epoch_loss, total_iter_num) print('\nTrain Epoch Loss: {:.4f}'.format(epoch_loss)) metrics = compute_errors(depth_n, output) print(metrics) for keys, values in metrics.items(): print(str(keys) + ':' + str(values)) # Record epoch's intermediate results LogProgress(model, writer, test_loader, niter) writer.add_scalar('Train/Loss.avg', losses.avg, epoch) # Save the model checkpoint every N epochs if (epoch % config.train.saveModelInterval) == 0: filename = os.path.join(CHECKPOINT_DIR, 'checkpoint-epoch-{:04d}.pth'.format(epoch)) if torch.cuda.device_count() > 1: model_params = model.module.state_dict() # Saving nn.DataParallel model else: model_params = model.state_dict() torch.save( { 'model_state_dict': model_params, 'optimizer_state_dict': optimizer.state_dict(), 'epoch': epoch, 'total_iter_num': total_iter_num, 'epoch_loss': epoch_loss, 'config': config_yaml }, filename)
def LogProgress(model, writer, test_loader, global_step): model.eval() sequential = test_loader sample_batched = next(iter(sequential)) image = torch.autograd.Variable(sample_batched['image'].cuda()) depth = torch.autograd.Variable( sample_batched['depth'].cuda(non_blocking=True)) # Normalize depth depth_n = DepthNorm(depth) # print() # if epoch == 0: # writer.add_image('Train.1.Image', vutils.make_grid(image.data, nrow=6, normalize=True), epoch) # if epoch == 0: writer.add_image( 'Train.2.Depth', colorize(vutils.make_grid(depth.data, nrow=6, normalize=False)), global_step) output = model(image) plt.imshow(output.detach().cpu().numpy().squeeze(), cmap='plasma') plt.show() # Predict # output = model(image) edge = sobel_(output) pred_edge = sobel_(depth_n) # Compute the loss l_sobel = nn.L1Loss()(edge, pred_edge) l_depth = nn.L1Loss()(output, depth_n) l_ssim = torch.clamp( (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1) if torch.isnan(l_sobel): print(1) print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim)) writer.add_scalar('Test/L1', l_depth.item(), global_step) writer.add_scalar('Test/SSIM', l_ssim.item(), global_step) writer.add_scalar('Test/EDGE', l_sobel.item(), global_step) # normal = Norm(output) # normal = (normal + 1.0) / 2. # normal = normal[3, :, :, :] # normal = normal.detach().cpu().numpy().astype(np.uint8) # normal = np.transpose(normal, (1, 2, 0)) # print() # plt.imshow(normal) # plt.show() output = DepthNorm(output) # writer.add_image('Train.3.Normal', vutils.make_grid(normal.data, nrow=6, normalize=False), epoch) writer.add_image( 'Test.2.Ours', colorize(vutils.make_grid(output.data, nrow=6, normalize=False)), global_step) writer.add_image( 'Test.3.Diff', colorize( vutils.make_grid(torch.abs(output - depth).data, nrow=6, normalize=False)), global_step) del image del depth del output del edge del pred_edge