def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# set random seeds
if args.seed is not None:
print('set random seed to', args.seed)
set_random_seed(args.seed, deterministic=args.deterministic)
# build the model and load checkpoint
model = build_model(
cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
_ = load_checkpoint(model, args.checkpoint, map_location='cpu')
# sanity check for models without ema
if not model.use_ema:
args.sample_model = 'orig'
if args.sample_model == 'ema':
generator = model.generator_ema
else:
generator = model.generator
mmcv.print_log(f'Sampling model: {args.sample_model}', 'mmgen')
mmcv.print_log(f'Show mode: {args.show_mode}', 'mmgen')
mmcv.print_log(f'Samples path: {args.samples_path}', 'mmgen')
generator.eval()
if not args.use_cpu:
generator = generator.cuda()
if args.show_mode == 'sequence':
assert args.endpoint >= 2
else:
assert args.endpoint >= 2 and args.endpoint % 2 == 0
kwargs = dict(max_batch_size=args.batch_size)
if args.sample_cfg is None:
args.sample_cfg = dict()
kwargs.update(args.sample_cfg)
# get noises corresponding to each endpoint
noise_batch = batch_inference(
generator,
None,
num_batches=args.endpoint,
dict_key='noise_batch' if args.space == 'z' else 'latent',
return_noise=True,
**kwargs)
if args.space == 'w':
kwargs['truncation_latent'] = generator.get_mean_latent()
kwargs['input_is_latent'] = True
if args.show_mode == 'sequence':
results = sample_from_path(generator, noise_batch[:-1, ],
noise_batch[1:, ], args.interval,
args.interp_mode, args.space, **kwargs)
else:
results = sample_from_path(generator, noise_batch[::2, ],
noise_batch[1::2, ], args.interval,
args.interp_mode, args.space, **kwargs)
# reorder results
results = torch.stack(results).permute(1, 0, 2, 3, 4)
_, _, ch, h, w = results.shape
results = results.reshape(-1, ch, h, w)
# rescale value range to [0, 1]
results = ((results + 1) / 2)
results = results[:, [2, 1, 0], ...]
results = results.clamp_(0, 1)
# save image
mmcv.mkdir_or_exist(args.samples_path)
if args.show_mode == 'sequence':
for i in range(results.shape[0]):
image = results[i:i + 1]
save_image(
image,
os.path.join(args.samples_path, '{:0>5d}'.format(i) + '.png'))
else:
save_image(
results,
os.path.join(args.samples_path, 'group.png'),
nrow=args.interval)
help='Other customized kwargs for sampling function')
# system args
parser.add_argument('--num-samples', type=int, default=2)
parser.add_argument('--sample-path', type=str, default=None)
parser.add_argument('--random-seed', type=int, default=2020)
args = parser.parse_args()
set_random_seed(args.random_seed)
cfg = mmcv.Config.fromfile(args.config)
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
mmcv.print_log('Building models and loading checkpoints', 'mmgen')
# build model
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
model.eval()
load_checkpoint(model, args.ckpt, map_location='cpu')
# get generator
if model.use_ema:
generator = model.generator_ema
else:
generator = model.generator
generator = generator.to(device)
dict(type='ImageToTensor', keys=['real_img'])
]
# insert flip aug
if args.flip:
pipeline.insert(
1, dict(type='Flip', keys=['real_img'],
direction='horizontal'))
# build dataloader
if args.imgsdir is not None:
dataset = UnconditionalImageDataset(args.imgsdir, pipeline)
elif args.data_cfg is not None:
# Please make sure the dataset will sample images in `RGB` order.
data_config = Config.fromfile(args.data_cfg)
subset_config = data_config.data.get(args.subset, None)
print_log(subset_config, 'mmgen')
dataset = build_dataset(subset_config)
else:
raise RuntimeError('Please provide imgsdir or data_cfg')
data_loader = build_dataloader(dataset,
args.batch_size,
4,
dist=False,
shuffle=(not args.no_shuffle))
mmcv.mkdir_or_exist(args.pkl_dir)
# build inception network
if args.inception_style == 'stylegan':
inception = torch.jit.load(args.inception_pth).eval().cuda()
def main():
args = parse_args()
# set cudnn_benchmark
cfg = Config.fromfile(args.config)
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# set random seeds
if args.seed is not None:
print('set random seed to', args.seed)
set_random_seed(args.seed, deterministic=args.deterministic)
os.makedirs(args.results_dir, exist_ok=True)
text_inputs = torch.cat([clip.tokenize(args.description)]).cuda()
model = init_model(args.config, args.checkpoint, device='cpu')
g_ema = model.generator_ema
g_ema.eval()
if not args.use_cpu:
g_ema = g_ema.cuda()
mean_latent = g_ema.get_mean_latent()
# if given proj_latent
if args.proj_latent is not None:
mmcv.print_log(f'Load projected latent: {args.proj_latent}', 'mmgen')
proj_file = torch.load(args.proj_latent)
proj_n = len(proj_file)
assert proj_n == 1
noise_batch = []
for img_path in proj_file:
noise_batch.append(proj_file[img_path]['latent'].unsqueeze(0))
latent_code_init = torch.cat(noise_batch, dim=0).cuda()
elif args.mode == 'edit':
latent_code_init_not_trunc = torch.randn(1, 512).cuda()
with torch.no_grad():
results = g_ema([latent_code_init_not_trunc],
return_latents=True,
truncation=args.truncation,
truncation_latent=mean_latent)
latent_code_init = results['latent']
else:
latent_code_init = mean_latent.detach().clone().repeat(1, 18, 1)
with torch.no_grad():
img_orig = g_ema([latent_code_init],
input_is_latent=True,
randomize_noise=False)
latent = latent_code_init.detach().clone()
latent.requires_grad = True
clip_loss = CLIPLoss(clip_model=dict(in_size=g_ema.out_size))
id_loss = FaceIdLoss(
facenet=dict(type='ArcFace', ir_se50_weights=None, device='cuda'))
optimizer = optim.Adam([latent], lr=args.lr)
pbar = tqdm(range(args.step))
mmcv.print_log(f'Description: {args.description}')
for i in pbar:
t = i / args.step
lr = get_lr(t, args.lr)
optimizer.param_groups[0]['lr'] = lr
img_gen = g_ema([latent], input_is_latent=True, randomize_noise=False)
img_gen = img_gen[:, [2, 1, 0], ...]
# clip loss
c_loss = clip_loss(image=img_gen, text=text_inputs)
if args.id_lambda > 0:
i_loss = id_loss(pred=img_gen, gt=img_orig)[0]
else:
i_loss = 0
if args.mode == 'edit':
l2_loss = ((latent_code_init - latent)**2).sum()
loss = c_loss + args.l2_lambda * l2_loss + args.id_lambda * i_loss
else:
loss = c_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.set_description((f'loss: {loss.item():.4f};'))
if args.save_interval > 0 and (i % args.save_interval == 0):
with torch.no_grad():
img_gen = g_ema([latent],
input_is_latent=True,
randomize_noise=False)
img_gen = img_gen[:, [2, 1, 0], ...]
torchvision.utils.save_image(img_gen,
os.path.join(
args.results_dir,
f'{str(i).zfill(5)}.png'),
normalize=True,
range=(-1, 1))
if args.mode == 'edit':
img_orig = img_orig[:, [2, 1, 0], ...]
final_result = torch.cat([img_orig, img_gen])
else:
final_result = img_gen
torchvision.utils.save_image(final_result.detach().cpu(),
os.path.join(args.results_dir,
'final_result.png'),
normalize=True,
scale_each=True,
range=(-1, 1))
def train_step(self,
data_batch,
optimizer,
ddp_reducer=None,
running_status=None):
"""Train step function.
This function implements the standard training iteration for
asynchronous adversarial training. Namely, in each iteration, we first
update discriminator and then compute loss for generator with the newly
updated discriminator.
As for distributed training, we use the ``reducer`` from ddp to
synchronize the necessary params in current computational graph.
Args:
data_batch (dict): Input data from dataloader.
optimizer (dict): Dict contains optimizer for generator and
discriminator.
ddp_reducer (:obj:`Reducer` | None, optional): Reducer from ddp.
It is used to prepare for ``backward()`` in ddp. Defaults to
None.
running_status (dict | None, optional): Contains necessary basic
information for training, e.g., iteration number. Defaults to
None.
Returns:
dict: Contains 'log_vars', 'num_samples', and 'results'.
"""
# get data from data_batch
real_imgs = data_batch['real_img']
# If you adopt ddp, this batch size is local batch size for each GPU.
batch_size = real_imgs.shape[0]
# get running status
if running_status is not None:
curr_iter = running_status['iteration']
else:
# dirty walkround for not providing running status
if not hasattr(self, 'iteration'):
self.iteration = 0
curr_iter = self.iteration
# check if optimizer from model
if hasattr(self, 'optimizer'):
optimizer = self.optimizer
# update current stage
self.curr_stage = int(
min(sum(self.cum_nkimgs <= self.shown_nkimg.item()),
len(self.scales) - 1))
self.curr_scale = self.scales[self.curr_stage]
self._curr_scale_int = self._next_scale_int.clone()
# add new scale and update training status
if self.curr_stage != self.prev_stage:
self.prev_stage = self.curr_stage
self._actual_nkimgs.append(self.shown_nkimg.item())
# reset optimizer
if self.reset_optim_for_new_scale:
optim_cfg = deepcopy(self.train_cfg['optimizer_cfg'])
optim_cfg['generator']['lr'] = self.g_lr_schedule.get(
str(self.curr_scale[0]), self.g_lr_base)
optim_cfg['discriminator']['lr'] = self.d_lr_schedule.get(
str(self.curr_scale[0]), self.d_lr_base)
self.optimizer = build_optimizers(self, optim_cfg)
optimizer = self.optimizer
mmcv.print_log('Reset optimizer for new scale', logger='mmgen')
# update training configs, like transition weight for torgb layers.
# get current transition weight for interpolating two torgb layers
if self.curr_stage == 0:
transition_weight = 1.
else:
transition_weight = (
self.shown_nkimg.item() -
self._actual_nkimgs[-1]) / self.transition_kimgs
# clip to [0, 1]
transition_weight = min(max(transition_weight, 0.), 1.)
self._curr_transition_weight = torch.tensor(transition_weight).to(
self._curr_transition_weight)
# resize real image to target scale
if real_imgs.shape[2:] == self.curr_scale:
pass
elif real_imgs.shape[2] >= self.curr_scale[0] and real_imgs.shape[
3] >= self.curr_scale[1]:
real_imgs = self.interp_real_to(real_imgs, size=self.curr_scale)
else:
raise RuntimeError(
f'The scale of real image {real_imgs.shape[2:]} is smaller '
f'than current scale {self.curr_scale}.')
# disc training
set_requires_grad(self.discriminator, True)
optimizer['discriminator'].zero_grad()
# TODO: add noise sampler to customize noise sampling
with torch.no_grad():
fake_imgs = self.generator(None,
num_batches=batch_size,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight)
# disc pred for fake imgs and real_imgs
disc_pred_fake = self.discriminator(
fake_imgs,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight)
disc_pred_real = self.discriminator(
real_imgs,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight)
# get data dict to compute losses for disc
data_dict_ = dict(
iteration=curr_iter,
gen=self.generator,
disc=self.discriminator,
disc_pred_fake=disc_pred_fake,
disc_pred_real=disc_pred_real,
fake_imgs=fake_imgs,
real_imgs=real_imgs,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight,
gen_partial=partial(self.generator,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight),
disc_partial=partial(self.discriminator,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight))
loss_disc, log_vars_disc = self._get_disc_loss(data_dict_)
# prepare for backward in ddp. If you do not call this function before
# back propagation, the ddp will not dynamically find the used params
# in current computation.
if ddp_reducer is not None:
ddp_reducer.prepare_for_backward(_find_tensors(loss_disc))
loss_disc.backward()
optimizer['discriminator'].step()
# update training log status
if dist.is_initialized():
_batch_size = batch_size * dist.get_world_size()
else:
if 'batch_size' not in running_status:
raise RuntimeError(
'You should offer "batch_size" in running status for PGGAN'
)
_batch_size = running_status['batch_size']
self.shown_nkimg += (_batch_size / 1000.)
log_vars_disc.update(
dict(shown_nkimg=self.shown_nkimg.item(),
curr_scale=self.curr_scale[0],
transition_weight=transition_weight))
# skip generator training if only train discriminator for current
# iteration
if (curr_iter + 1) % self.disc_steps != 0:
results = dict(fake_imgs=fake_imgs.cpu(),
real_imgs=real_imgs.cpu())
outputs = dict(log_vars=log_vars_disc,
num_samples=batch_size,
results=results)
if hasattr(self, 'iteration'):
self.iteration += 1
return outputs
# generator training
set_requires_grad(self.discriminator, False)
optimizer['generator'].zero_grad()
# TODO: add noise sampler to customize noise sampling
fake_imgs = self.generator(None,
num_batches=batch_size,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight)
disc_pred_fake_g = self.discriminator(
fake_imgs,
curr_scale=self.curr_scale[0],
transition_weight=transition_weight)
data_dict_ = dict(iteration=curr_iter,
gen=self.generator,
disc=self.discriminator,
fake_imgs=fake_imgs,
disc_pred_fake_g=disc_pred_fake_g)
loss_gen, log_vars_g = self._get_gen_loss(data_dict_)
# prepare for backward in ddp. If you do not call this function before
# back propagation, the ddp will not dynamically find the used params
# in current computation.
if ddp_reducer is not None:
ddp_reducer.prepare_for_backward(_find_tensors(loss_gen))
loss_gen.backward()
optimizer['generator'].step()
log_vars = {}
log_vars.update(log_vars_g)
log_vars.update(log_vars_disc)
log_vars.update({'batch_size': batch_size})
results = dict(fake_imgs=fake_imgs.cpu(), real_imgs=real_imgs.cpu())
outputs = dict(log_vars=log_vars,
num_samples=batch_size,
results=results)
if hasattr(self, 'iteration'):
self.iteration += 1
# check if a new scale will be added in the next iteration
_curr_stage = int(
min(sum(self.cum_nkimgs <= self.shown_nkimg.item()),
len(self.scales) - 1))
# in the next iteration, we will switch to a new scale
if _curr_stage != self.curr_stage:
# `self._next_scale_int` is updated at the end of `train_step`
self._next_scale_int = self._next_scale_int * 2
return outputs
def single_gpu_evaluation(model,
data_loader,
metrics,
logger,
basic_table_info,
batch_size,
samples_path=None,
**kwargs):
"""Evaluate model with a single gpu.
This method evaluate model with a single gpu and displays eval progress
bar.
Args:
model (nn.Module): Model to be tested.
data_loader (nn.Dataloader): PyTorch data loader.
metrics (list): List of metric objects.
logger (Logger): logger used to record results of evaluation.
batch_size (int): Batch size of images fed into metrics.
basic_table_info (dict): Dictionary containing the basic information \
of the metric table include training configuration and ckpt.
samples_path (str): Used to save generated images. If it's none, we'll
give it a default directory and delete it after finishing the
evaluation. Default to None.
kwargs (dict): Other arguments.
"""
# eval special metric online only
special_metric_name = ['PPL']
for metric in metrics:
assert metric.name not in special_metric_name, 'Please eval '\
f'{metric.name} online'
delete_samples_path = False
if samples_path:
mmcv.mkdir_or_exist(samples_path)
else:
temp_path = './work_dirs/temp_samples'
# if temp_path exists, add suffix
suffix = 1
samples_path = temp_path
while os.path.exists(samples_path):
samples_path = temp_path + '_' + str(suffix)
suffix += 1
os.makedirs(samples_path)
delete_samples_path = True
# sample images
num_exist = len(
list(
mmcv.scandir(samples_path,
suffix=('.jpg', '.png', '.jpeg', '.JPEG'))))
if basic_table_info['num_samples'] > 0:
max_num_images = basic_table_info['num_samples']
else:
max_num_images = max(metric.num_images for metric in metrics)
num_needed = max(max_num_images - num_exist, 0)
if num_needed > 0:
mmcv.print_log(f'Sample {num_needed} fake images for evaluation',
'mmgen')
# define mmcv progress bar
pbar = mmcv.ProgressBar(num_needed)
# if no images, `num_exist` should be zero
for begin in range(num_exist, num_needed, batch_size):
end = min(begin + batch_size, max_num_images)
fakes = model(None,
num_batches=end - begin,
return_loss=False,
sample_model=basic_table_info['sample_model'],
**kwargs)
pbar.update(end - begin)
for i in range(end - begin):
images = fakes[i:i + 1]
images = ((images + 1) / 2)
images = images[:, [2, 1, 0], ...]
images = images.clamp_(0, 1)
image_name = str(begin + i) + '.png'
save_image(images, os.path.join(samples_path, image_name))
if num_needed > 0:
sys.stdout.write('\n')
# return if only save sampled images
if len(metrics) == 0:
return
# empty cache to release GPU memory
torch.cuda.empty_cache()
fake_dataloader = make_vanilla_dataloader(samples_path, batch_size)
for metric in metrics:
mmcv.print_log(f'Evaluate with {metric.name} metric.', 'mmgen')
metric.prepare()
# feed in real images
for data in data_loader:
reals = data['real_img']
num_left = metric.feed(reals, 'reals')
if num_left <= 0:
break
# feed in fake images
for data in fake_dataloader:
fakes = data['real_img']
num_left = metric.feed(fakes, 'fakes')
if num_left <= 0:
break
metric.summary()
table_str = make_metrics_table(basic_table_info['train_cfg'],
basic_table_info['ckpt'],
basic_table_info['sample_model'], metrics)
logger.info('\n' + table_str)
if delete_samples_path:
shutil.rmtree(samples_path)
def main():
args = parse_args()
assert args.out or args.eval or args.format_only or args.show \
or args.show_dir, \
('Please specify at least one operation (save/eval/format/show the '
'results / save the results) with the argument "--out", "--eval"'
', "--format-only", "--show" or "--show-dir"')
if args.eval and args.format_only:
raise ValueError('--eval and --format_only cannot be both specified')
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
raise ValueError('The output file must be a pkl file.')
cfg = Config.fromfile(args.config)
if cfg.get('USE_MMDET', False):
from mmdet.apis import multi_gpu_test, single_gpu_test
from mmdet.datasets import build_dataloader
from mmdet.models import build_detector as build_model
else:
from mmtrack.apis import multi_gpu_test, single_gpu_test
from mmtrack.datasets import build_dataloader
from mmtrack.models import build_model
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# cfg.model.pretrains = None
if hasattr(cfg.model, 'detector'):
cfg.model.detector.pretrained = None
cfg.data.test.test_mode = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# build the dataloader
dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(dataset,
samples_per_gpu=1,
workers_per_gpu=cfg.data.workers_per_gpu,
dist=distributed,
shuffle=False)
logger = get_logger('ParamsSearcher', log_file=args.log)
# get all cases
search_params = get_search_params(cfg.model.tracker, logger=logger)
combinations = [p for p in product(*search_params.values())]
search_cfgs = []
for c in combinations:
search_cfg = dotty(cfg.model.tracker.copy())
for i, k in enumerate(search_params.keys()):
search_cfg[k] = c[i]
search_cfgs.append(dict(search_cfg))
print_log(f'Totally {len(search_cfgs)} cases.', logger)
# init with the first one
cfg.model.tracker = search_cfgs[0].copy()
# build the model and load checkpoint
if cfg.get('test_cfg', False):
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
else:
model = build_model(cfg.model)
# We need call `init_weights()` to load pretained weights in MOT task.
model.init_weights()
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
wrap_fp16_model(model)
if args.checkpoint is not None:
checkpoint = load_checkpoint(model,
args.checkpoint,
map_location='cpu')
if 'meta' in checkpoint and 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['CLASSES']
if not hasattr(model, 'CLASSES'):
model.CLASSES = dataset.CLASSES
if args.fuse_conv_bn:
model = fuse_conv_bn(model)
if not distributed:
model = MMDataParallel(model, device_ids=[0])
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False)
print_log(f'Record {cfg.search_metrics}.', logger)
for i, search_cfg in enumerate(search_cfgs):
if not distributed:
model.module.tracker = build_tracker(search_cfg)
outputs = single_gpu_test(model, data_loader, args.show,
args.show_dir)
else:
model.module.tracker = build_tracker(search_cfg)
outputs = multi_gpu_test(model, data_loader, args.tmpdir,
args.gpu_collect)
rank, _ = get_dist_info()
if rank == 0:
if args.out:
print(f'\nwriting results to {args.out}')
mmcv.dump(outputs, args.out)
kwargs = {} if args.eval_options is None else args.eval_options
if args.format_only:
dataset.format_results(outputs, **kwargs)
if args.eval:
eval_kwargs = cfg.get('evaluation', {}).copy()
# hard-code way to remove EvalHook args
for key in ['interval', 'tmpdir', 'start', 'gpu_collect']:
eval_kwargs.pop(key, None)
eval_kwargs.update(dict(metric=args.eval, **kwargs))
results = dataset.evaluate(outputs, **eval_kwargs)
_records = []
for k in cfg.search_metrics:
if isinstance(results[k], float):
_records.append(f'{(results[k]):.3f}')
else:
_records.append(f'{(results[k])}')
print_log(f'{combinations[i]}: {_records}', logger)
def train_model(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
meta=None):
logger = get_root_logger(cfg.log_level)
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [
build_dataloader(
ds,
cfg.data.samples_per_gpu,
cfg.data.workers_per_gpu,
# cfg.gpus will be ignored if distributed
len(cfg.gpu_ids),
dist=distributed,
persistent_workers=cfg.data.get('persistent_workers', False),
seed=cfg.seed) for ds in dataset
]
# dirty code for use apex amp
# apex.amp request that models should be in cuda device before
# initialization.
if cfg.get('apex_amp', None):
assert distributed, (
'Currently, apex.amp is only supported with DDP training.')
model = model.cuda()
# build optimizer
if cfg.optimizer:
optimizer = build_optimizers(model, cfg.optimizer)
# In GANs, we allow building optimizer in GAN model.
else:
optimizer = None
_use_apex_amp = False
if cfg.get('apex_amp', None):
model, optimizer = apex_amp_initialize(model, optimizer,
**cfg.apex_amp)
_use_apex_amp = True
# put model on gpus
if distributed:
find_unused_parameters = cfg.get('find_unused_parameters', False)
use_ddp_wrapper = cfg.get('use_ddp_wrapper', False)
# Sets the `find_unused_parameters` parameter in
# torch.nn.parallel.DistributedDataParallel
if use_ddp_wrapper:
mmcv.print_log('Use DDP Wrapper.', 'mmgen')
model = DistributedDataParallelWrapper(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
model = MMDataParallel(
model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
# allow users to define the runner
if cfg.get('runner', None):
runner = build_runner(
cfg.runner,
dict(
model=model,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=logger,
use_apex_amp=_use_apex_amp,
meta=meta))
else:
runner = IterBasedRunner(
model,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=logger,
meta=meta)
# set if use dynamic ddp in training
# is_dynamic_ddp=cfg.get('is_dynamic_ddp', False))
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
# In GANs, we can directly optimize parameter in `train_step` function.
if cfg.get('optimizer_cfg', None) is None:
optimizer_config = None
elif fp16_cfg is not None:
raise NotImplementedError('Fp16 has not been supported.')
# optimizer_config = Fp16OptimizerHook(
# **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
# default to use OptimizerHook
elif distributed and 'type' not in cfg.optimizer_config:
optimizer_config = OptimizerHook(**cfg.optimizer_config)
else:
optimizer_config = cfg.optimizer_config
# update `out_dir` in ckpt hook
if cfg.checkpoint_config is not None:
cfg.checkpoint_config['out_dir'] = os.path.join(
cfg.work_dir, cfg.checkpoint_config.get('out_dir', 'ckpt'))
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config,
cfg.get('momentum_config', None))
# # DistSamplerSeedHook should be used with EpochBasedRunner
# if distributed:
# runner.register_hook(DistSamplerSeedHook())
# In general, we do NOT adopt standard evaluation hook in GAN training.
# Thus, if you want a eval hook, you need further define the key of
# 'evaluation' in the config.
# register eval hooks
if validate and cfg.get('evaluation', None) is not None:
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
# Support batch_size > 1 in validation
val_loader_cfg = {
'samples_per_gpu': 1,
'shuffle': False,
'workers_per_gpu': cfg.data.workers_per_gpu,
**cfg.data.get('val_data_loader', {})
}
val_dataloader = build_dataloader(
val_dataset, dist=distributed, **val_loader_cfg)
eval_cfg = deepcopy(cfg.get('evaluation'))
priority = eval_cfg.pop('priority', 'LOW')
eval_cfg.update(dict(dist=distributed, dataloader=val_dataloader))
eval_hook = build_from_cfg(eval_cfg, HOOKS)
runner.register_hook(eval_hook, priority=priority)
# user-defined hooks
if cfg.get('custom_hooks', None):
custom_hooks = cfg.custom_hooks
assert isinstance(custom_hooks, list), \
f'custom_hooks expect list type, but got {type(custom_hooks)}'
for hook_cfg in cfg.custom_hooks:
assert isinstance(hook_cfg, dict), \
'Each item in custom_hooks expects dict type, but got ' \
f'{type(hook_cfg)}'
hook_cfg = hook_cfg.copy()
priority = hook_cfg.pop('priority', 'NORMAL')
hook = build_from_cfg(hook_cfg, HOOKS)
runner.register_hook(hook, priority=priority)
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_iters)
def single_gpu_online_evaluation(model, data_loader, metrics, logger,
basic_table_info, batch_size, **kwargs):
"""Evaluate model with a single gpu in online mode.
This method evaluate model with a single gpu and displays eval progress
bar. Different form `single_gpu_evaluation`, this function will not save
the images or read images from disks. Namely, there do not exist any IO
operations in this function. Thus, in general, `online` mode will achieve a
faster evaluation. However, this mode will take much more memory cost.
Therefore This evaluation function is recommended to evaluate your model
with a single metric.
Args:
model (nn.Module): Model to be tested.
data_loader (nn.Dataloader): PyTorch data loader.
metrics (list): List of metric objects.
logger (Logger): logger used to record results of evaluation.
batch_size (int): Batch size of images fed into metrics.
basic_table_info (dict): Dictionary containing the basic information \
of the metric table include training configuration and ckpt.
kwargs (dict): Other arguments.
"""
# separate metrics into special metrics and vanilla metrics.
# For vanilla metrics, images are generated in a random way, and are
# shared by these metrics. For special metrics like 'PPL', images are
# generated in a metric-special way and not shared between different
# metrics.
special_metrics = []
vanilla_metrics = []
special_metric_name = ['PPL']
for metric in metrics:
if metric.name in special_metric_name:
special_metrics.append(metric)
else:
vanilla_metrics.append(metric)
max_num_images = 0 if len(vanilla_metrics) == 0 else max(
metric.num_images for metric in vanilla_metrics)
for metric in vanilla_metrics:
mmcv.print_log(f'Feed reals to {metric.name} metric.', 'mmgen')
metric.prepare()
pbar = mmcv.ProgressBar(metric.num_real_need)
# feed in real images
for data in data_loader:
reals = data['real_img']
num_feed = metric.feed(reals, 'reals')
if num_feed <= 0:
break
pbar.update(num_feed)
# finish the pbar stdout
sys.stdout.write('\n')
mmcv.print_log(f'Sample {max_num_images} fake images for evaluation',
'mmgen')
# define mmcv progress bar
max_num_images = 0 if len(vanilla_metrics) == 0 else max(
metric.num_fake_need for metric in vanilla_metrics)
pbar = mmcv.ProgressBar(max_num_images)
# sampling fake images and directly send them to metrics
for begin in range(0, max_num_images, batch_size):
end = min(begin + batch_size, max_num_images)
fakes = model(None,
num_batches=end - begin,
return_loss=False,
sample_model=basic_table_info['sample_model'],
**kwargs)
pbar.update(end - begin)
fakes = fakes[:end - begin]
for metric in vanilla_metrics:
# feed in fake images
_ = metric.feed(fakes, 'fakes')
# finish the pbar stdout
sys.stdout.write('\n')
# feed special metric
for metric in special_metrics:
metric.prepare()
fakedata_iterator = iter(
metric.get_sampler(model.module, batch_size,
basic_table_info['sample_model']))
pbar = mmcv.ProgressBar(metric.num_images)
for fakes in fakedata_iterator:
num_left = metric.feed(fakes, 'fakes')
pbar.update(fakes.shape[0])
if num_left <= 0:
break
# finish the pbar stdout
sys.stdout.write('\n')
for metric in metrics:
metric.summary()
table_str = make_metrics_table(basic_table_info['train_cfg'],
basic_table_info['ckpt'],
basic_table_info['sample_model'], metrics)
logger.info('\n' + table_str)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# set random seeds
if args.seed is not None:
print('set random seed to', args.seed)
set_random_seed(args.seed, deterministic=args.deterministic)
# build the model and load checkpoint
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
_ = load_checkpoint(model, args.checkpoint, map_location='cpu')
# sanity check for models without ema
if not model.use_ema:
args.sample_model = 'orig'
if args.sample_model == 'ema':
generator = model.generator_ema
else:
generator = model.generator
mmcv.print_log(f'Sampling model: {args.sample_model}', 'mmgen')
mmcv.print_log(f'Show mode: {args.show_mode}', 'mmgen')
mmcv.print_log(f'Samples path: {args.samples_path}', 'mmgen')
generator.eval()
if not args.use_cpu:
generator = generator.cuda()
# if given proj_latent, reset args.endpoint
if args.proj_latent is not None:
mmcv.print_log(f'Load projected latent: {args.proj_latent}', 'mmgen')
proj_file = torch.load(args.proj_latent)
proj_n = len(proj_file)
setattr(args, 'endpoint', proj_n)
assert args.space == 'w', 'Projected latent are w or w-plus latent.'
noise_batch = []
for img_path in proj_file:
noise_batch.append(proj_file[img_path]['latent'].unsqueeze(0))
noise_batch = torch.cat(noise_batch, dim=0).cuda()
if args.use_cpu:
noise_batch = noise_batch.to('cpu')
if args.show_mode == 'sequence':
assert args.endpoint >= 2
else:
assert args.endpoint >= 2 and args.endpoint % 2 == 0,\
'''We need paired images in group mode,
so keep endpoint an even number'''
kwargs = dict(max_batch_size=args.batch_size)
if args.sample_cfg is None:
args.sample_cfg = dict()
kwargs.update(args.sample_cfg)
# remind users to fixed injected noise
if kwargs.get('randomize_noise', 'True'):
mmcv.print_log(
'''Hint: For Style-Based GAN, you can add
`--sample-cfg randomize_noise=False` to fix injected noises''',
'mmgen')
# get noises corresponding to each endpoint
if not args.proj_latent:
noise_batch = batch_inference(
generator,
None,
num_batches=args.endpoint,
dict_key='noise_batch' if args.space == 'z' else 'latent',
return_noise=True,
**kwargs)
if args.space == 'w':
kwargs['truncation_latent'] = generator.get_mean_latent()
kwargs['input_is_latent'] = True
if args.show_mode == 'sequence':
results = sample_from_path(generator, noise_batch[:-1, ],
noise_batch[1:, ], args.interval,
args.interp_mode, args.space, **kwargs)
else:
results = sample_from_path(generator, noise_batch[::2, ],
noise_batch[1::2, ], args.interval,
args.interp_mode, args.space, **kwargs)
# reorder results
results = torch.stack(results).permute(1, 0, 2, 3, 4)
_, _, ch, h, w = results.shape
results = results.reshape(-1, ch, h, w)
# rescale value range to [0, 1]
results = ((results + 1) / 2)
results = results[:, [2, 1, 0], ...]
results = results.clamp_(0, 1)
# save image
mmcv.mkdir_or_exist(args.samples_path)
if args.show_mode == 'sequence':
if args.export_video:
# render video.
video_out = imageio.get_writer(os.path.join(
args.samples_path, 'lerp.mp4'),
mode='I',
fps=60,
codec='libx264',
bitrate='12M')
video_out = layout_grid(video_out, results)
video_out.close()
else:
for i in range(results.shape[0]):
image = results[i:i + 1]
save_image(
image,
os.path.join(args.samples_path,
'{:0>5d}'.format(i) + '.png'))
else:
if args.export_video:
# render video.
video_out = imageio.get_writer(os.path.join(
args.samples_path, 'lerp.mp4'),
mode='I',
fps=60,
codec='libx264',
bitrate='12M')
n_pair = args.endpoint // 2
grid_w, grid_h = crack_integer(n_pair)
video_out = layout_grid(video_out,
results,
grid_h=grid_h,
grid_w=grid_w)
video_out.close()
else:
save_image(results,
os.path.join(args.samples_path, 'group.png'),
nrow=args.interval)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# set random seeds
if args.seed is not None:
print('set random seed to', args.seed)
set_random_seed(args.seed, deterministic=args.deterministic)
# build the model and load checkpoint
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
_ = load_checkpoint(model, args.checkpoint, map_location='cpu')
# sanity check for models without ema
if not model.use_ema:
args.sample_model = 'orig'
if args.sample_model == 'ema':
generator = model.generator_ema
else:
generator = model.generator
mmcv.print_log(f'Sampling model: {args.sample_model}', 'mmgen')
generator.eval()
device = 'cpu'
if not args.use_cpu:
generator = generator.cuda()
device = 'cuda'
img_size = min(generator.out_size, 256)
transform = transforms.Compose([
transforms.Resize(img_size),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])
# read images
imgs = []
for imgfile in args.files:
img = Image.open(imgfile).convert('RGB')
img = transform(img)
img = img[[2, 1, 0], ...]
imgs.append(img)
imgs = torch.stack(imgs, 0).to(device)
# get mean and standard deviation of style latents
with torch.no_grad():
noise_sample = torch.randn(args.n_mean_latent,
generator.style_channels,
device=device)
latent_out = generator.style_mapping(noise_sample)
latent_mean = latent_out.mean(0)
latent_std = ((latent_out - latent_mean).pow(2).sum() /
args.n_mean_latent)**0.5
latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(
imgs.shape[0], 1)
if args.w_plus:
latent_in = latent_in.unsqueeze(1).repeat(1, generator.num_latents, 1)
latent_in.requires_grad = True
# define lpips loss
percept = PerceptualLoss(use_gpu=device.startswith('cuda'))
# initialize layer noises
noises_single = generator.make_injected_noise()
noises = []
for noise in noises_single:
noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_())
for noise in noises:
noise.requires_grad = True
optimizer = optim.Adam([latent_in] + noises, lr=args.lr)
pbar = tqdm(range(args.total_iters))
# run optimization
for i in pbar:
t = i / args.total_iters
lr = get_lr(t, args.lr, args.lr_rampdown, args.lr_rampup)
optimizer.param_groups[0]['lr'] = lr
noise_strength = latent_std * args.noise * max(
0, 1 - t / args.noise_ramp)**2
latent_n = latent_noise(latent_in, noise_strength.item())
img_gen = generator([latent_n],
input_is_latent=True,
injected_noise=noises)
batch, channel, height, width = img_gen.shape
if height > 256:
factor = height // 256
img_gen = img_gen.reshape(batch, channel, height // factor, factor,
width // factor, factor)
img_gen = img_gen.mean([3, 5])
p_loss = percept(img_gen, imgs).sum()
n_loss = noise_regularize(noises)
mse_loss = F.mse_loss(img_gen, imgs)
loss = p_loss + args.noise_regularize * n_loss + args.mse * mse_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
noise_normalize_(noises)
pbar.set_description(
f' perceptual: {p_loss.item():.4f}, noise regularize:'
f'{n_loss.item():.4f}, mse: {mse_loss.item():.4f}, lr: {lr:.4f}')
results = generator([latent_in.detach().clone()],
input_is_latent=True,
injected_noise=noises)
# rescale value range to [0, 1]
results = ((results + 1) / 2)
results = results[:, [2, 1, 0], ...]
results = results.clamp_(0, 1)
mmcv.mkdir_or_exist(args.results_path)
# save projection results
result_file = OrderedDict()
for i, input_name in enumerate(args.files):
noise_single = []
for noise in noises:
noise_single.append(noise[i:i + 1])
result_file[input_name] = {
'img': img_gen[i],
'latent': latent_in[i],
'injected_noise': noise_single,
}
img_name = os.path.splitext(
os.path.basename(input_name))[0] + '-project.png'
save_image(results[i], os.path.join(args.results_path, img_name))
torch.save(result_file, os.path.join(args.results_path,
'project_result.pt'))
def test_print_log_exception():
with pytest.raises(TypeError):
print_log('welcome', logger=0)
def test_print_log_silent(capsys, caplog):
print_log('welcome', logger='silent')
out, _ = capsys.readouterr()
assert out == ''
assert len(caplog.records) == 0
def test_print_log_print(capsys):
print_log('welcome', logger=None)
out, _ = capsys.readouterr()
assert out == 'welcome\n'
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
rank, _ = get_dist_info()
dirname = os.path.dirname(args.checkpoint)
ckpt = os.path.basename(args.checkpoint)
if 'http' in args.checkpoint:
log_path = None
else:
log_name = ckpt.split('.')[0] + '_eval_log' + '.txt'
log_path = os.path.join(dirname, log_name)
logger = get_root_logger(log_file=log_path,
log_level=cfg.log_level,
file_mode='a')
logger.info('evaluation')
# set random seeds
if args.seed is not None:
if rank == 0:
mmcv.print_log(f'set random seed to {args.seed}', 'mmgen')
set_random_seed(args.seed, deterministic=args.deterministic)
# build the model and load checkpoint
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
# sanity check for models without ema
if not model.use_ema:
args.sample_model = 'orig'
mmcv.print_log(f'Sampling model: {args.sample_model}', 'mmgen')
model.eval()
if not distributed:
_ = load_checkpoint(model, args.checkpoint, map_location='cpu')
model = MMDataParallel(model, device_ids=[0])
# build metrics
if args.eval:
if args.eval[0] == 'none':
# only sample images
metrics = []
assert args.num_samples is not None and args.num_samples > 0
else:
metrics = [
build_metric(cfg.metrics[metric]) for metric in args.eval
]
else:
metrics = [
build_metric(cfg.metrics[metric]) for metric in cfg.metrics
]
basic_table_info = dict(train_cfg=os.path.basename(cfg._filename),
ckpt=ckpt,
sample_model=args.sample_model)
if len(metrics) == 0:
basic_table_info['num_samples'] = args.num_samples
data_loader = None
else:
basic_table_info['num_samples'] = -1
# build the dataloader
if cfg.data.get('test', None) and cfg.data.test.get(
'imgs_root', None):
dataset = build_dataset(cfg.data.test)
elif cfg.data.get('val', None) and cfg.data.val.get(
'imgs_root', None):
dataset = build_dataset(cfg.data.val)
elif cfg.data.get('train', None):
# we assume that the train part should work well
dataset = build_dataset(cfg.data.train)
else:
raise RuntimeError('There is no valid dataset config to run, '
'please check your dataset configs.')
data_loader = build_dataloader(dataset,
samples_per_gpu=args.batch_size,
workers_per_gpu=cfg.data.get(
'val_workers_per_gpu',
cfg.data.workers_per_gpu),
dist=distributed,
shuffle=True)
if args.sample_cfg is None:
args.sample_cfg = dict()
# online mode will not save samples
if args.online and len(metrics) > 0:
single_gpu_online_evaluation(model, data_loader, metrics, logger,
basic_table_info, args.batch_size,
**args.sample_cfg)
else:
single_gpu_evaluation(model, data_loader, metrics, logger,
basic_table_info, args.batch_size,
args.samples_path, **args.sample_cfg)
else:
raise NotImplementedError("We hasn't implemented multi gpu eval yet.")
def after_train_iter(self, runner):
"""The behavior after each train iteration.
Args:
runner (``mmcv.runner.BaseRunner``): The runner.
"""
if not self.every_n_iters(runner, self.interval):
return
runner.model.eval()
# sample fake images
max_num_images = max(metric.num_images for metric in self.metrics)
for metric in self.metrics:
if metric.num_real_feeded >= metric.num_real_need:
continue
mmcv.print_log(f'Feed reals to {metric.name} metric.', 'mmgen')
# feed in real images
for data in self.dataloader:
reals = data['real_img']
num_feed = metric.feed(reals, 'reals')
if num_feed <= 0:
break
mmcv.print_log(f'Sample {max_num_images} fake images for evaluation',
'mmgen')
batch_size = self.dataloader.batch_size
rank, ws = get_dist_info()
total_batch_size = batch_size * ws
# define mmcv progress bar
if rank == 0:
pbar = mmcv.ProgressBar(max_num_images)
# sampling fake images and directly send them to metrics
for _ in range(0, max_num_images, total_batch_size):
with torch.no_grad():
fakes = runner.model(
None,
num_batches=batch_size,
return_loss=False,
**self.sample_kwargs)
for metric in self.metrics:
# feed in fake images
num_left = metric.feed(fakes, 'fakes')
if num_left <= 0:
break
if rank == 0:
pbar.update(total_batch_size)
runner.log_buffer.clear()
# a dirty walkround to change the line at the end of pbar
if rank == 0:
sys.stdout.write('\n')
for metric in self.metrics:
with torch.no_grad():
metric.summary()
for name, val in metric._result_dict.items():
runner.log_buffer.output[name] = val
# record best metric and save the best ckpt
if self.save_best_ckpt and name == self.best_metric:
self._save_best_ckpt(runner, val)
runner.log_buffer.ready = True
runner.model.train()
# clear all current states for next evaluation
for metric in self.metrics:
metric.clear()