def _render_reprojections(self, batch): recon_camera = Camera.vcat( (batch['in_gt']['camera'], batch['out_gt']['camera']), batch_size=self.batch_size) if self._recon_params['generator_input_depth']: depth_noise = self._depth_noise_dist.sample( batch['in']['depth'].size()).to(self.device) depth_real_in = (batch['in']['depth'] + depth_noise).clamp(-1, 1) else: depth_real_in = None # Create IBR images. with torch.set_grad_enabled(self.train_recon): z_obj, z_extra = self._sculptor.encode( self._fuser, camera=batch['in']['camera'], color=batch['in']['image'], depth=depth_real_in, mask=batch['in']['mask'], data_parallel=self.data_parallel) fake, _, _ = self._photographer.decode( z_obj, recon_camera, data_parallel=self.data_parallel) # Reshape things to BVCHW. sections = (self.num_input_views, self.num_output_views) depth_fake_in, depth_fake_out = torch.split(fake['depth'], sections, dim=1) mask_fake_in, mask_fake_out = torch.split(fake['mask'], sections, dim=1) image_reproj, depth_reproj, cam_dists_r, cam_dists_t = ibr.reproject_views_batch( image_in=batch['in']['image'], depth_in=depth_fake_in, depth_out=depth_fake_out, camera_in=batch['in']['camera'], camera_out=batch['out_gt']['camera']) image_reproj = image_reproj * mask_fake_out.unsqueeze(2) depth_reproj = (depth_reproj + 1.0) * mask_fake_out.unsqueeze(2) - 1.0 return ( image_reproj, depth_reproj, mask_fake_out.contiguous(), depth_fake_out.contiguous(), cam_dists_r, cam_dists_t, )
def run_iteration(self, batch, train, is_step): self.mark_time() # Update depth criterion k if applicable. if 'hard_' in self.g_depth_recon_loss_type: self._g_depth_recon_criterion.k = int( self._g_depth_recon_k_scheduler.get(self.epoch)) batch = process_batch(batch, self.cube_size, self.camera_dist, self._sculptor.in_size, self.device, self.random_orientation) if self.reconstruct_input: recon_camera = Camera.vcat( (batch['in_gt']['camera'], batch['out_gt']['camera']), batch_size=self.batch_size) recon_mask = torch.cat( (batch['in_gt']['mask'], batch['out_gt']['mask']), dim=1) recon_image = torch.cat( (batch['in_gt']['image'], batch['out_gt']['image']), dim=1) recon_depth = torch.cat( (batch['in_gt']['depth'], batch['out_gt']['depth']), dim=1) else: recon_camera = batch['out_gt']['camera'] recon_mask = batch['out_gt']['mask'] recon_image = batch['out_gt']['image'] recon_depth = batch['out_gt']['depth'] if not self.color_random_background or self.crop_random_background: batch['in']['image'] = batch['in']['image'] * batch['in']['mask'] if not self.depth_random_background or self.crop_random_background: batch['in']['depth'] = mask_normalized_depth( batch['in']['depth'], batch['in']['mask']) depth_in = None if self.generator_input_depth: depth_noise = self._depth_noise_dist.sample( batch['in']['depth'].size()).to(self.device) depth_in = (batch['in']['depth'] + depth_noise).clamp(-1, 1) data_process_time = self.mark_time() with autocast(): # Evaluate generator. z_obj, z_extra = self._sculptor.encode( self._fuser, camera=batch['in']['camera'], color=batch['in']['image'], depth=depth_in, mask=batch['in']['mask'], data_parallel=self.data_parallel) fake_image, fake_depth, fake_mask, fake_mask_logits, fake_vox_depth = \ self._run_photographer(z_obj, recon_camera, recon_mask) if 'blend_weights' in z_extra: z_weights = z_extra['blend_weights'] else: z_weights = None # Train discriminator. if self._discriminator: d_real, d_fake_d, d_fake_g = self._run_discriminator( fake_image, fake_depth, fake_mask, recon_image, recon_depth, recon_mask) loss_d_real = multiscale_lsgan_loss(d_real, 1) loss_d_fake = multiscale_lsgan_loss(d_fake_d, 0) loss_d = loss_d_real + loss_d_fake loss_g_gan = multiscale_lsgan_loss(d_fake_g, 1) if train: loss_d.backward() if is_step: self._optimizers['discriminator'].step() self.plotter.put_scalar('loss/discriminator/real', loss_d_real) self.plotter.put_scalar('loss/discriminator/fake', loss_d_fake) self.plotter.put_scalar('loss/discriminator/total', loss_d) else: loss_g_gan = torch.tensor(0.0, device=self.device) # Train generator. if self.predict_color: loss_g_color_recon = reduce_loss( self._g_color_recon_criterion(fake_image, recon_image)) else: loss_g_color_recon = torch.tensor(0.0, device=self.device) if self.predict_depth or self.use_occlusion_depth: loss_g_depth_recon = reduce_loss( self._g_depth_recon_criterion(fake_depth, recon_depth)) else: loss_g_depth_recon = torch.tensor(0.0, device=self.device) if self.predict_mask: if self.g_mask_recon_loss_type == 'binary_cross_entropy': y_mask = fake_mask_logits else: y_mask = fake_mask loss_g_mask_recon = reduce_loss( self._g_mask_recon_criterion(y_mask, recon_mask)) loss_g_mask_beta = beta_prior_loss( fake_mask, alpha=self.g_mask_beta_loss_param, beta=self.g_mask_beta_loss_param) else: loss_g_mask_recon = torch.tensor(0.0, device=self.device) loss_g_mask_beta = torch.tensor(0.0, device=self.device) loss_g = (self.g_gan_loss_weight * loss_g_gan + self.g_color_recon_loss_weight * loss_g_color_recon + self.g_depth_recon_loss_weight * loss_g_depth_recon + self.g_mask_recon_loss_weight * loss_g_mask_recon + self.g_mask_beta_loss_weight * loss_g_mask_beta) / self.batch_groups if train: if self.kwargs.get('use_amp', False): self._scaler.scale(loss_g).backward() else: loss_g.backward() if is_step: if self.kwargs.get('use_amp', False): self._scaler.step(self._optimizers['generator']) self._scaler.update() else: self._optimizers['generator'].step() with torch.no_grad(): if self.predict_depth: self.plotter.put_scalar('error/depth/l1', F.l1_loss(fake_depth, recon_depth)) if self.reconstruct_input: self.plotter.put_scalar( 'error/depth/input_l1', F.l1_loss(fake_depth[:, :self.num_input_views], batch['in_gt']['depth'])) self.plotter.put_scalar( 'error/depth/output_l1', F.l1_loss(fake_depth[:, self.num_input_views:], batch['out_gt']['depth'])) if self.predict_mask: self.plotter.put_scalar( 'error/mask/cross_entropy', F.binary_cross_entropy_with_logits(fake_mask_logits, recon_mask)) self.plotter.put_scalar('error/mask/l1', F.l1_loss(fake_mask, recon_mask)) compute_time = self.mark_time() self.plotter.put_scalar('loss/generator/gan', loss_g_gan) self.plotter.put_scalar('loss/generator/recon/color', loss_g_color_recon) self.plotter.put_scalar('loss/generator/recon/depth', loss_g_depth_recon) self.plotter.put_scalar('loss/generator/recon/mask', loss_g_mask_recon) self.plotter.put_scalar('loss/generator/recon/mask_beta', loss_g_mask_beta) self.plotter.put_scalar('loss/generator/total', loss_g) self.plotter.put_scalar('params/input_noise_weight', self.input_noise_weight) if hasattr(self._g_depth_recon_criterion, 'k'): self.plotter.put_scalar('params/depth_loss_k', self._g_depth_recon_criterion.k) self.plotter.put_scalar('time/data_process', data_process_time) self.plotter.put_scalar('time/compute', compute_time) plot_scalar_time = self.mark_time() self.plotter.put_scalar('time/plot/scalars', plot_scalar_time) if self.plotter.is_it_time_yet('histogram'): if self.predict_color: self.plotter.put_histogram('image_fake', fake_image) self.plotter.put_histogram('image_real', recon_image) if self.predict_mask: self.plotter.put_histogram('mask_fake', fake_mask) self.plotter.put_histogram('z_obj', z_obj) if z_weights is not None: self.plotter.put_histogram('z_weights', z_weights) plot_histogram_time = self.mark_time() self.plotter.put_scalar('time/plot/histogram', plot_histogram_time) if self.plotter.is_it_time_yet('show'): self.plotter.put_image( 'inputs', viz.make_grid([ gan_denormalize(batch['in']['image']), viz.colorize_depth(batch['in']['depth']) if self.generator_input_depth else None, viz.colorize_tensor(batch['in']['mask']) if self.generator_input_mask else None, ], row_size=4, stride=2, output_size=64)) with torch.no_grad(): self.plotter.put_image( 'reconstruction', viz.make_grid([ gan_denormalize(recon_image), gan_denormalize(fake_image) if (fake_image is not None) else None, viz.colorize_depth(recon_depth), viz.colorize_depth(fake_depth) if (fake_depth is not None) else None, viz.colorize_tensor( (recon_depth.cpu() - fake_depth.cpu()).abs()) if (fake_depth is not None) else None, viz.colorize_tensor(recon_mask), viz.colorize_tensor(fake_mask) if (fake_mask is not None) else None, viz.colorize_tensor( (recon_mask.cpu() - fake_mask.cpu()).abs()) if (fake_mask is not None) else None, ], stride=8)) plot_images_time = self.mark_time() self.plotter.put_scalar('time/plot/images', plot_images_time)