# 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 json
import os
import time
import torch
import torchvision
import wandb
from torch.cuda.amp import GradScaler, autocast
from tqdm import tqdm
from imaginaire.utils.distributed import is_master, master_only
from imaginaire.utils.distributed import master_only_print as print
from imaginaire.utils.io import save_pilimage_in_jpeg
from imaginaire.utils.meters import Meter
from imaginaire.utils.misc import to_cuda, to_device, requires_grad, to_channels_last
from imaginaire.utils.model_average import (calibrate_batch_norm_momentum,
reset_batch_norm)
from imaginaire.utils.visualization import tensor2pilimage
[docs]class BaseTrainer(object):
r"""Base trainer. We expect that all trainers inherit this class.
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(BaseTrainer, self).__init__()
print('Setup trainer.')
# Initialize models and data loaders.
self.cfg = cfg
self.net_G = net_G
if cfg.trainer.model_average_config.enabled:
# Two wrappers (DDP + model average).
self.net_G_module = self.net_G.module.module
else:
# One wrapper (DDP)
self.net_G_module = self.net_G.module
self.val_data_loader = val_data_loader
self.is_inference = train_data_loader is None
self.net_D = net_D
self.opt_G = opt_G
self.opt_D = opt_D
self.sch_G = sch_G
self.sch_D = sch_D
self.train_data_loader = train_data_loader
if self.cfg.trainer.channels_last:
self.net_G = self.net_G.to(memory_format=torch.channels_last)
self.net_D = self.net_D.to(memory_format=torch.channels_last)
# Initialize amp.
if self.cfg.trainer.amp_config.enabled:
print("Using automatic mixed precision training.")
self.scaler_G = GradScaler(**vars(self.cfg.trainer.amp_config))
self.scaler_D = GradScaler(**vars(self.cfg.trainer.amp_config))
# In order to check whether the discriminator/generator has
# skipped the last parameter update due to gradient overflow.
self.last_step_count_G = 0
self.last_step_count_D = 0
self.skipped_G = False
self.skipped_D = False
# Initialize data augmentation policy.
self.aug_policy = cfg.trainer.aug_policy
print("Augmentation policy: {}".format(self.aug_policy))
# Initialize loss functions.
# All loss names have weights. Some have criterion modules.
# Mapping from loss names to criterion modules.
self.criteria = torch.nn.ModuleDict()
# Mapping from loss names to loss weights.
self.weights = dict()
self.losses = dict(gen_update=dict(), dis_update=dict())
self.gen_losses = self.losses['gen_update']
self.dis_losses = self.losses['dis_update']
self._init_loss(cfg)
for loss_name, loss_weight in self.weights.items():
print("Loss {:<20} Weight {}".format(loss_name, loss_weight))
if loss_name in self.criteria.keys() and \
self.criteria[loss_name] is not None:
self.criteria[loss_name].to('cuda')
if self.is_inference:
# The initialization steps below can be skipped during inference.
return
# Initialize logging attributes.
self.current_iteration = 0
self.current_epoch = 0
self.start_iteration_time = None
self.start_epoch_time = None
self.elapsed_iteration_time = 0
self.time_iteration = None
self.time_epoch = None
self.best_fid = None
if self.cfg.speed_benchmark:
self.accu_gen_forw_iter_time = 0
self.accu_gen_loss_iter_time = 0
self.accu_gen_back_iter_time = 0
self.accu_gen_step_iter_time = 0
self.accu_gen_avg_iter_time = 0
self.accu_dis_forw_iter_time = 0
self.accu_dis_loss_iter_time = 0
self.accu_dis_back_iter_time = 0
self.accu_dis_step_iter_time = 0
# Initialize tensorboard and hparams.
self._init_tensorboard()
self._init_hparams()
# Initialize validation parameters.
self.val_sample_size = getattr(cfg.trainer, 'val_sample_size', 50000)
self.kid_num_subsets = getattr(cfg.trainer, 'kid_num_subsets', 10)
self.kid_subset_size = self.val_sample_size // self.kid_num_subsets
self.metrics_path = os.path.join(torch.hub.get_dir(), 'metrics')
self.best_metrics = {}
self.eval_networks = getattr(cfg.trainer, 'eval_network', ['clean_inception'])
if self.cfg.metrics_iter is None:
self.cfg.metrics_iter = self.cfg.snapshot_save_iter
if self.cfg.metrics_epoch is None:
self.cfg.metrics_epoch = self.cfg.snapshot_save_epoch
# AWS credentials.
if hasattr(cfg, 'aws_credentials_file'):
with open(cfg.aws_credentials_file) as fin:
self.credentials = json.load(fin)
else:
self.credentials = None
if 'TORCH_HOME' not in os.environ:
os.environ['TORCH_HOME'] = os.path.join(
os.environ['HOME'], ".cache")
def _init_tensorboard(self):
r"""Initialize the tensorboard. Different algorithms might require
different performance metrics. Hence, custom tensorboard
initialization might be necessary.
"""
# Logging frequency: self.cfg.logging_iter
self.meters = {}
# Logging frequency: self.cfg.snapshot_save_iter
self.metric_meters = {}
# Logging frequency: self.cfg.image_display_iter
self.image_meter = Meter('images', reduce=False)
def _init_hparams(self):
r"""Initialize a dictionary of hyperparameters that we want to monitor
in the HParams dashboard in tensorBoard.
"""
self.hparam_dict = {}
def _write_tensorboard(self):
r"""Write values to tensorboard. By default, we will log the time used
per iteration, time used per epoch, generator learning rate, and
discriminator learning rate. We will log all the losses as well as
custom meters.
"""
# Logs that are shared by all models.
self._write_to_meters({'time/iteration': self.time_iteration,
'time/epoch': self.time_epoch,
'optim/gen_lr': self.sch_G.get_last_lr()[0],
'optim/dis_lr': self.sch_D.get_last_lr()[0]},
self.meters,
reduce=False)
# Logs for loss values. Different models have different losses.
self._write_loss_meters()
# Other custom logs.
self._write_custom_meters()
def _write_loss_meters(self):
r"""Write all loss values to tensorboard."""
for update, losses in self.losses.items():
# update is 'gen_update' or 'dis_update'.
assert update == 'gen_update' or update == 'dis_update'
for loss_name, loss in losses.items():
if loss is not None:
full_loss_name = update + '/' + loss_name
if full_loss_name not in self.meters.keys():
# Create a new meter if it doesn't exist.
self.meters[full_loss_name] = Meter(
full_loss_name, reduce=True)
self.meters[full_loss_name].write(loss.item())
def _write_custom_meters(self):
r"""Dummy member function to be overloaded by the child class.
In the child class, you can write down whatever you want to track.
"""
pass
@staticmethod
def _write_to_meters(data, meters, reduce=True):
r"""Write values to meters."""
if reduce or is_master():
for key, value in data.items():
if key not in meters:
meters[key] = Meter(key, reduce=reduce)
meters[key].write(value)
def _flush_meters(self, meters):
r"""Flush all meters using the current iteration."""
for meter in meters.values():
meter.flush(self.current_iteration)
def _pre_save_checkpoint(self):
r"""Implement the things you want to do before saving a checkpoint.
For example, you can compute the K-mean features (pix2pixHD) before
saving the model weights to a checkpoint.
"""
pass
[docs] def save_checkpoint(self, current_epoch, current_iteration):
r"""Save network weights, optimizer parameters, scheduler parameters
to a checkpoint.
"""
self._pre_save_checkpoint()
_save_checkpoint(self.cfg,
self.net_G, self.net_D,
self.opt_G, self.opt_D,
self.sch_G, self.sch_D,
current_epoch, current_iteration)
[docs] def load_checkpoint(self, cfg, checkpoint_path, resume=None, load_sch=True):
r"""Load network weights, optimizer parameters, scheduler parameters
from a checkpoint.
Args:
cfg (obj): Global configuration.
checkpoint_path (str): Path to the checkpoint.
resume (bool or None): If not ``None``, will determine whether or
not to load optimizers in addition to network weights.
"""
if os.path.exists(checkpoint_path):
# If checkpoint_path exists, we will load its weights to
# initialize our network.
if resume is None:
resume = False
elif os.path.exists(os.path.join(cfg.logdir, 'latest_checkpoint.txt')):
# This is for resuming the training from the previously saved
# checkpoint.
fn = os.path.join(cfg.logdir, 'latest_checkpoint.txt')
with open(fn, 'r') as f:
line = f.read().splitlines()
checkpoint_path = os.path.join(cfg.logdir, line[0].split(' ')[-1])
if resume is None:
resume = True
else:
# checkpoint not found and not specified. We will train
# everything from scratch.
current_epoch = 0
current_iteration = 0
print('No checkpoint found.')
resume = False
return resume, current_epoch, current_iteration
# Load checkpoint
checkpoint = torch.load(
checkpoint_path, map_location=lambda storage, loc: storage)
current_epoch = 0
current_iteration = 0
if resume:
self.net_G.load_state_dict(checkpoint['net_G'], strict=self.cfg.trainer.strict_resume)
if not self.is_inference:
self.net_D.load_state_dict(checkpoint['net_D'], strict=self.cfg.trainer.strict_resume)
if 'opt_G' in checkpoint:
current_epoch = checkpoint['current_epoch']
current_iteration = checkpoint['current_iteration']
self.opt_G.load_state_dict(checkpoint['opt_G'])
self.opt_D.load_state_dict(checkpoint['opt_D'])
if load_sch:
self.sch_G.load_state_dict(checkpoint['sch_G'])
self.sch_D.load_state_dict(checkpoint['sch_D'])
else:
if self.cfg.gen_opt.lr_policy.iteration_mode:
self.sch_G.last_epoch = current_iteration
else:
self.sch_G.last_epoch = current_epoch
if self.cfg.dis_opt.lr_policy.iteration_mode:
self.sch_D.last_epoch = current_iteration
else:
self.sch_D.last_epoch = current_epoch
print('Load from: {}'.format(checkpoint_path))
else:
print('Load network weights only.')
else:
try:
self.net_G.load_state_dict(checkpoint['net_G'], strict=self.cfg.trainer.strict_resume)
if 'net_D' in checkpoint:
self.net_D.load_state_dict(checkpoint['net_D'], strict=self.cfg.trainer.strict_resume)
except Exception:
if self.cfg.trainer.model_average_config.enabled:
net_G_module = self.net_G.module.module
else:
net_G_module = self.net_G.module
if hasattr(net_G_module, 'load_pretrained_network'):
net_G_module.load_pretrained_network(self.net_G, checkpoint['net_G'])
print('Load generator weights only.')
else:
raise ValueError('Checkpoint cannot be loaded.')
print('Done with loading the checkpoint.')
return resume, current_epoch, current_iteration
[docs] def start_of_epoch(self, current_epoch):
r"""Things to do before an epoch.
Args:
current_epoch (int): Current number of epoch.
"""
self._start_of_epoch(current_epoch)
self.current_epoch = current_epoch
self.start_epoch_time = time.time()
[docs] 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._start_of_iteration(data, current_iteration)
data = to_cuda(data)
if self.cfg.trainer.channels_last:
data = to_channels_last(data)
self.current_iteration = current_iteration
if not self.is_inference:
self.net_D.train()
self.net_G.train()
# torch.cuda.synchronize()
self.start_iteration_time = time.time()
return data
[docs] def end_of_iteration(self, data, current_epoch, current_iteration):
r"""Things to do after an iteration.
Args:
data (dict): Data used for the current iteration.
current_epoch (int): Current number of epoch.
current_iteration (int): Current number of iteration.
"""
self.current_iteration = current_iteration
self.current_epoch = current_epoch
# Update the learning rate policy for the generator if operating in the
# iteration mode.
if self.cfg.gen_opt.lr_policy.iteration_mode:
self.sch_G.step()
# Update the learning rate policy for the discriminator if operating in
# the iteration mode.
if self.cfg.dis_opt.lr_policy.iteration_mode:
self.sch_D.step()
# Accumulate time
# torch.cuda.synchronize()
self.elapsed_iteration_time += time.time() - self.start_iteration_time
# Logging.
if current_iteration % self.cfg.logging_iter == 0:
ave_t = self.elapsed_iteration_time / self.cfg.logging_iter
self.time_iteration = ave_t
print('Iteration: {}, average iter time: '
'{:6f}.'.format(current_iteration, ave_t))
self.elapsed_iteration_time = 0
if self.cfg.speed_benchmark:
# Below code block only needed when analyzing computation
# bottleneck.
print('\tGenerator FWD time {:6f}'.format(
self.accu_gen_forw_iter_time / self.cfg.logging_iter))
print('\tGenerator LOS time {:6f}'.format(
self.accu_gen_loss_iter_time / self.cfg.logging_iter))
print('\tGenerator BCK time {:6f}'.format(
self.accu_gen_back_iter_time / self.cfg.logging_iter))
print('\tGenerator STP time {:6f}'.format(
self.accu_gen_step_iter_time / self.cfg.logging_iter))
print('\tGenerator AVG time {:6f}'.format(
self.accu_gen_avg_iter_time / self.cfg.logging_iter))
print('\tDiscriminator FWD time {:6f}'.format(
self.accu_dis_forw_iter_time / self.cfg.logging_iter))
print('\tDiscriminator LOS time {:6f}'.format(
self.accu_dis_loss_iter_time / self.cfg.logging_iter))
print('\tDiscriminator BCK time {:6f}'.format(
self.accu_dis_back_iter_time / self.cfg.logging_iter))
print('\tDiscriminator STP time {:6f}'.format(
self.accu_dis_step_iter_time / self.cfg.logging_iter))
print('{:6f}'.format(ave_t))
self.accu_gen_forw_iter_time = 0
self.accu_gen_loss_iter_time = 0
self.accu_gen_back_iter_time = 0
self.accu_gen_step_iter_time = 0
self.accu_gen_avg_iter_time = 0
self.accu_dis_forw_iter_time = 0
self.accu_dis_loss_iter_time = 0
self.accu_dis_back_iter_time = 0
self.accu_dis_step_iter_time = 0
self._end_of_iteration(data, current_epoch, current_iteration)
# Save everything to the checkpoint.
if current_iteration % self.cfg.snapshot_save_iter == 0:
if current_iteration >= self.cfg.snapshot_save_start_iter:
self.save_checkpoint(current_epoch, current_iteration)
# Compute metrics.
if current_iteration % self.cfg.metrics_iter == 0:
self.save_image(self._get_save_path('images', 'jpg'), data)
self.write_metrics()
# Compute image to be saved.
elif current_iteration % self.cfg.image_save_iter == 0:
self.save_image(self._get_save_path('images', 'jpg'), data)
elif current_iteration % self.cfg.image_display_iter == 0:
image_path = os.path.join(self.cfg.logdir, 'images', 'current.jpg')
self.save_image(image_path, data)
# Logging.
self._write_tensorboard()
if current_iteration % self.cfg.logging_iter == 0:
# Write all logs to tensorboard.
self._flush_meters(self.meters)
from torch.distributed import barrier
import torch.distributed as dist
if dist.is_initialized():
barrier()
[docs] def end_of_epoch(self, data, current_epoch, current_iteration):
r"""Things to do after an epoch.
Args:
data (dict): Data used for the current iteration.
current_epoch (int): Current number of epoch.
current_iteration (int): Current number of iteration.
"""
# Update the learning rate policy for the generator if operating in the
# epoch mode.
self.current_iteration = current_iteration
self.current_epoch = current_epoch
if not self.cfg.gen_opt.lr_policy.iteration_mode:
self.sch_G.step()
# Update the learning rate policy for the discriminator if operating
# in the epoch mode.
if not self.cfg.dis_opt.lr_policy.iteration_mode:
self.sch_D.step()
elapsed_epoch_time = time.time() - self.start_epoch_time
# Logging.
print('Epoch: {}, total time: {:6f}.'.format(current_epoch,
elapsed_epoch_time))
self.time_epoch = elapsed_epoch_time
self._end_of_epoch(data, current_epoch, current_iteration)
# Save everything to the checkpoint.
if current_iteration % self.cfg.snapshot_save_iter == 0:
if current_epoch >= self.cfg.snapshot_save_start_epoch:
self.save_checkpoint(current_epoch, current_iteration)
# Compute metrics.
if current_iteration % self.cfg.metrics_iter == 0:
self.save_image(self._get_save_path('images', 'jpg'), data)
self.write_metrics()
[docs] def pre_process(self, data):
r"""Custom data pre-processing function. Utilize this function if you
need to preprocess your data before sending it to the generator and
discriminator.
Args:
data (dict): Data used for the current iteration.
"""
[docs] def recalculate_batch_norm_statistics(self, data_loader, averaged=True):
r"""Update the statistics in the moving average model.
Args:
data_loader (torch.utils.data.DataLoader): Data loader for
estimating the statistics.
averaged (Boolean): True/False, we recalculate batch norm statistics for EMA/regular
"""
if not self.cfg.trainer.model_average_config.enabled:
return
if averaged:
net_G = self.net_G.module.averaged_model
else:
net_G = self.net_G_module
model_average_iteration = \
self.cfg.trainer.model_average_config.num_batch_norm_estimation_iterations
if model_average_iteration == 0:
return
with torch.no_grad():
# Accumulate bn stats..
net_G.train()
# Reset running stats.
net_G.apply(reset_batch_norm)
for cal_it, cal_data in enumerate(data_loader):
if cal_it >= model_average_iteration:
print('Done with {} iterations of updating batch norm '
'statistics'.format(model_average_iteration))
break
cal_data = to_device(cal_data, 'cuda')
cal_data = self.pre_process(cal_data)
# Averaging over all batches
net_G.apply(calibrate_batch_norm_momentum)
net_G(cal_data)
[docs] def save_image(self, path, data):
r"""Compute visualization images and save them to the disk.
Args:
path (str): Location of the file.
data (dict): Data used for the current iteration.
"""
self.net_G.eval()
vis_images = self._get_visualizations(data)
if is_master() and vis_images is not None:
vis_images = torch.cat(
[img for img in vis_images if img is not None], dim=3).float()
vis_images = (vis_images + 1) / 2
print('Save output images to {}'.format(path))
vis_images.clamp_(0, 1)
os.makedirs(os.path.dirname(path), exist_ok=True)
image_grid = torchvision.utils.make_grid(
vis_images, nrow=1, padding=0, normalize=False)
if self.cfg.trainer.image_to_tensorboard:
self.image_meter.write_image(image_grid, self.current_iteration)
torchvision.utils.save_image(image_grid, path, nrow=1)
wandb.log({os.path.splitext(os.path.basename(path))[0]: [wandb.Image(path)]})
[docs] def write_metrics(self):
r"""Write metrics to the tensorboard."""
cur_fid = self._compute_fid()
if cur_fid is not None:
if self.best_fid is not None:
self.best_fid = min(self.best_fid, cur_fid)
else:
self.best_fid = cur_fid
metric_dict = {'FID': cur_fid, 'best_FID': self.best_fid}
self._write_to_meters(metric_dict, self.metric_meters, reduce=False)
self._flush_meters(self.metric_meters)
def _get_save_path(self, subdir, ext):
r"""Get the image save path.
Args:
subdir (str): Sub-directory under the main directory for saving
the outputs.
ext (str): Filename extension for the image (e.g., jpg, png, ...).
Return:
(str): image filename to be used to save the visualization results.
"""
subdir_path = os.path.join(self.cfg.logdir, subdir)
if not os.path.exists(subdir_path):
os.makedirs(subdir_path, exist_ok=True)
return os.path.join(
subdir_path, 'epoch_{:05}_iteration_{:09}.{}'.format(
self.current_epoch, self.current_iteration, ext))
def _get_outputs(self, net_D_output, real=True):
r"""Return output values. Note that when the gan mode is relativistic.
It will do the difference before returning.
Args:
net_D_output (dict):
real_outputs (tensor): Real output values.
fake_outputs (tensor): Fake output values.
real (bool): Return real or fake.
"""
def _get_difference(a, b):
r"""Get difference between two lists of tensors or two tensors.
Args:
a: list of tensors or tensor
b: list of tensors or tensor
"""
out = list()
for x, y in zip(a, b):
if isinstance(x, list):
res = _get_difference(x, y)
else:
res = x - y
out.append(res)
return out
if real:
if self.cfg.trainer.gan_relativistic:
return _get_difference(net_D_output['real_outputs'], net_D_output['fake_outputs'])
else:
return net_D_output['real_outputs']
else:
if self.cfg.trainer.gan_relativistic:
return _get_difference(net_D_output['fake_outputs'], net_D_output['real_outputs'])
else:
return net_D_output['fake_outputs']
def _start_of_epoch(self, current_epoch):
r"""Operations to do before starting an epoch.
Args:
current_epoch (int): Current number of epoch.
"""
pass
def _start_of_iteration(self, data, current_iteration):
r"""Operations to do before starting an iteration.
Args:
data (dict): Data used for the current iteration.
current_iteration (int): Current epoch number.
Returns:
(dict): Data used for the current iteration. They might be
processed by the custom _start_of_iteration function.
"""
return data
def _end_of_iteration(self, data, current_epoch, current_iteration):
r"""Operations to do after an iteration.
Args:
data (dict): Data used for the current iteration.
current_epoch (int): Current number of epoch.
current_iteration (int): Current epoch number.
"""
pass
def _end_of_epoch(self, data, current_epoch, current_iteration):
r"""Operations to do after an epoch.
Args:
data (dict): Data used for the current iteration.
current_epoch (int): Current number of epoch.
current_iteration (int): Current epoch number.
"""
pass
def _get_visualizations(self, data):
r"""Compute visualization outputs.
Args:
data (dict): Data used for the current iteration.
"""
return None
def _compute_fid(self):
r"""FID computation function to be overloaded."""
return None
def _init_loss(self, cfg):
r"""Every trainer should implement its own init loss function."""
raise NotImplementedError
[docs] def gen_update(self, data):
r"""Update the generator.
Args:
data (dict): Data used for the current iteration.
"""
update_finished = False
while not update_finished:
# Set requires_grad flags.
requires_grad(self.net_G_module, True)
requires_grad(self.net_D, False)
# Compute the loss.
self._time_before_forward()
with autocast(enabled=self.cfg.trainer.amp_config.enabled):
total_loss = self.gen_forward(data)
if total_loss is None:
return
# Zero-grad and backpropagate the loss.
self.opt_G.zero_grad(set_to_none=True)
self._time_before_backward()
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
)
self.gen_grad_norm = total_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.")
if getattr(self.cfg.gen_opt, 'skip_grad', False):
print(f"Skip gradient update.")
self.opt_G.zero_grad(set_to_none=True)
self.scaler_G.step(self.opt_G)
self.scaler_G.update()
break
# else:
# print(f"Clip gradient norm to "
# f"{self.cfg.gen_opt.clip_grad_norm}.")
# Perform an optimizer step.
self._time_before_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!")
if not torch.isfinite(total_loss):
print("Generator loss is not finite. Skip this iteration!")
update_finished = True
else:
self.last_step_count_G = self.opt_G._step_count
update_finished = True
self._extra_gen_step(data)
# Update model average.
self._time_before_model_avg()
if self.cfg.trainer.model_average_config.enabled:
self.net_G.module.update_average()
self._detach_losses()
self._time_before_leave_gen()
[docs] def gen_forward(self, data):
r"""Every trainer should implement its own generator forward."""
raise NotImplementedError
def _extra_gen_step(self, data):
pass
[docs] def dis_update(self, data):
r"""Update the discriminator.
Args:
data (dict): Data used for the current iteration.
"""
update_finished = False
while not update_finished:
# Set requires_grad flags.
requires_grad(self.net_G_module, False)
requires_grad(self.net_D, True)
# Compute the loss.
self._time_before_forward()
with autocast(enabled=self.cfg.trainer.amp_config.enabled):
total_loss = self.dis_forward(data)
if total_loss is None:
return
# 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
)
self.dis_grad_norm = total_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.")
if getattr(self.cfg.dis_opt, 'skip_grad', False):
print(f"Skip gradient update.")
self.opt_D.zero_grad(set_to_none=True)
self.scaler_D.step(self.opt_D)
self.scaler_D.update()
continue
else:
print(f"Clip gradient norm 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!")
if not torch.isfinite(total_loss):
print("Discriminator loss is not finite. "
"Skip this iteration!")
update_finished = True
else:
self.last_step_count_D = self.opt_D._step_count
update_finished = True
self._extra_dis_step(data)
self._detach_losses()
self._time_before_leave_dis()
[docs] def dis_forward(self, data):
r"""Every trainer should implement its own discriminator forward."""
raise NotImplementedError
def _extra_dis_step(self, data):
pass
[docs] def test(self, data_loader, output_dir, inference_args):
r"""Compute results images for a batch of input data and save the
results in the specified folder.
Args:
data_loader (torch.utils.data.DataLoader): PyTorch dataloader.
output_dir (str): Target location for saving the output image.
"""
if self.cfg.trainer.model_average_config.enabled:
net_G = self.net_G.module.averaged_model
else:
net_G = self.net_G.module
net_G.eval()
print('# of samples %d' % len(data_loader))
for it, data in enumerate(tqdm(data_loader)):
data = self.start_of_iteration(data, current_iteration=-1)
with torch.no_grad():
output_images, file_names = \
net_G.inference(data, **vars(inference_args))
for output_image, file_name in zip(output_images, file_names):
fullname = os.path.join(output_dir, file_name + '.jpg')
output_image = tensor2pilimage(output_image.clamp_(-1, 1),
minus1to1_normalized=True)
save_pilimage_in_jpeg(fullname, output_image)
def _get_total_loss(self, gen_forward):
r"""Return the total loss to be backpropagated.
Args:
gen_forward (bool): If ``True``, backpropagates the generator loss,
otherwise the discriminator loss.
"""
losses = self.gen_losses if gen_forward else self.dis_losses
total_loss = torch.tensor(0., device=torch.device('cuda'))
# Iterates over all possible losses.
for loss_name in self.weights:
# If it is for the current model (gen/dis).
if loss_name in losses:
# Multiply it with the corresponding weight
# and add it to the total loss.
total_loss += losses[loss_name] * self.weights[loss_name]
losses['total'] = total_loss # logging purpose
return total_loss
def _detach_losses(self):
r"""Detach all logging variables to prevent potential memory leak."""
for loss_name in self.gen_losses:
self.gen_losses[loss_name] = self.gen_losses[loss_name].detach()
for loss_name in self.dis_losses:
self.dis_losses[loss_name] = self.dis_losses[loss_name].detach()
def _time_before_forward(self):
r"""
Record time before applying forward.
"""
if self.cfg.speed_benchmark:
torch.cuda.synchronize()
self.forw_time = time.time()
def _time_before_loss(self):
r"""
Record time before computing loss.
"""
if self.cfg.speed_benchmark:
torch.cuda.synchronize()
self.loss_time = time.time()
def _time_before_backward(self):
r"""
Record time before applying backward.
"""
if self.cfg.speed_benchmark:
torch.cuda.synchronize()
self.back_time = time.time()
def _time_before_step(self):
r"""
Record time before updating the weights
"""
if self.cfg.speed_benchmark:
torch.cuda.synchronize()
self.step_time = time.time()
def _time_before_model_avg(self):
r"""
Record time before applying model average.
"""
if self.cfg.speed_benchmark:
torch.cuda.synchronize()
self.avg_time = time.time()
def _time_before_leave_gen(self):
r"""
Record forward, backward, loss, and model average time for the
generator update.
"""
if self.cfg.speed_benchmark:
torch.cuda.synchronize()
end_time = time.time()
self.accu_gen_forw_iter_time += self.loss_time - self.forw_time
self.accu_gen_loss_iter_time += self.back_time - self.loss_time
self.accu_gen_back_iter_time += self.step_time - self.back_time
self.accu_gen_step_iter_time += self.avg_time - self.step_time
self.accu_gen_avg_iter_time += end_time - self.avg_time
def _time_before_leave_dis(self):
r"""
Record forward, backward, loss time for the discriminator update.
"""
if self.cfg.speed_benchmark:
torch.cuda.synchronize()
end_time = time.time()
self.accu_dis_forw_iter_time += self.loss_time - self.forw_time
self.accu_dis_loss_iter_time += self.back_time - self.loss_time
self.accu_dis_back_iter_time += self.step_time - self.back_time
self.accu_dis_step_iter_time += end_time - self.step_time
@master_only
def _save_checkpoint(cfg,
net_G, net_D,
opt_G, opt_D,
sch_G, sch_D,
current_epoch, current_iteration):
r"""Save network weights, optimizer parameters, scheduler parameters
in the checkpoint.
Args:
cfg (obj): Global configuration.
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.
current_epoch (int): Current epoch.
current_iteration (int): Current iteration.
"""
latest_checkpoint_path = 'epoch_{:05}_iteration_{:09}_checkpoint.pt'.format(
current_epoch, current_iteration)
save_path = os.path.join(cfg.logdir, latest_checkpoint_path)
torch.save(
{
'net_G': net_G.state_dict(),
'net_D': net_D.state_dict(),
'opt_G': opt_G.state_dict(),
'opt_D': opt_D.state_dict(),
'sch_G': sch_G.state_dict(),
'sch_D': sch_D.state_dict(),
'current_epoch': current_epoch,
'current_iteration': current_iteration,
},
save_path,
)
fn = os.path.join(cfg.logdir, 'latest_checkpoint.txt')
with open(fn, 'wt') as f:
f.write('latest_checkpoint: %s' % latest_checkpoint_path)
print('Save checkpoint to {}'.format(save_path))
return save_path