Source code for imaginaire.trainers.vid2vid

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
import os

from torch.cuda.amp import autocast
import imageio
import numpy as np
import torch
from tqdm import tqdm

from imaginaire.evaluation.fid import compute_fid
from imaginaire.losses import (FeatureMatchingLoss, FlowLoss, GANLoss,
                               PerceptualLoss)
from imaginaire.model_utils.fs_vid2vid import (concat_frames, detach,
                                               get_fg_mask,
                                               pre_process_densepose, resample)
from imaginaire.trainers.base import BaseTrainer
from imaginaire.utils.distributed import is_master
from imaginaire.utils.distributed import master_only_print as print
from imaginaire.utils.misc import get_nested_attr, split_labels, to_cuda
from imaginaire.utils.visualization import (tensor2flow, tensor2im, tensor2label)
from imaginaire.utils.visualization.pose import tensor2pose


[docs]class Trainer(BaseTrainer): r"""Initialize vid2vid trainer. Args: cfg (obj): Global configuration. net_G (obj): Generator network. net_D (obj): Discriminator network. opt_G (obj): Optimizer for the generator network. opt_D (obj): Optimizer for the discriminator network. sch_G (obj): Scheduler for the generator optimizer. sch_D (obj): Scheduler for the discriminator optimizer. train_data_loader (obj): Train data loader. val_data_loader (obj): Validation data loader. """ def __init__(self, cfg, net_G, net_D, opt_G, opt_D, sch_G, sch_D, train_data_loader, val_data_loader): super(Trainer, self).__init__(cfg, net_G, net_D, opt_G, opt_D, sch_G, sch_D, train_data_loader, val_data_loader) # Below is for testing setting, the FID computation during training # is just for getting a quick idea of the performance. It does not # equal to the final performance evaluation. # Below, we will determine how many videos that we want to do # evaluation, and the length of each video. # It is better to keep the number of videos to be multiple of 8 so # that all the GPUs in a node will contribute equally to the # evaluation. None of them is idol. self.sample_size = ( getattr(cfg.trainer, 'num_videos_to_test', 64), getattr(cfg.trainer, 'num_frames_per_video', 10) ) self.sequence_length = 1 if not self.is_inference: self.train_dataset = self.train_data_loader.dataset self.sequence_length_max = \ min(getattr(cfg.data.train, 'max_sequence_length', 100), self.train_dataset.sequence_length_max) self.Tensor = torch.cuda.FloatTensor self.has_fg = getattr(cfg.data, 'has_foreground', False) self.net_G_output = self.data_prev = None self.net_G_module = self.net_G.module if self.cfg.trainer.model_average_config.enabled: self.net_G_module = self.net_G_module.module def _assign_criteria(self, name, criterion, weight): r"""Assign training loss terms. Args: name (str): Loss name criterion (obj): Loss object. weight (float): Loss weight. It should be non-negative. """ self.criteria[name] = criterion self.weights[name] = weight def _init_loss(self, cfg): r"""Initialize training loss terms. In vid2vid, in addition to the GAN loss, feature matching loss, and perceptual loss used in pix2pixHD, we also add temporal GAN (and feature matching) loss, and flow warping loss. Optionally, we can also add an additional face discriminator for the face region. Args: cfg (obj): Global configuration. """ self.criteria = dict() self.weights = dict() trainer_cfg = cfg.trainer loss_weight = cfg.trainer.loss_weight # GAN loss and feature matching loss. self._assign_criteria('GAN', GANLoss(trainer_cfg.gan_mode), loss_weight.gan) self._assign_criteria('FeatureMatching', FeatureMatchingLoss(), loss_weight.feature_matching) # Perceptual loss. perceptual_loss = cfg.trainer.perceptual_loss self._assign_criteria('Perceptual', PerceptualLoss( network=perceptual_loss.mode, layers=perceptual_loss.layers, weights=perceptual_loss.weights, num_scales=getattr(perceptual_loss, 'num_scales', 1)), loss_weight.perceptual) # L1 Loss. if getattr(loss_weight, 'L1', 0) > 0: self._assign_criteria('L1', torch.nn.L1Loss(), loss_weight.L1) # Whether to add an additional discriminator for specific regions. self.add_dis_cfg = getattr(self.cfg.dis, 'additional_discriminators', None) if self.add_dis_cfg is not None: for name in self.add_dis_cfg: add_dis_cfg = self.add_dis_cfg[name] self.weights['GAN_' + name] = add_dis_cfg.loss_weight self.weights['FeatureMatching_' + name] = \ loss_weight.feature_matching # Temporal GAN loss. self.num_temporal_scales = get_nested_attr(self.cfg.dis, 'temporal.num_scales', 0) for s in range(self.num_temporal_scales): self.weights['GAN_T%d' % s] = loss_weight.temporal_gan self.weights['FeatureMatching_T%d' % s] = \ loss_weight.feature_matching # Flow loss. It consists of three parts: L1 loss compared to GT, # warping loss when used to warp images, and loss on the occlusion mask. self.use_flow = hasattr(cfg.gen, 'flow') if self.use_flow: self.criteria['Flow'] = FlowLoss(cfg) self.weights['Flow'] = self.weights['Flow_L1'] = \ self.weights['Flow_Warp'] = \ self.weights['Flow_Mask'] = loss_weight.flow # Other custom losses. self._define_custom_losses() def _define_custom_losses(self): r"""All other custom losses are defined here.""" pass def _start_of_epoch(self, current_epoch): r"""Things to do before an epoch. When current_epoch is smaller than $(single_frame_epoch), we only train a single frame and the generator is just an image generator. After that, we start doing temporal training and train multiple frames. We will double the number of training frames every $(num_epochs_temporal_step) epochs. Args: current_epoch (int): Current number of epoch. """ cfg = self.cfg # Only generates one frame at the beginning of training if current_epoch < cfg.single_frame_epoch: self.train_dataset.sequence_length = 1 # Then add the temporal network to generator, and train multiple frames. elif current_epoch == cfg.single_frame_epoch: self.init_temporal_network() # Double the length of training sequence every few epochs. temp_epoch = current_epoch - cfg.single_frame_epoch if temp_epoch > 0: sequence_length = \ cfg.data.train.initial_sequence_length * \ (2 ** (temp_epoch // cfg.num_epochs_temporal_step)) sequence_length = min(sequence_length, self.sequence_length_max) if sequence_length > self.sequence_length: self.sequence_length = sequence_length self.train_dataset.set_sequence_length(sequence_length) print('------- Updating sequence length to %d -------' % sequence_length)
[docs] def init_temporal_network(self): r"""Initialize temporal training when beginning to train multiple frames. Set the sequence length to $(initial_sequence_length). """ self.tensorboard_init = False # Update training sequence length. self.sequence_length = self.cfg.data.train.initial_sequence_length if not self.is_inference: self.train_dataset.set_sequence_length(self.sequence_length) print('------ Now start training %d frames -------' % self.sequence_length)
def _start_of_iteration(self, data, current_iteration): r"""Things to do before an iteration. Args: data (dict): Data used for the current iteration. current_iteration (int): Current number of iteration. """ data = self.pre_process(data) return to_cuda(data)
[docs] def pre_process(self, data): r"""Do any data pre-processing here. Args: data (dict): Data used for the current iteration. """ data_cfg = self.cfg.data if hasattr(data_cfg, 'for_pose_dataset') and \ ('pose_maps-densepose' in data_cfg.input_labels): pose_cfg = data_cfg.for_pose_dataset data['label'] = pre_process_densepose(pose_cfg, data['label'], self.is_inference) return data
[docs] def post_process(self, data, net_G_output): r"""Do any postprocessing of the data / output here. Args: data (dict): Training data at the current iteration. net_G_output (dict): Output of the generator. """ return data, net_G_output
[docs] def gen_update(self, data): r"""Update the vid2vid generator. We update in the fashion of dis_update (frame 1), gen_update (frame 1), dis_update (frame 2), gen_update (frame 2), ... in each iteration. Args: data (dict): Training data at the current iteration. """ # Whether to reuse generator output for both gen_update and # dis_update. It saves time but consumes a bit more memory. reuse_gen_output = getattr(self.cfg.trainer, 'reuse_gen_output', True) past_frames = [None, None] net_G_output = None data_prev = None for t in range(self.sequence_length): data_t = self.get_data_t(data, net_G_output, data_prev, t) data_prev = data_t # Discriminator update. if reuse_gen_output: net_G_output = self.net_G(data_t) else: with torch.no_grad(): net_G_output = self.net_G(data_t) data_t, net_G_output = self.post_process(data_t, net_G_output) # Get losses and update D if image generated by network in training. if 'fake_images_source' not in net_G_output: net_G_output['fake_images_source'] = 'in_training' if net_G_output['fake_images_source'] != 'pretrained': net_D_output, _ = self.net_D(data_t, detach(net_G_output), past_frames) self.get_dis_losses(net_D_output) # Generator update. if not reuse_gen_output: net_G_output = self.net_G(data_t) data_t, net_G_output = self.post_process(data_t, net_G_output) # Get losses and update G if image generated by network in training. if 'fake_images_source' not in net_G_output: net_G_output['fake_images_source'] = 'in_training' if net_G_output['fake_images_source'] != 'pretrained': net_D_output, past_frames = \ self.net_D(data_t, net_G_output, past_frames) self.get_gen_losses(data_t, net_G_output, net_D_output) # update average if self.cfg.trainer.model_average_config.enabled: self.net_G.module.update_average()
[docs] def dis_update(self, data): r"""The update is already done in gen_update. Args: data (dict): Training data at the current iteration. """ pass
[docs] def reset(self): r"""Reset the trainer (for inference) at the beginning of a sequence. """ # print('Resetting trainer.') self.net_G_output = self.data_prev = None self.t = 0 self.test_in_model_average_mode = getattr( self, 'test_in_model_average_mode', self.cfg.trainer.model_average_config.enabled) if self.test_in_model_average_mode: net_G_module = self.net_G.module.averaged_model else: net_G_module = self.net_G.module if hasattr(net_G_module, 'reset'): net_G_module.reset()
[docs] def create_sequence_output_dir(self, output_dir, key): r"""Create output subdir for this sequence. Args: output_dir (str): Root output dir. key (str): LMDB key which contains sequence name and file name. Returns: output_dir (str): Output subdir for this sequence. seq_name (str): Name of this sequence. """ seq_dir = '/'.join(key.split('/')[:-1]) output_dir = os.path.join(output_dir, seq_dir) os.makedirs(output_dir, exist_ok=True) seq_name = seq_dir.replace('/', '-') return output_dir, seq_name
[docs] def test(self, test_data_loader, root_output_dir, inference_args): r"""Run inference on all sequences. Args: test_data_loader (object): Test data loader. root_output_dir (str): Location to dump outputs. inference_args (optional): Optional args. """ # Go over all sequences. loader = test_data_loader num_inference_sequences = loader.dataset.num_inference_sequences() for sequence_idx in range(num_inference_sequences): loader.dataset.set_inference_sequence_idx(sequence_idx) print('Seq id: %d, Seq length: %d' % (sequence_idx + 1, len(loader))) # Reset model at start of new inference sequence. self.reset() self.sequence_length = len(loader) # Go over all frames of this sequence. video = [] for idx, data in enumerate(tqdm(loader)): key = data['key']['images'][0][0] filename = key.split('/')[-1] # Create output dir for this sequence. if idx == 0: output_dir, seq_name = \ self.create_sequence_output_dir(root_output_dir, key) video_path = os.path.join(output_dir, '..', seq_name) # Get output and save images. data['img_name'] = filename data = self.start_of_iteration(data, current_iteration=-1) output = self.test_single(data, output_dir, inference_args) video.append(output) # Save output as mp4. imageio.mimsave(video_path + '.mp4', video, fps=15)
[docs] def test_single(self, data, output_dir=None, inference_args=None): r"""The inference function. If output_dir exists, also save the output image. Args: data (dict): Training data at the current iteration. output_dir (str): Save image directory. inference_args (obj): Inference args. """ if getattr(inference_args, 'finetune', False): if not getattr(self, 'has_finetuned', False): self.finetune(data, inference_args) net_G = self.net_G if self.test_in_model_average_mode: net_G = net_G.module.averaged_model net_G.eval() data_t = self.get_data_t(data, self.net_G_output, self.data_prev, 0) if self.is_inference or self.sequence_length > 1: self.data_prev = data_t # Generator forward. with torch.no_grad(): self.net_G_output = net_G(data_t) if output_dir is None: return self.net_G_output save_fake_only = getattr(inference_args, 'save_fake_only', False) if save_fake_only: image_grid = tensor2im(self.net_G_output['fake_images'])[0] else: vis_images = self.get_test_output_images(data) image_grid = np.hstack([np.vstack(im) for im in vis_images if im is not None]) if 'img_name' in data: save_name = data['img_name'].split('.')[0] + '.jpg' else: save_name = '%04d.jpg' % self.t output_filename = os.path.join(output_dir, save_name) os.makedirs(output_dir, exist_ok=True) imageio.imwrite(output_filename, image_grid) self.t += 1 return image_grid
[docs] def get_test_output_images(self, data): r"""Get the visualization output of test function. Args: data (dict): Training data at the current iteration. """ vis_images = [ self.visualize_label(data['label'][:, -1]), tensor2im(data['images'][:, -1]), tensor2im(self.net_G_output['fake_images']), ] return vis_images
[docs] def gen_frames(self, data, use_model_average=False): r"""Generate a sequence of frames given a sequence of data. Args: data (dict): Training data at the current iteration. use_model_average (bool): Whether to use model average for update or not. """ net_G_output = None # Previous generator output. data_prev = None # Previous data. if use_model_average: net_G = self.net_G.module.averaged_model else: net_G = self.net_G # Iterate through the length of sequence. all_info = {'inputs': [], 'outputs': []} for t in range(self.sequence_length): # Get the data at the current time frame. data_t = self.get_data_t(data, net_G_output, data_prev, t) data_prev = data_t # Generator forward. with torch.no_grad(): net_G_output = net_G(data_t) # Do any postprocessing if necessary. data_t, net_G_output = self.post_process(data_t, net_G_output) if t == 0: # Get the output at beginning of sequence for visualization. first_net_G_output = net_G_output all_info['inputs'].append(data_t) all_info['outputs'].append(net_G_output) return first_net_G_output, net_G_output, all_info
[docs] def get_gen_losses(self, data_t, net_G_output, net_D_output): r"""Compute generator losses. Args: data_t (dict): Training data at the current time t. net_G_output (dict): Output of the generator. net_D_output (dict): Output of the discriminator. """ update_finished = False while not update_finished: with autocast(enabled=self.cfg.trainer.amp_config.enabled): # Individual frame GAN loss and feature matching loss. self.gen_losses['GAN'], self.gen_losses['FeatureMatching'] = \ self.compute_gan_losses(net_D_output['indv'], dis_update=False) # Perceptual loss. self.gen_losses['Perceptual'] = self.criteria['Perceptual']( net_G_output['fake_images'], data_t['image']) # L1 loss. if getattr(self.cfg.trainer.loss_weight, 'L1', 0) > 0: self.gen_losses['L1'] = self.criteria['L1']( net_G_output['fake_images'], data_t['image']) # Raw (hallucinated) output image losses (GAN and perceptual). if 'raw' in net_D_output: raw_GAN_losses = self.compute_gan_losses( net_D_output['raw'], dis_update=False ) fg_mask = get_fg_mask(data_t['label'], self.has_fg) raw_perceptual_loss = self.criteria['Perceptual']( net_G_output['fake_raw_images'] * fg_mask, data_t['image'] * fg_mask) self.gen_losses['GAN'] += raw_GAN_losses[0] self.gen_losses['FeatureMatching'] += raw_GAN_losses[1] self.gen_losses['Perceptual'] += raw_perceptual_loss # Additional discriminator losses. if self.add_dis_cfg is not None: for name in self.add_dis_cfg: (self.gen_losses['GAN_' + name], self.gen_losses['FeatureMatching_' + name]) = \ self.compute_gan_losses(net_D_output[name], dis_update=False) # Flow and mask loss. if self.use_flow: (self.gen_losses['Flow_L1'], self.gen_losses['Flow_Warp'], self.gen_losses['Flow_Mask']) = self.criteria['Flow']( data_t, net_G_output, self.current_epoch) # Temporal GAN loss and feature matching loss. if self.cfg.trainer.loss_weight.temporal_gan > 0: if self.sequence_length > 1: for s in range(self.num_temporal_scales): loss_GAN, loss_FM = self.compute_gan_losses( net_D_output['temporal_%d' % s], dis_update=False ) self.gen_losses['GAN_T%d' % s] = loss_GAN self.gen_losses['FeatureMatching_T%d' % s] = loss_FM # Other custom losses. self._get_custom_gen_losses(data_t, net_G_output, net_D_output) # Sum all losses together. total_loss = self.Tensor(1).fill_(0) for key in self.gen_losses: if key != 'total': total_loss += self.gen_losses[key] * self.weights[key] self.gen_losses['total'] = total_loss # Zero-grad and backpropagate the loss. self.opt_G.zero_grad(set_to_none=True) self.scaler_G.scale(total_loss).backward() # Optionally clip gradient norm. if hasattr(self.cfg.gen_opt, 'clip_grad_norm'): self.scaler_G.unscale_(self.opt_G) total_norm = torch.nn.utils.clip_grad_norm_( self.net_G_module.parameters(), self.cfg.gen_opt.clip_grad_norm ) if torch.isfinite(total_norm) and \ total_norm > self.cfg.gen_opt.clip_grad_norm: print(f"Gradient norm of the generator ({total_norm}) " f"too large, clipping it to " f"{self.cfg.gen_opt.clip_grad_norm}.") # Perform an optimizer step. self.scaler_G.step(self.opt_G) self.scaler_G.update() # Whether the step above was skipped. if self.last_step_count_G == self.opt_G._step_count: print("Generator overflowed!") else: self.last_step_count_G = self.opt_G._step_count update_finished = True
def _get_custom_gen_losses(self, data_t, net_G_output, net_D_output): r"""All other custom generator losses go here. Args: data_t (dict): Training data at the current time t. net_G_output (dict): Output of the generator. net_D_output (dict): Output of the discriminator. """ pass
[docs] def get_dis_losses(self, net_D_output): r"""Compute discriminator losses. Args: net_D_output (dict): Output of the discriminator. """ update_finished = False while not update_finished: with autocast(enabled=self.cfg.trainer.amp_config.enabled): # Individual frame GAN loss. self.dis_losses['GAN'] = self.compute_gan_losses( net_D_output['indv'], dis_update=True ) # Raw (hallucinated) output image GAN loss. if 'raw' in net_D_output: raw_loss = self.compute_gan_losses(net_D_output['raw'], dis_update=True) self.dis_losses['GAN'] += raw_loss # Additional GAN loss. if self.add_dis_cfg is not None: for name in self.add_dis_cfg: self.dis_losses['GAN_' + name] = \ self.compute_gan_losses(net_D_output[name], dis_update=True) # Temporal GAN loss. if self.cfg.trainer.loss_weight.temporal_gan > 0: if self.sequence_length > 1: for s in range(self.num_temporal_scales): self.dis_losses['GAN_T%d' % s] = \ self.compute_gan_losses( net_D_output['temporal_%d' % s], dis_update=True ) # Other custom losses. self._get_custom_dis_losses(net_D_output) # Sum all losses together. total_loss = self.Tensor(1).fill_(0) for key in self.dis_losses: if key != 'total': total_loss += self.dis_losses[key] * self.weights[key] self.dis_losses['total'] = total_loss # Zero-grad and backpropagate the loss. self.opt_D.zero_grad(set_to_none=True) self._time_before_backward() self.scaler_D.scale(total_loss).backward() # Optionally clip gradient norm. if hasattr(self.cfg.dis_opt, 'clip_grad_norm'): self.scaler_D.unscale_(self.opt_D) total_norm = torch.nn.utils.clip_grad_norm_( self.net_D.parameters(), self.cfg.dis_opt.clip_grad_norm ) if torch.isfinite(total_norm) and \ total_norm > self.cfg.dis_opt.clip_grad_norm: print(f"Gradient norm of the discriminator ({total_norm}) " f"too large, clipping it to " f"{self.cfg.dis_opt.clip_grad_norm}.") # Perform an optimizer step. self._time_before_step() self.scaler_D.step(self.opt_D) self.scaler_D.update() # Whether the step above was skipped. if self.last_step_count_D == self.opt_D._step_count: print("Discriminator overflowed!") else: self.last_step_count_D = self.opt_D._step_count update_finished = True
def _get_custom_dis_losses(self, net_D_output): r"""All other custom losses go here. Args: net_D_output (dict): Output of the discriminator. """ pass
[docs] def compute_gan_losses(self, net_D_output, dis_update): r"""Compute GAN loss and feature matching loss. Args: net_D_output (dict): Output of the discriminator. dis_update (bool): Whether to update discriminator. """ if net_D_output['pred_fake'] is None: return self.Tensor(1).fill_(0) if dis_update else [ self.Tensor(1).fill_(0), self.Tensor(1).fill_(0)] if dis_update: # Get the GAN loss for real/fake outputs. GAN_loss = \ self.criteria['GAN'](net_D_output['pred_fake']['output'], False, dis_update=True) + \ self.criteria['GAN'](net_D_output['pred_real']['output'], True, dis_update=True) return GAN_loss else: # Get the GAN loss and feature matching loss for fake output. GAN_loss = self.criteria['GAN']( net_D_output['pred_fake']['output'], True, dis_update=False) FM_loss = self.criteria['FeatureMatching']( net_D_output['pred_fake']['features'], net_D_output['pred_real']['features']) return GAN_loss, FM_loss
[docs] def get_data_t(self, data, net_G_output, data_prev, t): r"""Get data at current time frame given the sequence of data. Args: data (dict): Training data for current iteration. net_G_output (dict): Output of the generator (for previous frame). data_prev (dict): Data for previous frame. t (int): Current time. """ label = data['label'][:, t] image = data['images'][:, t] if data_prev is not None: # Concat previous labels/fake images to the ones before. num_frames_G = self.cfg.data.num_frames_G prev_labels = concat_frames(data_prev['prev_labels'], data_prev['label'], num_frames_G - 1) prev_images = concat_frames( data_prev['prev_images'], net_G_output['fake_images'].detach(), num_frames_G - 1) else: prev_labels = prev_images = None data_t = dict() data_t['label'] = label data_t['image'] = image data_t['prev_labels'] = prev_labels data_t['prev_images'] = prev_images data_t['real_prev_image'] = data['images'][:, t - 1] if t > 0 else None return data_t
def _end_of_iteration(self, data, current_epoch, current_iteration): r"""Print the errors to console.""" if not torch.distributed.is_initialized(): if current_iteration % self.cfg.logging_iter == 0: message = '(epoch: %d, iters: %d) ' % (current_epoch, current_iteration) for k, v in self.gen_losses.items(): if k != 'total': message += '%s: %.3f, ' % (k, v) message += '\n' for k, v in self.dis_losses.items(): if k != 'total': message += '%s: %.3f, ' % (k, v) print(message)
[docs] def write_metrics(self): r"""If moving average model presents, we have two meters one for regular FID and one for average FID. If no moving average model, we just report average FID. """ if self.cfg.trainer.model_average_config.enabled: regular_fid, average_fid = self._compute_fid() if regular_fid is None or average_fid is None: return metric_dict = {'FID/average': average_fid, 'FID/regular': regular_fid} self._write_to_meters(metric_dict, self.metric_meters, reduce=False) else: regular_fid = self._compute_fid() if regular_fid is None: return metric_dict = {'FID/regular': regular_fid} self._write_to_meters(metric_dict, self.metric_meters, reduce=False) self._flush_meters(self.metric_meters)
def _compute_fid(self): r"""Compute FID values.""" self.net_G.eval() self.net_G_output = None # Due to complicated video evaluation procedure we are using, we will # pass the trainer to the evaluation code instead of the # generator network. # net_G_for_evaluation = self.net_G trainer = self self.test_in_model_average_mode = False regular_fid_path = self._get_save_path('regular_fid', 'npy') few_shot = True if 'few_shot' in self.cfg.data.type else False regular_fid_value = compute_fid(regular_fid_path, self.val_data_loader, trainer, sample_size=self.sample_size, is_video=True, few_shot_video=few_shot) print('Epoch {:05}, Iteration {:09}, Regular FID {}'.format( self.current_epoch, self.current_iteration, regular_fid_value)) if self.cfg.trainer.model_average_config.enabled: # Due to complicated video evaluation procedure we are using, # we will pass the trainer to the evaluation code instead of the # generator network. # avg_net_G_for_evaluation = self.net_G.module.averaged_model trainer_avg_mode = self self.test_in_model_average_mode = True # The above flag will be reset after computing FID. fid_path = self._get_save_path('average_fid', 'npy') few_shot = True if 'few_shot' in self.cfg.data.type else False fid_value = compute_fid(fid_path, self.val_data_loader, trainer_avg_mode, sample_size=self.sample_size, is_video=True, few_shot_video=few_shot) print('Epoch {:05}, Iteration {:09}, Average FID {}'.format( self.current_epoch, self.current_iteration, fid_value)) self.net_G.float() return regular_fid_value, fid_value else: self.net_G.float() return regular_fid_value
[docs] def visualize_label(self, label): r"""Visualize the input label when saving to image. Args: label (tensor): Input label tensor. """ cfgdata = self.cfg.data if hasattr(cfgdata, 'for_pose_dataset'): label = tensor2pose(self.cfg, label) elif hasattr(cfgdata, 'input_labels') and \ 'seg_maps' in cfgdata.input_labels: for input_type in cfgdata.input_types: if 'seg_maps' in input_type: num_labels = cfgdata.one_hot_num_classes.seg_maps label = tensor2label(label, num_labels) elif getattr(cfgdata, 'label_channels', 1) > 3: label = tensor2im(label.sum(1, keepdim=True)) else: label = tensor2im(label) return label
[docs] def save_image(self, path, data): r"""Save the output images to path. Note when the generate_raw_output is FALSE. Then, first_net_G_output['fake_raw_images'] is None and will not be displayed. In model average mode, we will plot the flow visualization twice. Args: path (str): Save path. data (dict): Training data for current iteration. """ self.net_G.eval() if self.cfg.trainer.model_average_config.enabled: self.net_G.module.averaged_model.eval() self.net_G_output = None with torch.no_grad(): first_net_G_output, net_G_output, all_info = self.gen_frames(data) if self.cfg.trainer.model_average_config.enabled: first_net_G_output_avg, net_G_output_avg, _ = self.gen_frames( data, use_model_average=True) # Visualize labels. label_lengths = self.train_data_loader.dataset.get_label_lengths() labels = split_labels(data['label'], label_lengths) vis_labels_start, vis_labels_end = [], [] for key, value in labels.items(): if key == 'seg_maps': vis_labels_start.append(self.visualize_label(value[:, -1])) vis_labels_end.append(self.visualize_label(value[:, 0])) else: vis_labels_start.append(tensor2im(value[:, -1])) vis_labels_end.append(tensor2im(value[:, 0])) if is_master(): vis_images = [ *vis_labels_start, tensor2im(data['images'][:, -1]), tensor2im(net_G_output['fake_images']), tensor2im(net_G_output['fake_raw_images'])] if self.cfg.trainer.model_average_config.enabled: vis_images += [ tensor2im(net_G_output_avg['fake_images']), tensor2im(net_G_output_avg['fake_raw_images'])] if self.sequence_length > 1: vis_images_first = [ *vis_labels_end, tensor2im(data['images'][:, 0]), tensor2im(first_net_G_output['fake_images']), tensor2im(first_net_G_output['fake_raw_images']) ] if self.cfg.trainer.model_average_config.enabled: vis_images_first += [ tensor2im(first_net_G_output_avg['fake_images']), tensor2im(first_net_G_output_avg['fake_raw_images'])] if self.use_flow: flow_gt, conf_gt = self.criteria['Flow'].flowNet( data['images'][:, -1], data['images'][:, -2]) warped_image_gt = resample(data['images'][:, -1], flow_gt) vis_images_first += [ tensor2flow(flow_gt), tensor2im(conf_gt, normalize=False), tensor2im(warped_image_gt), ] vis_images += [ tensor2flow(net_G_output['fake_flow_maps']), tensor2im(net_G_output['fake_occlusion_masks'], normalize=False), tensor2im(net_G_output['warped_images']), ] if self.cfg.trainer.model_average_config.enabled: vis_images_first += [ tensor2flow(flow_gt), tensor2im(conf_gt, normalize=False), tensor2im(warped_image_gt), ] vis_images += [ tensor2flow(net_G_output_avg['fake_flow_maps']), tensor2im(net_G_output_avg['fake_occlusion_masks'], normalize=False), tensor2im(net_G_output_avg['warped_images'])] vis_images = [[np.vstack((im_first, im)) for im_first, im in zip(imgs_first, imgs)] for imgs_first, imgs in zip(vis_images_first, vis_images) if imgs is not None] image_grid = np.hstack([np.vstack(im) for im in vis_images if im is not None]) print('Save output images to {}'.format(path)) os.makedirs(os.path.dirname(path), exist_ok=True) imageio.imwrite(path, image_grid) # Gather all outputs for dumping into video. if self.sequence_length > 1: output_images = [] for item in all_info['outputs']: output_images.append(tensor2im(item['fake_images'])[0]) imageio.mimwrite(os.path.splitext(path)[0] + '.mp4', output_images, fps=2, macro_block_size=None) self.net_G.float()