def validate(val_loader, disp_net, pose_exp_net, epoch, logger, output_writers=[]): global args batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) log_outputs = len(output_writers) > 0 w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight poses = np.zeros( ((len(val_loader) - 1) * args.batch_size * (args.sequence_length - 1), 6)) # switch to evaluate mode disp_net.eval() pose_exp_net.eval() end = time.time() for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(val_loader): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_imgs_var = [ Variable(img.cuda(), volatile=True) for img in ref_imgs ] intrinsics_var = Variable(intrinsics.cuda(), volatile=True) intrinsics_inv_var = Variable(intrinsics_inv.cuda(), volatile=True) # compute output disp = disp_net(tgt_img_var) depth = 1 / disp explainability_mask, pose = pose_exp_net(tgt_img_var, ref_imgs_var) loss_1 = photometric_reconstruction_loss(tgt_img_var, ref_imgs_var, intrinsics_var, intrinsics_inv_var, depth, explainability_mask, pose) loss_1 = loss_1.data[0] if w2 > 0: loss_2 = explainability_loss(explainability_mask).data[0] else: loss_2 = 0 loss_3 = smooth_loss(disp).data[0] if log_outputs and i % 100 == 0 and i / 100 < len( output_writers): # log first output of every 100 batch index = int(i // 100) if epoch == 0: for j, ref in enumerate(ref_imgs): output_writers[index].add_image('val Input {}'.format(j), tensor2array(tgt_img[0]), 0) output_writers[index].add_image('val Input {}'.format(j), tensor2array(ref[0]), 1) output_writers[index].add_image( 'val Dispnet Output Normalized', tensor2array(disp.data[0].cpu(), max_value=None, colormap='bone'), epoch) output_writers[index].add_image( 'val Depth Output', tensor2array(1. / disp.data[0].cpu(), max_value=10), epoch) # log warped images along with explainability mask for j, ref in enumerate(ref_imgs_var): ref_warped = inverse_warp(ref[:1], depth[:1, 0], pose[:1, j], intrinsics_var[:1], intrinsics_inv_var[:1])[0] output_writers[index].add_image( 'val Warped Outputs {}'.format(j), tensor2array(ref_warped.data.cpu()), epoch) output_writers[index].add_image( 'val Diff Outputs {}'.format(j), tensor2array( 0.5 * (tgt_img_var[0] - ref_warped).abs().data.cpu()), epoch) if explainability_mask is not None: output_writers[index].add_image( 'val Exp mask Outputs {}'.format(j), tensor2array(explainability_mask[0, j].data.cpu(), max_value=1, colormap='bone'), epoch) if log_outputs and i < len(val_loader) - 1: step = args.batch_size * (args.sequence_length - 1) poses[i * step:(i + 1) * step] = pose.data.cpu().view(-1, 6).numpy() loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 losses.update([loss, loss_1, loss_2]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() logger.valid_bar.update(i) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Loss {}'.format( batch_time, losses)) if log_outputs: output_writers[0].add_histogram('val poses_tx', poses[:, 0], epoch) output_writers[0].add_histogram('val poses_ty', poses[:, 1], epoch) output_writers[0].add_histogram('val poses_tz', poses[:, 2], epoch) output_writers[0].add_histogram('val poses_rx', poses[:, 3], epoch) output_writers[0].add_histogram('val poses_ry', poses[:, 4], epoch) output_writers[0].add_histogram('val poses_rz', poses[:, 5], epoch) return losses.avg
def train(args, train_loader, disp_net, pose_exp_net, optimizer, epoch_size, logger, train_writer): global n_iter, device batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight # switch to train mode disp_net.train() pose_exp_net.train() end = time.time() logger.train_bar.update(0) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): # measure data loading time data_time.update(time.time() - end) tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) # compute output disparities = disp_net(tgt_img) depth = [1 / disp for disp in disparities] explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1 = photometric_reconstruction_loss(tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) if w2 > 0: loss_2 = explainability_loss(explainability_mask) else: loss_2 = 0 loss_3 = smooth_loss(depth) loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 if i > 0 and n_iter % args.print_freq == 0: train_writer.add_scalar('photometric_error', loss_1.item(), n_iter) if w2 > 0: train_writer.add_scalar('explanability_loss', loss_2.item(), n_iter) train_writer.add_scalar('disparity_smoothness_loss', loss_3.item(), n_iter) train_writer.add_scalar('total_loss', loss.item(), n_iter) if args.training_output_freq > 0 and n_iter % args.training_output_freq == 0: train_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) with torch.no_grad(): for k, scaled_depth in enumerate(depth): train_writer.add_image( 'train Dispnet Output Normalized {}'.format(k), tensor2array(disparities[k][0], max_value=None, colormap='magma'), n_iter) train_writer.add_image( 'train Depth Output Normalized {}'.format(k), tensor2array(1 / disparities[k][0], max_value=None), n_iter) b, _, h, w = scaled_depth.size() downscale = tgt_img.size(2) / h tgt_img_scaled = F.interpolate(tgt_img, (h, w), mode='area') ref_imgs_scaled = [ F.interpolate(ref_img, (h, w), mode='area') for ref_img in ref_imgs ] intrinsics_scaled = torch.cat( (intrinsics[:, 0:2] / downscale, intrinsics[:, 2:]), dim=1) # log warped images along with explainability mask for j, ref in enumerate(ref_imgs_scaled): ref_warped = inverse_warp( ref, scaled_depth[:, 0], pose[:, j], intrinsics_scaled, rotation_mode=args.rotation_mode, padding_mode=args.padding_mode)[0] train_writer.add_image( 'train Warped Outputs {} {}'.format(k, j), tensor2array(ref_warped), n_iter) train_writer.add_image( 'train Diff Outputs {} {}'.format(k, j), tensor2array( 0.5 * (tgt_img_scaled[0] - ref_warped).abs()), n_iter) if explainability_mask[k] is not None: train_writer.add_image( 'train Exp mask Outputs {} {}'.format(k, j), tensor2array(explainability_mask[k][0, j], max_value=1, colormap='bone'), n_iter) # record loss and EPE losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([ loss.item(), loss_1.item(), loss_2.item() if w2 > 0 else 0, loss_3.item() ]) logger.train_bar.update(i + 1) if i % args.print_freq == 0: logger.train_writer.write('Train: Time {} Data {} Loss {}'.format( batch_time, data_time, losses)) if i >= epoch_size - 1: break n_iter += 1 return losses.avg[0]
def validate_without_gt(args, val_loader, depth_net, pose_net, epoch, logger, output_writers=[], **env): global device batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) w1, w2, w3 = args.photo_loss_weight, args.smooth_loss_weight, args.ssim if args.log_output: poses_values = np.zeros(((len(val_loader) - 1) * args.test_batch_size * (args.sequence_length - 1), 6)) disp_values = np.zeros( ((len(val_loader) - 1) * args.test_batch_size * 3)) # switch to evaluate mode depth_net.eval() pose_net.eval() upsample_depth_net = models.UpSampleNet(depth_net, args.network_input_size) end = time.time() logger.valid_bar.update(0) for i, sample in enumerate(val_loader): log_output = i < len(output_writers) imgs = torch.stack(sample['imgs'], dim=1).to(device) intrinsics = sample['intrinsics'].to(device) intrinsics_inv = sample['intrinsics_inv'].to(device) if epoch == 1 and log_output: for j, img in enumerate(sample['imgs']): output_writers[i].add_image('val Input', tensor2array(img[0]), j) batch_size, seq = imgs.size()[:2] if args.network_input_size is not None: h, w = args.network_input_size downsample_imgs = F.interpolate(imgs, (3, h, w), mode='area') poses = pose_net(downsample_imgs) # [B, seq, 6] else: poses = pose_net(imgs) pose_matrices = pose_vec2mat(poses, args.rotation_mode) # [B, seq, 3, 4] mid_index = (args.sequence_length - 1) // 2 tgt_imgs = imgs[:, mid_index] # [B, 3, H, W] tgt_poses = pose_matrices[:, mid_index] # [B, 3, 4] compensated_poses = compensate_pose( pose_matrices, tgt_poses) # [B, seq, 3, 4] tgt_poses are now neutral pose ref_indices = list(range(args.sequence_length)) ref_indices.remove(mid_index) loss_1 = 0 loss_2 = 0 for ref_index in ref_indices: prior_imgs = imgs[:, ref_index] prior_poses = compensated_poses[:, ref_index] # [B, 3, 4] prior_imgs_compensated = inverse_rotate(prior_imgs, prior_poses[:, :, :3], intrinsics, intrinsics_inv) input_pair = torch.cat([prior_imgs_compensated, tgt_imgs], dim=1) # [B, 6, W, H] predicted_magnitude = prior_poses[:, :, -1:].norm( p=2, dim=1, keepdim=True).unsqueeze(1) # [B, 1, 1, 1] scale_factor = args.nominal_displacement / predicted_magnitude normalized_translation = compensated_poses[:, :, :, -1:] * scale_factor # [B, seq, 3, 1] new_pose_matrices = torch.cat( [compensated_poses[:, :, :, :-1], normalized_translation], dim=-1) depth = upsample_depth_net(input_pair) disparity = 1 / depth total_indices = torch.arange(seq).long().unsqueeze(0).expand( batch_size, seq).to(device) tgt_id = total_indices[:, mid_index] ref_indices = total_indices[ total_indices != tgt_id.unsqueeze(1)].view( batch_size, seq - 1) photo_loss, diff_maps, warped_imgs = photometric_reconstruction_loss( imgs, tgt_id, ref_indices, depth, new_pose_matrices, intrinsics, intrinsics_inv, args.rotation_mode, ssim_weight=w3) loss_1 += photo_loss if log_output: output_writers[i].add_image( 'val Dispnet Output Normalized {}'.format(ref_index), tensor2array(disparity[0], max_value=None, colormap='bone'), epoch) output_writers[i].add_image( 'val Depth Output {}'.format(ref_index), tensor2array(depth[0].cpu(), max_value=args.max_depth), epoch) for j, (diff, warped) in enumerate(zip(diff_maps, warped_imgs)): output_writers[i].add_image( 'val Warped Outputs {} {}'.format(j, ref_index), tensor2array(warped[0]), epoch) output_writers[i].add_image( 'val Diff Outputs {} {}'.format(j, ref_index), tensor2array(diff[0].abs() - 1), epoch) loss_2 += texture_aware_smooth_loss( disparity, tgt_imgs if args.texture_loss else None) if args.log_output and i < len(val_loader) - 1: step = args.test_batch_size * (args.sequence_length - 1) poses_values[i * step:(i + 1) * step] = poses[:, :-1].cpu().view( -1, 6).numpy() step = args.test_batch_size * 3 disp_unraveled = disparity.cpu().view(args.test_batch_size, -1) disp_values[i * step:(i + 1) * step] = torch.cat([ disp_unraveled.min(-1)[0], disp_unraveled.median(-1)[0], disp_unraveled.max(-1)[0] ]).numpy() loss = w1 * loss_1 + w2 * loss_2 losses.update([loss.item(), loss_1.item(), loss_2.item()]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() logger.valid_bar.update(i + 1) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Loss {}'.format( batch_time, losses)) if args.log_output: rot_coeffs = ['rx', 'ry', 'rz'] if args.rotation_mode == 'euler' else [ 'qx', 'qy', 'qz' ] tr_coeffs = ['tx', 'ty', 'tz'] for k, (coeff_name) in enumerate(tr_coeffs + rot_coeffs): output_writers[0].add_histogram('val poses_{}'.format(coeff_name), poses_values[:, k], epoch) output_writers[0].add_histogram('disp_values', disp_values, epoch) logger.valid_bar.update(len(val_loader)) return OrderedDict( zip(['Total loss', 'Photo loss', 'Smooth loss'], losses.avg))
def validate_without_gt(args, val_loader, disp_net, pose_exp_net, epoch, logger, output_writers=[]): batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) log_outputs = len(output_writers) > 0 w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight poses = np.zeros( ((len(val_loader) - 1) * args.batch_size * (args.sequence_length - 1), 6)) disp_values = np.zeros(((len(val_loader) - 1) * args.batch_size * 3)) # switch to evaluate mode disp_net.eval() pose_exp_net.eval() end = time.time() logger.valid_bar.update(0) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(val_loader): tgt_img_var = Variable(tgt_img.cuda(), volatile=True) ref_imgs_var = [ Variable(img.cuda(), volatile=True) for img in ref_imgs ] intrinsics_var = Variable(intrinsics.cuda(), volatile=True) intrinsics_inv_var = Variable(intrinsics_inv.cuda(), volatile=True) # compute output disp = disp_net(tgt_img_var) depth = 1 / disp explainability_mask, pose = pose_exp_net(tgt_img_var, ref_imgs_var) loss_1 = photometric_reconstruction_loss(tgt_img_var, ref_imgs_var, intrinsics_var, intrinsics_inv_var, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) loss_1 = loss_1.data[0] if w2 > 0: loss_2 = explainability_loss(explainability_mask).data[0] else: loss_2 = 0 loss_3 = smooth_loss(disp).data[0] if log_outputs and i % 100 == 0 and i / 100 < len( output_writers): # log first output of every 100 batch index = int(i // 100) if epoch == 0: for j, ref in enumerate(ref_imgs): output_writers[index].add_image('val Input {}'.format(j), tensor2array(tgt_img[0]), 0) output_writers[index].add_image('val Input {}'.format(j), tensor2array(ref[0]), 1) output_writers[index].add_image( 'val Dispnet Output Normalized', tensor2array(disp.data[0].cpu(), max_value=None, colormap='bone'), epoch) output_writers[index].add_image( 'val Depth Output', tensor2array(1. / disp.data[0].cpu(), max_value=10), epoch) # log warped images along with explainability mask for j, ref in enumerate(ref_imgs_var): ref_warped = inverse_warp(ref[:1], depth[:1, 0], pose[:1, j], intrinsics_var[:1], intrinsics_inv_var[:1], rotation_mode=args.rotation_mode, padding_mode=args.padding_mode)[0] output_writers[index].add_image( 'val Warped Outputs {}'.format(j), tensor2array(ref_warped.data.cpu()), epoch) output_writers[index].add_image( 'val Diff Outputs {}'.format(j), tensor2array( 0.5 * (tgt_img_var[0] - ref_warped).abs().data.cpu()), epoch) if explainability_mask is not None: output_writers[index].add_image( 'val Exp mask Outputs {}'.format(j), tensor2array(explainability_mask[0, j].data.cpu(), max_value=1, colormap='bone'), epoch) if log_outputs and i < len(val_loader) - 1: step = args.batch_size * (args.sequence_length - 1) poses[i * step:(i + 1) * step] = pose.data.cpu().view(-1, 6).numpy() step = args.batch_size * 3 disp_unraveled = disp.data.cpu().view(args.batch_size, -1) disp_values[i * step:(i + 1) * step] = torch.cat([ disp_unraveled.min(-1)[0], disp_unraveled.median(-1)[0], disp_unraveled.max(-1)[0] ]).numpy() loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 losses.update([loss, loss_1, loss_2]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() logger.valid_bar.update(i + 1) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Loss {}'.format( batch_time, losses)) if log_outputs: prefix = 'valid poses' coeffs_names = ['tx', 'ty', 'tz'] if args.rotation_mode == 'euler': coeffs_names.extend(['rx', 'ry', 'rz']) elif args.rotation_mode == 'quat': coeffs_names.extend(['qx', 'qy', 'qz']) for i in range(poses.shape[1]): output_writers.add_histogram( '{} {}'.format(prefix, coeffs_names[i]), poses[:, i], epoch) output_writers[0].add_histogram('disp_values', disp_values, epoch) logger.valid_bar.update(len(val_loader)) return losses.avg, ['Total loss', 'Photo loss', 'Exp loss']
def train(train_loader, disp_net, pose_exp_net, optimizer, epoch_size, logger, train_writer): global args, n_iter batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight # switch to train mode disp_net.train() pose_exp_net.train() end = time.time() for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): # measure data loading time data_time.update(time.time() - end) tgt_img_var = Variable(tgt_img.cuda()) ref_imgs_var = [Variable(img.cuda()) for img in ref_imgs] intrinsics_var = Variable(intrinsics.cuda()) intrinsics_inv_var = Variable(intrinsics_inv.cuda()) # compute output disparities = disp_net(tgt_img_var) depth = [1 / disp for disp in disparities] explainability_mask, pose = pose_exp_net(tgt_img_var, ref_imgs_var) loss_1 = photometric_reconstruction_loss(tgt_img_var, ref_imgs_var, intrinsics_var, intrinsics_inv_var, depth, explainability_mask, pose) loss_2 = explainability_loss(explainability_mask) loss_3 = smooth_loss(disparities) loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 train_writer.add_scalar('photometric_error', loss_1.data[0], n_iter) train_writer.add_scalar('explanability_loss', loss_2.data[0], n_iter) train_writer.add_scalar('disparity_smoothness_loss', loss_3.data[0], n_iter) train_writer.add_scalar('total_loss', loss.data[0], n_iter) if n_iter % 200 == 0 and args.log_output: train_writer.add_image('train Input', tensor2array(ref_imgs[0][0]), n_iter - 1) train_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) train_writer.add_image('train Input', tensor2array(ref_imgs[1][0]), n_iter + 1) for k, scaled_depth in enumerate(depth): train_writer.add_image( 'train Dispnet Output {}'.format(k), tensor2array(disparities[k].data[0].cpu(), max_value=10, colormap='bone'), n_iter) train_writer.add_image( 'train Depth Output Normalized {}'.format(k), tensor2array(1 / disparities[k].data[0].cpu(), max_value=None), n_iter) b, _, h, w = scaled_depth.size() downscale = tgt_img_var.size(2) / h tgt_img_scaled = nn.functional.adaptive_avg_pool2d( tgt_img_var, (h, w)) ref_imgs_scaled = [ nn.functional.adaptive_avg_pool2d(ref_img, (h, w)) for ref_img in ref_imgs_var ] intrinsics_scaled = torch.cat( (intrinsics_var[:, 0:2] / downscale, intrinsics_var[:, 2:]), dim=1) intrinsics_scaled_inv = torch.cat( (intrinsics_inv_var[:, :, 0:2] * downscale, intrinsics_inv_var[:, :, 2:]), dim=2) # log warped images along with explainability mask for j, ref in enumerate(ref_imgs_scaled): ref_warped = inverse_warp(ref, scaled_depth[:, 0], pose[:, j], intrinsics_scaled, intrinsics_scaled_inv)[0] train_writer.add_image( 'train Warped Outputs {} {}'.format(k, j), tensor2array(ref_warped.data.cpu(), max_value=1), n_iter) train_writer.add_image( 'train Diff Outputs {} {}'.format(k, j), tensor2array( 0.5 * (tgt_img_scaled[0] - ref_warped).abs().data.cpu()), n_iter) train_writer.add_image( 'train Exp mask Outputs {} {}'.format(k, j), tensor2array(explainability_mask[k][0, j].data.cpu(), max_value=1, colormap='bone'), n_iter) # record loss and EPE losses.update(loss.data[0], args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow( [loss.data[0], loss_1.data[0], loss_2.data[0], loss_3.data[0]]) logger.train_bar.update(i) if i % args.print_freq == 0: logger.train_writer.write( 'Train: Time {batch_time.val:.3f} ({batch_time.avg:.3f}) ' 'Data {data_time.val:.3f} ({data_time.avg:.3f}) ' 'Loss {loss.val:.4f} ({loss.avg:.4f}) '.format( batch_time=batch_time, data_time=data_time, loss=losses)) if i >= epoch_size - 1: break n_iter += 1 return losses.avg
def train_one_epoch(args, train_loader, depth_net, pose_net, optimizer, epoch, n_iter, logger, training_writer, **env): global device logger.reset_train_bar() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3 = args.photo_loss_weight, args.smooth_loss_weight, args.ssim e1, e2 = args.training_milestones # switch to train mode depth_net.train() pose_net.train() upsample_depth_net = models.UpSampleNet(depth_net, args.network_input_size) end = time.time() logger.train_bar.update(0) for i, sample in enumerate(train_loader): log_losses = i > 0 and n_iter % args.print_freq == 0 log_output = args.training_output_freq > 0 and n_iter % args.training_output_freq == 0 # measure data loading time data_time.update(time.time() - end) imgs = torch.stack(sample['imgs'], dim=1).to(device) intrinsics = sample['intrinsics'].to(device) intrinsics_inv = sample['intrinsics_inv'].to(device) batch_size, seq = imgs.size()[:2] if args.network_input_size is not None: h, w = args.network_input_size downsample_imgs = F.interpolate(imgs, (3, h, w), mode='area') poses = pose_net(downsample_imgs) # [B, seq, 6] else: poses = pose_net(imgs) pose_matrices = pose_vec2mat(poses, args.rotation_mode) # [B, seq, 3, 4] total_indices = torch.arange(seq).long().to(device).unsqueeze( 0).expand(batch_size, seq) batch_range = torch.arange(batch_size).long().to(device) ''' for each element of the batch select a random picture in the sequence to which we will compute the depth, all poses are then converted so that pose of this very picture is exactly identity. At first this image is always in the middle of the sequence''' if epoch > e2: tgt_id = torch.floor(torch.rand(batch_size) * seq).long().to(device) else: tgt_id = torch.zeros(batch_size).long().to( device) + args.sequence_length // 2 ''' Select what other picture we are going to feed DepthNet, it must not be the same as tgt_id. At first, it's always first picture of the sequence, it is randomly chosen when first training milestone is reached ''' ref_indices = total_indices[total_indices != tgt_id.unsqueeze(1)].view( batch_size, seq - 1) if epoch > e1: prior_id = torch.floor(torch.rand(batch_size) * (seq - 1)).long().to(device) else: prior_id = torch.zeros(batch_size).long().to(device) prior_id = ref_indices[batch_range, prior_id] tgt_imgs = imgs[batch_range, tgt_id] # [B, 3, H, W] tgt_poses = pose_matrices[batch_range, tgt_id] # [B, 3, 4] prior_imgs = imgs[batch_range, prior_id] compensated_poses = compensate_pose( pose_matrices, tgt_poses) # [B, seq, 3, 4] tgt_poses are now neutral pose prior_poses = compensated_poses[batch_range, prior_id] # [B, 3, 4] if args.supervise_pose: from_GT = invert_mat(sample['pose']).to(device) compensated_GT_poses = compensate_pose( from_GT, from_GT[batch_range, tgt_id]) prior_GT_poses = compensated_GT_poses[batch_range, prior_id] prior_imgs_compensated = inverse_rotate(prior_imgs, prior_GT_poses[:, :, :-1], intrinsics, intrinsics_inv) else: prior_imgs_compensated = inverse_rotate(prior_imgs, prior_poses[:, :, :-1], intrinsics, intrinsics_inv) input_pair = torch.cat([prior_imgs_compensated, tgt_imgs], dim=1) # [B, 6, W, H] depth = upsample_depth_net(input_pair) # depth = [sample['depth'].to(device).unsqueeze(1) * 3 / abs(tgt_id[0] - prior_id[0])] # depth.append(torch.nn.functional.interpolate(depth[0], scale_factor=2)) disparities = [1 / d for d in depth] predicted_magnitude = prior_poses[:, :, -1:].norm(p=2, dim=1, keepdim=True).unsqueeze(1) scale_factor = args.nominal_displacement / (predicted_magnitude + 1e-5) normalized_translation = compensated_poses[:, :, :, -1:] * scale_factor # [B, seq_length-1, 3] new_pose_matrices = torch.cat( [compensated_poses[:, :, :, :-1], normalized_translation], dim=-1) biggest_scale = depth[0].size(-1) loss_1 = 0 for k, scaled_depth in enumerate(depth): size_ratio = scaled_depth.size(-1) / biggest_scale loss, diff_maps, warped_imgs = photometric_reconstruction_loss( imgs, tgt_id, ref_indices, scaled_depth, new_pose_matrices, intrinsics, intrinsics_inv, args.rotation_mode, ssim_weight=w3) loss_1 += loss * size_ratio if log_output: training_writer.add_image( 'train Dispnet Output Normalized scale {}'.format(k), tensor2array(disparities[k][0], max_value=None, colormap='bone'), n_iter) training_writer.add_image( 'train Depth Output scale {}'.format(k), tensor2array(scaled_depth[0], max_value=args.max_depth), n_iter) for j, (diff, warped) in enumerate(zip(diff_maps, warped_imgs)): training_writer.add_image( 'train Warped Outputs {} {}'.format(k, j), tensor2array(warped[0]), n_iter) training_writer.add_image( 'train Diff Outputs {} {}'.format(k, j), tensor2array(diff.abs()[0] - 1), n_iter) loss_2 = texture_aware_smooth_loss( depth, tgt_imgs if args.texture_loss else None) loss = w1 * loss_1 + w2 * loss_2 if args.supervise_pose: loss += (from_GT[:, :, :, :3] - pose_matrices[:, :, :, :3]).abs().mean() if log_losses: training_writer.add_scalar('photometric_error', loss_1.item(), n_iter) training_writer.add_scalar('disparity_smoothness_loss', loss_2.item(), n_iter) training_writer.add_scalar('total_loss', loss.item(), n_iter) if log_output: nominal_translation_magnitude = poses[:, -2, :3].norm(p=2, dim=-1) # last pose is always identity and penultimate translation magnitude is always 1, so you don't need to log them for j in range(args.sequence_length - 2): trans_mag = poses[:, j, :3].norm(p=2, dim=-1) training_writer.add_histogram( 'tr {}'.format(j), (trans_mag / nominal_translation_magnitude).detach().cpu().numpy(), n_iter) for j in range(args.sequence_length - 1): # TODO log a better value : this is magnitude of vector (yaw, pitch, roll) which is not a physical value rot_mag = poses[:, j, 3:].norm(p=2, dim=-1) training_writer.add_histogram('rot {}'.format(j), rot_mag.detach().cpu().numpy(), n_iter) training_writer.add_image('train Input', tensor2array(tgt_imgs[0]), n_iter) # record loss for average meter losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([loss.item(), loss_1.item(), loss_2.item()]) logger.train_bar.update(i + 1) if i % args.print_freq == 0: logger.train_writer.write('Train: Time {} Data {} Loss {}'.format( batch_time, data_time, losses)) if i >= args.epoch_size - 1: break n_iter += 1 return losses.avg[0], n_iter
def validate_without_gt(args, val_loader, disp_net, pose_exp_net, epoch, tb_writer, sample_nb_to_log=3): global device batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) log_outputs = sample_nb_to_log > 0 w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight poses = np.zeros( ((len(val_loader) - 1) * args.batch_size * (args.sequence_length - 1), 6)) disp_values = np.zeros(((len(val_loader) - 1) * args.batch_size * 3)) # switch to evaluate mode disp_net.eval() pose_exp_net.eval() end = time.time() validate_pbar = tqdm( total=len(val_loader), bar_format='{desc} {percentage:3.0f}%|{bar}| {postfix}') validate_pbar.set_description( 'valid: Loss *.**** *.*****.****(*.**** *.**** *.****)') validate_pbar.set_postfix_str('<Time *.***(*.***)>') for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(val_loader): tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) intrinsics_inv = intrinsics_inv.to(device) # compute output disp = disp_net(tgt_img) depth = 1 / disp explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1, warped, diff = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) loss_1 = loss_1.item() if w2 > 0: loss_2 = explainability_loss(explainability_mask).item() else: loss_2 = 0 loss_3 = smooth_loss(depth).item() if log_outputs and i < sample_nb_to_log - 1: # log first output of first batches if epoch == 0: for j, ref in enumerate(ref_imgs): tb_writer.add_image('val Input {}/{}'.format(j, i), tensor2array(tgt_img[0]), 0) tb_writer.add_image('val Input {}/{}'.format(j, i), tensor2array(ref[0]), 1) log_output_tensorboard(tb_writer, 'val', i, '', epoch, 1. / disp, disp, warped, diff, explainability_mask) if log_outputs and i < len(val_loader) - 1: step = args.batch_size * (args.sequence_length - 1) poses[i * step:(i + 1) * step] = pose.cpu().view(-1, 6).numpy() step = args.batch_size * 3 disp_unraveled = disp.cpu().view(args.batch_size, -1) disp_values[i * step:(i + 1) * step] = torch.cat([ disp_unraveled.min(-1)[0], disp_unraveled.median(-1)[0], disp_unraveled.max(-1)[0] ]).numpy() loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 losses.update([loss, loss_1, loss_2]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() validate_pbar.clear() validate_pbar.update(1) validate_pbar.set_description('valid: Loss {}'.format(losses)) validate_pbar.set_postfix_str('<Time {}>'.format(batch_time)) validate_pbar.close() if log_outputs: prefix = 'valid poses' coeffs_names = ['tx', 'ty', 'tz'] if args.rotation_mode == 'euler': coeffs_names.extend(['rx', 'ry', 'rz']) elif args.rotation_mode == 'quat': coeffs_names.extend(['qx', 'qy', 'qz']) for i in range(poses.shape[1]): tb_writer.add_histogram('{} {}'.format(prefix, coeffs_names[i]), poses[:, i], epoch) tb_writer.add_histogram('disp_values', disp_values, epoch) time.sleep(0.2) else: time.sleep(1) return losses.avg, ['Total loss', 'Photo loss', 'Exp loss']
def validate_without_gt(args, val_loader, disp_net, pose_exp_net, epoch, logger, tb_writer, sample_nb_to_log=3): global device batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) log_outputs = sample_nb_to_log > 0 w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight poses = np.zeros( ((len(val_loader) - 1) * args.batch_size * (args.sequence_length - 1), 6)) disp_values = np.zeros(((len(val_loader) - 1) * args.batch_size * 3)) # switch to evaluate mode disp_net.eval() pose_exp_net.eval() end = time.time() logger.valid_bar.update(0) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(val_loader): tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) intrinsics_inv = intrinsics_inv.to(device) # compute output disp = disp_net(tgt_img) depth = 1 / disp explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1, warped, diff = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) loss_1 = loss_1.item() if w2 > 0: loss_2 = explainability_loss(explainability_mask).item() else: loss_2 = 0 loss_3 = smooth_loss(depth).item() if log_outputs and i < sample_nb_to_log - 1: # log first output of first batches if epoch == 0: for j, ref in enumerate(ref_imgs): tb_writer.add_image('val Input {}/{}'.format(j, i), tensor2array(tgt_img[0]), 0) tb_writer.add_image('val Input {}/{}'.format(j, i), tensor2array(ref[0]), 1) log_output_tensorboard(tb_writer, 'val', i, '', epoch, 1. / disp, disp, warped[0], diff[0], explainability_mask) if log_outputs and i < len(val_loader) - 1: step = args.batch_size * (args.sequence_length - 1) poses[i * step:(i + 1) * step] = pose.cpu().view(-1, 6).numpy() step = args.batch_size * 3 disp_unraveled = disp.cpu().view(args.batch_size, -1) disp_values[i * step:(i + 1) * step] = torch.cat([ disp_unraveled.min(-1)[0], disp_unraveled.median(-1)[0], disp_unraveled.max(-1)[0] ]).numpy() loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 if args.with_photocon_loss: batch_size = pose.size()[0] homo_row = torch.tensor([[0, 0, 0, 1]], dtype=torch.float).to(device) homo_row = homo_row.unsqueeze(0).expand(batch_size, -1, -1) T21 = pose_vec2mat(pose[:, 0]) T21 = torch.cat((T21, homo_row), 1) T12 = torch.inverse(T21) T23 = pose_vec2mat(pose[:, 1]) T23 = torch.cat((T23, homo_row), 1) T13 = torch.matmul(T23, T12) #[B,4,4] # print("----",T13.size()) # target = 1(ref_imgs[0]) and ref = 3(ref_imgs[1]) ref_img_warped, valid_points = inverse_warp_posemat( ref_imgs[1], depth[:, 0], T13, intrinsics, args.rotation_mode, args.padding_mode) diff = (ref_imgs[0] - ref_img_warped) * valid_points.unsqueeze(1).float() loss_4 = diff.abs().mean() loss += loss_4 losses.update([loss, loss_1, loss_2]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() logger.valid_bar.update(i + 1) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Loss {}'.format( batch_time, losses)) if log_outputs: prefix = 'valid poses' coeffs_names = ['tx', 'ty', 'tz'] if args.rotation_mode == 'euler': coeffs_names.extend(['rx', 'ry', 'rz']) elif args.rotation_mode == 'quat': coeffs_names.extend(['qx', 'qy', 'qz']) for i in range(poses.shape[1]): tb_writer.add_histogram('{} {}'.format(prefix, coeffs_names[i]), poses[:, i], epoch) tb_writer.add_histogram('disp_values', disp_values, epoch) logger.valid_bar.update(len(val_loader)) return losses.avg, [ 'Validation Total loss', 'Validation Photo loss', 'Validation Exp loss' ]
def validate_without_gt(args, val_loader, depth_net, pose_net, epoch, logger, tb_writer, sample_nb_to_log, **env): global device batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) w1, w2, w3 = args.photo_loss_weight, args.smooth_loss_weight, args.ssim if args.log_output: poses_values = np.zeros(((len(val_loader) - 1) * args.test_batch_size * (args.sequence_length - 1), 6)) disp_values = np.zeros( ((len(val_loader) - 1) * args.test_batch_size * 3)) # switch to evaluate mode depth_net.eval() pose_net.eval() end = time.time() logger.valid_bar.update(0) for i, sample in enumerate(val_loader): log_output = i < sample_nb_to_log imgs = torch.stack(sample['imgs'], dim=1).to(device) intrinsics = sample['intrinsics'].to(device) if epoch == 1 and log_output: for j, img in enumerate(sample['imgs']): tb_writer.add_image('val Input/{}'.format(i), tensor2array(img[0]), j) batch_size, seq = imgs.size()[:2] poses = pose_net(imgs) pose_matrices = pose_vec2mat(poses, args.rotation_mode) # [B, seq, 3, 4] mid_index = (args.sequence_length - 1) // 2 tgt_imgs = imgs[:, mid_index] # [B, 3, H, W] tgt_poses = pose_matrices[:, mid_index] # [B, 3, 4] compensated_poses = compensate_pose( pose_matrices, tgt_poses) # [B, seq, 3, 4] tgt_poses are now neutral pose ref_ids = list(range(args.sequence_length)) ref_ids.remove(mid_index) loss_1 = 0 loss_2 = 0 for ref_index in ref_ids: prior_imgs = imgs[:, ref_index] prior_poses = compensated_poses[:, ref_index] # [B, 3, 4] prior_imgs_compensated = inverse_rotate(prior_imgs, prior_poses[:, :, :3], intrinsics) input_pair = torch.cat([prior_imgs_compensated, tgt_imgs], dim=1) # [B, 6, W, H] predicted_magnitude = prior_poses[:, :, -1:].norm( p=2, dim=1, keepdim=True).unsqueeze(1) # [B, 1, 1, 1] scale_factor = args.nominal_displacement / predicted_magnitude normalized_translation = compensated_poses[:, :, :, -1:] * scale_factor # [B, seq, 3, 1] new_pose_matrices = torch.cat( [compensated_poses[:, :, :, :-1], normalized_translation], dim=-1) depth = depth_net(input_pair) disparity = 1 / depth tgt_id = torch.full((batch_size, ), ref_index, dtype=torch.int64, device=device) ref_ids_tensor = torch.tensor(ref_ids, dtype=torch.int64, device=device).expand( batch_size, -1) photo_loss, *to_log = photometric_reconstruction_loss( imgs, tgt_id, ref_ids_tensor, depth, new_pose_matrices, intrinsics, args.rotation_mode, ssim_weight=w3, upsample=args.upscale) loss_1 += photo_loss if log_output: log_output_tensorboard(tb_writer, "train", i, ref_index, epoch, depth[0], disparity[0], *to_log) loss_2 += grad_diffusion_loss(disparity, tgt_imgs, args.kappa) if args.log_output and i < len(val_loader) - 1: step = args.test_batch_size * (args.sequence_length - 1) poses_values[i * step:(i + 1) * step] = poses[:, :-1].cpu().view( -1, 6).numpy() step = args.test_batch_size * 3 disp_unraveled = disparity.cpu().view(args.test_batch_size, -1) disp_values[i * step:(i + 1) * step] = torch.cat([ disp_unraveled.min(-1)[0], disp_unraveled.median(-1)[0], disp_unraveled.max(-1)[0] ]).numpy() loss = w1 * loss_1 + w2 * loss_2 losses.update([loss.item(), loss_1.item(), loss_2.item()]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() logger.valid_bar.update(i + 1) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Loss {}'.format( batch_time, losses)) if args.log_output: rot_coeffs = ['rx', 'ry', 'rz'] if args.rotation_mode == 'euler' else [ 'qx', 'qy', 'qz' ] tr_coeffs = ['tx', 'ty', 'tz'] for k, (coeff_name) in enumerate(tr_coeffs + rot_coeffs): tb_writer.add_histogram('val poses_{}'.format(coeff_name), poses_values[:, k], epoch) tb_writer.add_histogram('disp_values', disp_values, epoch) logger.valid_bar.update(len(val_loader)) return OrderedDict( zip(['Total loss', 'Photo loss', 'Smooth loss'], losses.avg))
def train(args, train_loader, disp_net, pose_exp_net, optimizer, epoch_size, tb_writer): global n_iter, device batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight # switch to train mode disp_net.train() pose_exp_net.train() end = time.time() train_pbar = tqdm(total=min(len(train_loader), args.epoch_size), bar_format='{desc} {percentage:3.0f}%|{bar}| {postfix}') train_pbar.set_description('Train: Total Loss=#.####(#.####)') train_pbar.set_postfix_str('<TIME: op=#.###(#.###) DataFlow=#.###(#.###)>') for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): log_losses = i > 0 and n_iter % args.print_freq == 0 log_output = args.training_output_freq > 0 and n_iter % args.training_output_freq == 0 # measure DataFlow loading time data_time.update(time.time() - end) tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) # compute output disparities = disp_net(tgt_img) depth = [1 / disp for disp in disparities] explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1, warped, diff = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) if w2 > 0: loss_2 = explainability_loss(explainability_mask) else: loss_2 = 0 loss_3 = smooth_loss(depth) loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 if loss < 0.0005: abc = 0 if log_losses: tb_writer.add_scalar('photometric_error', loss_1.item(), n_iter) if w2 > 0: tb_writer.add_scalar('explanability_loss', loss_2.item(), n_iter) tb_writer.add_scalar('disparity_smoothness_loss', loss_3.item(), n_iter) tb_writer.add_scalar('total_loss', loss.item(), n_iter) if log_output: tb_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) for k, scaled_maps in enumerate( zip(depth, disparities, warped, diff, explainability_mask)): log_output_tensorboard(tb_writer, "train", 0, k, n_iter, *scaled_maps) # record loss and EPE losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([ loss.item(), loss_1.item(), loss_2.item() if w2 > 0 else 0, loss_3.item() ]) train_pbar.clear() train_pbar.update(1) train_pbar.set_description('Train: Total Loss={}'.format(losses)) train_pbar.set_postfix_str('<TIME: op={} DataFlow={}>'.format( batch_time, data_time)) if i >= epoch_size - 1: break n_iter += 1 train_pbar.close() time.sleep(1) return losses.avg[0]
def train_one_epoch(args, train_loader, depth_net, pose_net, optimizer, epoch, n_iter, logger, tb_writer, **env): global device logger.reset_train_bar() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3 = args.photo_loss_weight, args.smooth_loss_weight, args.ssim e1, e2 = args.training_milestones # switch to train mode depth_net.train() pose_net.train() end = time.time() logger.train_bar.update(0) for i, sample in enumerate(train_loader): log_losses = i > 0 and n_iter % args.print_freq == 0 log_output = args.training_output_freq > 0 and n_iter % args.training_output_freq == 0 # measure data loading time data_time.update(time.time() - end) imgs = torch.stack(sample['imgs'], dim=1).to(device) intrinsics = sample['intrinsics'].to(device) batch_size, seq = imgs.size()[:2] if args.network_input_size is not None: h, w = args.network_input_size downsample_imgs = F.interpolate(imgs, (3, h, w), mode='area') poses = pose_net(downsample_imgs) # [B, seq, 6] else: poses = pose_net(imgs) pose_matrices = pose_vec2mat(poses, args.rotation_mode) # [B, seq, 3, 4] total_indices = torch.arange(seq, dtype=torch.int64, device=device).expand(batch_size, seq) batch_range = torch.arange(batch_size, dtype=torch.int64, device=device) ''' for each element of the batch select a random picture in the sequence to which we will compute the depth, all poses are then converted so that pose of this very picture is exactly identity. At first this image is always in the middle of the sequence''' if epoch > e2: tgt_id = torch.randint(0, seq, (batch_size, ), device=device) else: tgt_id = torch.full_like(batch_range, args.sequence_length // 2) ref_ids = total_indices[total_indices != tgt_id.unsqueeze(1)].view( batch_size, seq - 1) ''' Select what other picture we are going to feed DepthNet, it must not be the same as tgt_id. At first, it's always first picture of the sequence, it is randomly chosen when first training milestone is reached ''' if epoch > e1: probs = torch.ones_like(total_indices, dtype=torch.float32) probs[batch_range, tgt_id] = args.same_ratio prior_id = torch.multinomial(probs, 1)[:, 0] else: prior_id = torch.zeros_like(batch_range) # Treat the case of prior_id == tgt_id and the depth must be max_depth, regardless of apparent movement tgt_imgs = imgs[batch_range, tgt_id] # [B, 3, H, W] tgt_poses = pose_matrices[batch_range, tgt_id] # [B, 3, 4] prior_imgs = imgs[batch_range, prior_id] compensated_poses = compensate_pose( pose_matrices, tgt_poses) # [B, seq, 3, 4] tgt_poses are now neutral pose prior_poses = compensated_poses[batch_range, prior_id] # [B, 3, 4] if args.supervise_pose: from_GT = invert_mat(sample['pose']).to(device) compensated_GT_poses = compensate_pose( from_GT, from_GT[batch_range, tgt_id]) prior_GT_poses = compensated_GT_poses[batch_range, prior_id] prior_imgs_compensated = inverse_rotate(prior_imgs, prior_GT_poses[:, :, :-1], intrinsics) else: prior_imgs_compensated = inverse_rotate(prior_imgs, prior_poses[:, :, :-1], intrinsics) input_pair = torch.cat([prior_imgs_compensated, tgt_imgs], dim=1) # [B, 6, W, H] depth = depth_net(input_pair) # depth = [sample['depth'].to(device).unsqueeze(1) * 3 / abs(tgt_id[0] - prior_id[0])] # depth.append(torch.nn.functional.interpolate(depth[0], scale_factor=2)) disparities = [1 / d for d in depth] predicted_magnitude = prior_poses[:, :, -1:].norm(p=2, dim=1, keepdim=True).unsqueeze(1) scale_factor = args.nominal_displacement / (predicted_magnitude + 1e-5) normalized_translation = compensated_poses[:, :, :, -1:] * scale_factor # [B, seq_length-1, 3] new_pose_matrices = torch.cat( [compensated_poses[:, :, :, :-1], normalized_translation], dim=-1) biggest_scale = depth[0].size(-1) # Construct valid sequence to compute photometric error, # make the rest converge to max_depth because nothing moved vb = batch_range[prior_id != tgt_id] same_range = batch_range[prior_id == tgt_id] # batch of still pairs loss_1 = 0 loss_1_same = 0 for k, scaled_depth in enumerate(depth): size_ratio = scaled_depth.size(-1) / biggest_scale if len(same_range) > 0: # Frames are identical. The corresponding depth must be infinite. Here, we set it to max depth still_depth = scaled_depth[same_range] loss_same = F.smooth_l1_loss(still_depth / args.max_depth, torch.ones_like(still_depth)) else: loss_same = 0 loss_valid, *to_log = photometric_reconstruction_loss( imgs[vb], tgt_id[vb], ref_ids[vb], scaled_depth[vb], new_pose_matrices[vb], intrinsics[vb], args.rotation_mode, ssim_weight=w3, upsample=args.upscale) loss_1 += loss_valid * size_ratio loss_1_same += loss_same * size_ratio if log_output and len(vb) > 0: log_output_tensorboard(tb_writer, "train", 0, k, n_iter, scaled_depth[0], disparities[k][0], *to_log) loss_2 = grad_diffusion_loss(disparities, tgt_imgs, args.kappa) loss = w1 * (loss_1 + loss_1_same) + w2 * loss_2 if args.supervise_pose: loss += (from_GT[:, :, :, :3] - pose_matrices[:, :, :, :3]).abs().mean() if log_losses: tb_writer.add_scalar('photometric_error', loss_1.item(), n_iter) tb_writer.add_scalar('disparity_smoothness_loss', loss_2.item(), n_iter) tb_writer.add_scalar('total_loss', loss.item(), n_iter) if log_output and len(vb) > 0: valid_poses = poses[vb] nominal_translation_magnitude = valid_poses[:, -2, :3].norm(p=2, dim=-1) # Log the translation magnitude relative to translation magnitude between last and penultimate frames # for a perfectly constant displacement magnitude, you should get ratio of 2,3,4 and so forth. # last pose is always identity and penultimate translation magnitude is always 1, so you don't need to log them for j in range(args.sequence_length - 2): trans_mag = valid_poses[:, j, :3].norm(p=2, dim=-1) tb_writer.add_histogram( 'tr {}'.format(j), (trans_mag / nominal_translation_magnitude).detach().cpu().numpy(), n_iter) for j in range(args.sequence_length - 1): # TODO log a better value : this is magnitude of vector (yaw, pitch, roll) which is not a physical value rot_mag = valid_poses[:, j, 3:].norm(p=2, dim=-1) tb_writer.add_histogram('rot {}'.format(j), rot_mag.detach().cpu().numpy(), n_iter) tb_writer.add_image('train Input', tensor2array(tgt_imgs[0]), n_iter) # record loss for average meter losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([loss.item(), loss_1.item(), loss_2.item()]) logger.train_bar.update(i + 1) if i % args.print_freq == 0: logger.train_writer.write('Train: Time {} Data {} Loss {}'.format( batch_time, data_time, losses)) if i >= args.epoch_size - 1: break n_iter += 1 return losses.avg[0], n_iter
def train(odometry_net, depth_net, feat_extractor, train_loader, epoch, optimizer): global device global data_parallel if data_parallel: odometry_net.module.set_fix_method(nfp.FIX_AUTO) else: odometry_net.set_fix_method(nfp.FIX_AUTO) odometry_net.train() depth_net.train() feat_extractor.train() total_loss = 0 img_reconstruction_total = 0 f_reconstruction_total = 0 smooth_total = 0 for batch_idx, (img_R1, img_L2, img_R2, intrinsics, inv_intrinsics, raw_K, T_R2L) in tqdm(enumerate(train_loader), desc='Train epoch %d' % epoch, leave=False, ncols=80): img_R1 = img_R1.type(torch.FloatTensor).to(device) img_R2 = img_R2.type(torch.FloatTensor).to(device) img_L2 = img_L2.type(torch.FloatTensor).to(device) intrinsics = intrinsics.type(torch.FloatTensor).to(device) inv_intrinsics = inv_intrinsics.type(torch.FloatTensor).to(device) raw_K = raw_K.type(torch.FloatTensor).to(device) T_R2L = T_R2L.type(torch.FloatTensor).to(device) img_R = torch.cat((img_R2, img_R1), dim=1) inv_depth_img_R2 = depth_net(img_R2) T_2to1, _ = odometry_net(img_R) T_2to1 = T_2to1.view(T_2to1.size(0), -1) T_R2L = T_R2L.view(T_R2L.size(0), -1) depth = (1 / (inv_depth_img_R2 + 1e-4)).squeeze(1) img_reconstruction_error = photometric_reconstruction_loss( 0.004 * img_R2, 0.004 * img_R1, 0.004 * img_L2, depth, T_2to1, T_R2L, intrinsics, inv_intrinsics) smooth_error = smooth_loss(depth.unsqueeze(1)) imgs = torch.cat((img_L2, img_R2, img_R1), dim=0) feat = feat_extractor(imgs) batch_size = img_R1.size(0) f_L2, f_R2, f_R1 = feat[:batch_size, :, :, :], feat[ batch_size:batch_size * 2, :, :, :], feat[2 * batch_size:, :, :, :] feat_reconstruction_error = photometric_reconstruction_loss( f_R2, f_R1, f_L2, depth, T_2to1, T_R2L, intrinsics, inv_intrinsics) loss = img_reconstruction_error + 0.1 * feat_reconstruction_error + 10 * smooth_error total_loss += loss.item() img_reconstruction_total += img_reconstruction_error.item() f_reconstruction_total += feat_reconstruction_error.item() smooth_total += smooth_error.item() optimizer.zero_grad() loss.backward() optimizer.step() print( "Train epoch {}: loss: {:.9f} img-recon-loss: {:.9f} f-recon-loss: {:.9f} smooth-loss: {:.9f}" .format(epoch, total_loss / len(train_loader), img_reconstruction_total / len(train_loader), f_reconstruction_total / len(train_loader), smooth_total / len(train_loader)))
def train(args, train_loader, disp_net, pose_exp_net, optimizer, epoch_size, logger, train_writer): global n_iter, device batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight # switch to train mode disp_net.train() pose_exp_net.train() end = time.time() logger.train_bar.update(0) #train main cycle for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): #for (i, data) in enumerate(train_loader):#data(list): [tensor(B,3,H,W),list(B),(B,H,W),(b,h,w)] log_losses = i > 0 and n_iter % args.print_freq == 0 log_output = args.training_output_freq > 0 and n_iter % args.training_output_freq == 0 #1 measure data loading time data_time.update(time.time() - end) tgt_img = tgt_img.to(device) #(4,3,128,416) ref_imgs = [img.to(device) for img in ref_imgs] #batch size张图片的前一帧和后一帧 intrinsics = intrinsics.to(device) #(4,3,3) """forward and loss""" #2 compute output disparities = disp_net( tgt_img ) # lenth batch-size list of tensor(4,1,128,416) ,(4,1,64,208),(4,1,32,104),(4,1,16,52)] explainability_mask, pose = pose_exp_net( tgt_img, ref_imgs) #pose tensor(bs,sq-lenth-1,6), relative camera pose depth = [1 / disp for disp in disparities] #depth = fxT/(d) 成反比关系,简单取倒数 #3 loss compute loss_1, warped, diff = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) if w2 > 0: loss_2 = explainability_loss(explainability_mask) else: loss_2 = 0 loss_3 = smooth_loss(depth) loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 #4. 数据记录 tensorboard batch-record data, 而且不用初始化数据名称(自动初始化),直接往里面加 if log_losses: train_writer.add_scalar('photometric_error', loss_1.item(), n_iter) if w2 > 0: train_writer.add_scalar('explanabilityyyyyy_loss', loss_2.item(), n_iter) train_writer.add_scalar('disparity_smoothness_loss', loss_3.item(), n_iter) train_writer.add_scalar('total_loss', loss.item(), n_iter) if log_output: #数据弄到tensorboard可读文件里去, 名字就是events开头(defaulted) train_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) for k, scaled_maps in enumerate( zip(depth, disparities, warped, diff, explainability_mask)): log_output_tensorboard(train_writer, "train", k, n_iter, *scaled_maps) # record loss and EPE losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() #csv record with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([ loss.item(), loss_1.item(), loss_2.item() if w2 > 0 else 0, loss_3.item() ]) logger.train_bar.update(i + 1) if i % args.print_freq == 0: logger.train_writer.write('Train: Time {} Data {} Loss {}'.format( batch_time, data_time, losses)) if i >= epoch_size - 1: break n_iter += 1 return losses.avg[0]
def train(args, train_loader, disp_net, pose_exp_net, optimizer, epoch_size, logger, train_writer): global n_iter, device batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight # switch to train mode disp_net.train() pose_exp_net.train() end = time.time() logger.train_bar.update(0) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): log_losses = i > 0 and n_iter % args.print_freq == 0 log_output = args.training_output_freq > 0 and n_iter % args.training_output_freq == 0 # measure data loading time data_time.update(time.time() - end) tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) # compute output disparities = disp_net(tgt_img) depth = [1 / disp for disp in disparities] explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1, warped, diff = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) if w2 > 0: loss_2 = explainability_loss(explainability_mask) else: loss_2 = 0 loss_3 = inverse_depth_smooth_loss(depth, tgt_img) loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 if log_losses: train_writer.add_scalar('photometric_error', loss_1.item(), n_iter) if w2 > 0: train_writer.add_scalar('explanability_loss', loss_2.item(), n_iter) train_writer.add_scalar('disparity_smoothness_loss', loss_3.item(), n_iter) train_writer.add_scalar('total_loss', loss.item(), n_iter) if log_output: train_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) for k, scaled_maps in enumerate( zip(depth, disparities, warped, diff, explainability_mask)): log_output_tensorboard(train_writer, "train", k, n_iter, *scaled_maps) # record loss and EPE losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([ loss.item(), loss_1.item(), loss_2.item() if w2 > 0 else 0, loss_3.item() ]) logger.train_bar.update(i + 1) if i % args.print_freq == 0: logger.train_writer.write('Train: Time {} Data {} Loss {}'.format( batch_time, data_time, losses)) if i >= epoch_size - 1: break n_iter += 1 return losses.avg[0]
def validate_without_gt(args, val_loader, disp_net, pose_exp_net, epoch, logger, tb_writer, sample_nb_to_log=3): global device mse_l = torch.nn.MSELoss(reduction='mean') batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) log_outputs = sample_nb_to_log > 0 # Output the logs throughout the whole dataset batches_to_log = list( np.linspace(0, len(val_loader), sample_nb_to_log).astype(int)) w1, w2, w3, wf, wp = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight, args.flow_loss_weight, args.prior_loss_weight poses = np.zeros( ((len(val_loader) - 1) * args.batch_size * (args.sequence_length - 1), 6)) disp_values = np.zeros(((len(val_loader) - 1) * args.batch_size * 3)) # switch to evaluate mode disp_net.eval() pose_exp_net.eval() end = time.time() logger.valid_bar.update(0) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv, flow_maps) in enumerate(val_loader): tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) intrinsics_inv = intrinsics_inv.to(device) flow_maps = [flow_map.to(device) for flow_map in flow_maps] # compute output disp = disp_net(tgt_img) depth = 1 / disp explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1, warped, diff, grid = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) loss_1 = loss_1.item() if wf > 0: loss_f = flow_consistency_loss(grid, flow_maps, mse_l) else: loss_f = 0 if wp > 0: loss_p = ground_prior_loss(disp) else: loss_p = 0 if w2 > 0: loss_2 = explainability_loss(explainability_mask).item() else: loss_2 = 0 loss_3 = smooth_loss(depth).item() if log_outputs and i in batches_to_log: # log first output of wanted batches index = batches_to_log.index(i) if epoch == 0: for j, ref in enumerate(ref_imgs): tb_writer.add_image('val Input {}/{}'.format(j, index), tensor2array(tgt_img[0]), 0) tb_writer.add_image('val Input {}/{}'.format(j, index), tensor2array(ref[0]), 1) log_output_tensorboard(tb_writer, 'val', index, '', epoch, 1. / disp, disp, warped[0], diff[0], explainability_mask) if log_outputs and i < len(val_loader) - 1: step = args.batch_size * (args.sequence_length - 1) poses[i * step:(i + 1) * step] = pose.cpu().view(-1, 6).numpy() step = args.batch_size * 3 disp_unraveled = disp.cpu().view(args.batch_size, -1) disp_values[i * step:(i + 1) * step] = torch.cat([ disp_unraveled.min(-1)[0], disp_unraveled.median(-1)[0], disp_unraveled.max(-1)[0] ]).numpy() loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 + wf * loss_f + wp * loss_p losses.update([loss, loss_1, loss_2]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() logger.valid_bar.update(i + 1) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Loss {}'.format( batch_time, losses)) if log_outputs: prefix = 'valid poses' coeffs_names = ['tx', 'ty', 'tz'] if args.rotation_mode == 'euler': coeffs_names.extend(['rx', 'ry', 'rz']) elif args.rotation_mode == 'quat': coeffs_names.extend(['qx', 'qy', 'qz']) for i in range(poses.shape[1]): tb_writer.add_histogram('{} {}'.format(prefix, coeffs_names[i]), poses[:, i], epoch) tb_writer.add_histogram('disp_values', disp_values, epoch) logger.valid_bar.update(len(val_loader)) return losses.avg, [ 'Validation Total loss', 'Validation Photo loss', 'Validation Exp loss' ]
def validate_without_gt(args, val_loader, disp_net, pose_exp_net, epoch, logger, output_writers=[]): global device batch_time = AverageMeter() losses = AverageMeter(i=3, precision=4) log_outputs = len(output_writers) > 0 w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight poses = np.zeros( ((len(val_loader) - 1) * args.batch_size * (args.sequence_length - 1), 6)) disp_values = np.zeros(((len(val_loader) - 1) * args.batch_size * 3)) # switch to evaluate mode disp_net.eval() pose_exp_net.eval() end = time.time() logger.valid_bar.update(0) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(val_loader): tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) intrinsics_inv = intrinsics_inv.to(device) # compute output disp = disp_net(tgt_img) depth = 1 / disp explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1, warped, diff = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) loss_1 = loss_1.item() if w2 > 0: loss_2 = explainability_loss(explainability_mask).item() else: loss_2 = 0 loss_3 = smooth_loss(depth).item() if log_outputs and i < len( output_writers): # log first output of first batches if epoch == 0: for j, ref in enumerate(ref_imgs): output_writers[i].add_image('val Input {}'.format(j), tensor2array(tgt_img[0]), 0) output_writers[i].add_image('val Input {}'.format(j), tensor2array(ref[0]), 1) log_output_tensorboard(output_writers[i], 'val', '', epoch, 1. / disp, disp, warped, diff, explainability_mask) if log_outputs and i < len(val_loader) - 1: step = args.batch_size * (args.sequence_length - 1) poses[i * step:(i + 1) * step] = pose.cpu().view(-1, 6).numpy() step = args.batch_size * 3 disp_unraveled = disp.cpu().view(args.batch_size, -1) disp_values[i * step:(i + 1) * step] = torch.cat([ disp_unraveled.min(-1)[0], disp_unraveled.median(-1)[0], disp_unraveled.max(-1)[0] ]).numpy() loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 losses.update([loss, loss_1, loss_2]) # measure elapsed time batch_time.update(time.time() - end) end = time.time() logger.valid_bar.update(i + 1) if i % args.print_freq == 0: logger.valid_writer.write('valid: Time {} Loss {}'.format( batch_time, losses)) if log_outputs: prefix = 'valid poses' coeffs_names = ['tx', 'ty', 'tz'] if args.rotation_mode == 'euler': coeffs_names.extend(['rx', 'ry', 'rz']) elif args.rotation_mode == 'quat': coeffs_names.extend(['qx', 'qy', 'qz']) for i in range(poses.shape[1]): output_writers[0].add_histogram( '{} {}'.format(prefix, coeffs_names[i]), poses[:, i], epoch) output_writers[0].add_histogram('disp_values', disp_values, epoch) logger.valid_bar.update(len(val_loader)) return losses.avg, ['Total loss', 'Photo loss', 'Exp loss']
def train(args, train_loader, disp_net, pose_exp_net, optimizer, epoch_size, logger, tb_writer): global n_iter, device batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter(precision=4) w1, w2, w3 = args.photo_loss_weight, args.mask_loss_weight, args.smooth_loss_weight # switch to train mode disp_net.train() pose_exp_net.train() end = time.time() logger.train_bar.update(0) for i, (tgt_img, ref_imgs, intrinsics, intrinsics_inv) in enumerate(train_loader): log_losses = i > 0 and n_iter % args.print_freq == 0 log_output = args.training_output_freq > 0 and n_iter % args.training_output_freq == 0 # measure data loading time data_time.update(time.time() - end) tgt_img = tgt_img.to(device) ref_imgs = [img.to(device) for img in ref_imgs] intrinsics = intrinsics.to(device) # compute output disparities = disp_net(tgt_img) depth = [1 / disp for disp in disparities] # print("***",len(depth),depth[0].size()) explainability_mask, pose = pose_exp_net(tgt_img, ref_imgs) loss_1, warped, diff = photometric_reconstruction_loss( tgt_img, ref_imgs, intrinsics, depth, explainability_mask, pose, args.rotation_mode, args.padding_mode) if w2 > 0: loss_2 = explainability_loss(explainability_mask) else: loss_2 = 0 loss_3 = smooth_loss(depth) loss = w1 * loss_1 + w2 * loss_2 + w3 * loss_3 if args.with_photocon_loss: batch_size = pose.size()[0] homo_row = torch.tensor([[0, 0, 0, 1]], dtype=torch.float).to(device) homo_row = homo_row.unsqueeze(0).expand(batch_size, -1, -1) T21 = pose_vec2mat(pose[:, 0]) T21 = torch.cat((T21, homo_row), 1) T12 = torch.inverse(T21) T23 = pose_vec2mat(pose[:, 1]) T23 = torch.cat((T23, homo_row), 1) T13 = torch.matmul(T23, T12) #[B, 4, 4] # print("----",T13.size()) # target = 1 and ref = 3 ref_img_warped, valid_points = inverse_warp_posemat( ref_imgs[1], depth[0][:, 0], T13, intrinsics, args.rotation_mode, args.padding_mode) diff = (ref_imgs[0] - ref_img_warped) * valid_points.unsqueeze(1).float() loss_4 = diff.abs().mean() loss += loss_4 if log_losses: tb_writer.add_scalar('photometric_error', loss_1.item(), n_iter) if w2 > 0: tb_writer.add_scalar('explanability_loss', loss_2.item(), n_iter) tb_writer.add_scalar('disparity_smoothness_loss', loss_3.item(), n_iter) tb_writer.add_scalar('total_loss', loss.item(), n_iter) if log_output: tb_writer.add_image('train Input', tensor2array(tgt_img[0]), n_iter) for k, scaled_maps in enumerate( zip(depth, disparities, warped, diff, explainability_mask)): log_output_tensorboard(tb_writer, "train", 0, " {}".format(k), n_iter, *scaled_maps) # record loss and EPE losses.update(loss.item(), args.batch_size) # compute gradient and do Adam step optimizer.zero_grad() loss.backward() optimizer.step() # measure elapsed time batch_time.update(time.time() - end) end = time.time() with open(args.save_path / args.log_full, 'a') as csvfile: writer = csv.writer(csvfile, delimiter='\t') writer.writerow([ loss.item(), loss_1.item(), loss_2.item() if w2 > 0 else 0, loss_3.item() ]) logger.train_bar.update(i + 1) if i % args.print_freq == 0: logger.train_writer.write('Train: Time {} Data {} Loss {}'.format( batch_time, data_time, losses)) if i >= epoch_size - 1: break n_iter += 1 return losses.avg[0]