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 main():
args = parse_args()
cfg = Config.fromfile(args.config)
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:
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)
assert isinstance(model, BaseTranslationModel)
# 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()
_ = 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]
# get source domain and target domain
target_domain = args.target_domain
if target_domain is None:
target_domain = model.module._default_domain
source_domain = model.module.get_other_domains(target_domain)[0]
basic_table_info = dict(
train_cfg=os.path.basename(cfg._filename),
ckpt=ckpt,
sample_model=args.sample_model,
source_domain=source_domain,
target_domain=target_domain)
# build the dataloader
if len(metrics) == 0:
basic_table_info['num_samples'] = args.num_samples
data_loader = None
else:
basic_table_info['num_samples'] = -1
if cfg.data.get('test', None):
dataset = build_dataset(cfg.data.test)
else:
dataset = build_dataset(cfg.data.train)
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=False,
shuffle=True)
if args.online:
single_gpu_online_evaluation(model, data_loader, metrics, logger,
basic_table_info, args.batch_size)
else:
single_gpu_evaluation(model, data_loader, metrics, logger,
basic_table_info, args.batch_size,
args.samples_path)
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)
generator.eval()
mmcv.print_log('Calculating or loading eigen vectors', 'mmgen')
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 = {}
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_pix2pix():
model_cfg = dict(type='Pix2Pix',
generator=dict(type='UnetGenerator',
in_channels=3,
out_channels=3,
num_down=8,
base_channels=64,
norm_cfg=dict(type='BN'),
use_dropout=True,
init_cfg=dict(type='normal', gain=0.02)),
discriminator=dict(type='PatchDiscriminator',
in_channels=6,
base_channels=64,
num_conv=3,
norm_cfg=dict(type='BN'),
init_cfg=dict(type='normal',
gain=0.02)),
gan_loss=dict(type='GANLoss',
gan_type='vanilla',
real_label_val=1.0,
fake_label_val=0,
loss_weight=1.0),
pixel_loss=dict(type='L1Loss',
loss_weight=100.0,
reduction='mean'))
train_cfg = None
test_cfg = None
# build synthesizer
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
# test checking gan loss cannot be None
with pytest.raises(AssertionError):
bad_model_cfg = copy.deepcopy(model_cfg)
bad_model_cfg['gan_loss'] = None
_ = build_model(bad_model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)
# test attributes
assert synthesizer.__class__.__name__ == 'Pix2Pix'
assert isinstance(synthesizer.generator, UnetGenerator)
assert isinstance(synthesizer.discriminator, PatchDiscriminator)
assert isinstance(synthesizer.gan_loss, GANLoss)
assert isinstance(synthesizer.pixel_loss, L1Loss)
assert synthesizer.train_cfg is None
assert synthesizer.test_cfg is None
# prepare data
inputs = torch.rand(1, 3, 256, 256)
targets = torch.rand(1, 3, 256, 256)
data_batch = {'img_a': inputs, 'img_b': targets}
img_meta = {}
img_meta['img_a_path'] = 'img_a_path'
img_meta['img_b_path'] = 'img_b_path'
data_batch['meta'] = [img_meta]
# prepare optimizer
optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.5, 0.999))
optimizer = {
'generator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generator').parameters())),
'discriminator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminator').parameters()))
}
# test forward_dummy
with torch.no_grad():
output = synthesizer.forward_dummy(data_batch['img_a'])
assert torch.is_tensor(output)
assert output.size() == (1, 3, 256, 256)
# test forward_test
with torch.no_grad():
outputs = synthesizer(inputs, targets, [img_meta], test_mode=True)
assert torch.equal(outputs['real_a'], data_batch['img_a'])
assert torch.equal(outputs['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
# val_step
with torch.no_grad():
outputs = synthesizer.val_step(data_batch)
assert torch.equal(outputs['real_a'], data_batch['img_a'])
assert torch.equal(outputs['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
# test forward_train
outputs = synthesizer(inputs, targets, [img_meta], test_mode=False)
assert torch.equal(outputs['real_a'], data_batch['img_a'])
assert torch.equal(outputs['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
# test train_step
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in [
'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_pixel'
]:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_a'], data_batch['img_a'])
assert torch.equal(outputs['results']['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['results']['fake_b'])
assert outputs['results']['fake_b'].size() == (1, 3, 256, 256)
# test train_step and forward_test (gpu)
if torch.cuda.is_available():
synthesizer = synthesizer.cuda()
optimizer = {
'generator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generator').parameters())),
'discriminator':
obj_from_dict(
optim_cfg, torch.optim,
dict(
params=getattr(synthesizer, 'discriminator').parameters()))
}
data_batch_cuda = copy.deepcopy(data_batch)
data_batch_cuda['img_a'] = inputs.cuda()
data_batch_cuda['img_b'] = targets.cuda()
data_batch_cuda['meta'] = [DC(img_meta, cpu_only=True).data]
# forward_test
with torch.no_grad():
outputs = synthesizer(data_batch_cuda['img_a'],
data_batch_cuda['img_b'],
data_batch_cuda['meta'],
test_mode=True)
assert torch.equal(outputs['real_a'], data_batch_cuda['img_a'].cpu())
assert torch.equal(outputs['real_b'], data_batch_cuda['img_b'].cpu())
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
# val_step
with torch.no_grad():
outputs = synthesizer.val_step(data_batch_cuda)
assert torch.equal(outputs['real_a'], data_batch_cuda['img_a'].cpu())
assert torch.equal(outputs['real_b'], data_batch_cuda['img_b'].cpu())
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
# test forward_train
outputs = synthesizer(data_batch_cuda['img_a'],
data_batch_cuda['img_b'],
data_batch_cuda['meta'],
test_mode=False)
assert torch.equal(outputs['real_a'], data_batch_cuda['img_a'])
assert torch.equal(outputs['real_b'], data_batch_cuda['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
# train_step
outputs = synthesizer.train_step(data_batch_cuda, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in [
'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g',
'loss_pixel'
]:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_a'],
data_batch_cuda['img_a'].cpu())
assert torch.equal(outputs['results']['real_b'],
data_batch_cuda['img_b'].cpu())
assert torch.is_tensor(outputs['results']['fake_b'])
assert outputs['results']['fake_b'].size() == (1, 3, 256, 256)
# test disc_steps and disc_init_steps
data_batch['img_a'] = inputs.cpu()
data_batch['img_b'] = targets.cpu()
train_cfg = dict(disc_steps=2, disc_init_steps=2)
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
optimizer = {
'generator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generator').parameters())),
'discriminator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminator').parameters()))
}
# iter 0, 1
for i in range(2):
assert synthesizer.step_counter == i
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
assert outputs['log_vars'].get('loss_gan_g') is None
assert outputs['log_vars'].get('loss_pixel') is None
for v in ['loss_gan_d_fake', 'loss_gan_d_real']:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_a'], data_batch['img_a'])
assert torch.equal(outputs['results']['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['results']['fake_b'])
assert outputs['results']['fake_b'].size() == (1, 3, 256, 256)
assert synthesizer.step_counter == i + 1
# iter 2, 3, 4, 5
for i in range(2, 6):
assert synthesizer.step_counter == i
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
log_check_list = [
'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_pixel'
]
if i % 2 == 1:
assert outputs['log_vars'].get('loss_gan_g') is None
assert outputs['log_vars'].get('loss_pixel') is None
log_check_list.remove('loss_gan_g')
log_check_list.remove('loss_pixel')
for v in log_check_list:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_a'], data_batch['img_a'])
assert torch.equal(outputs['results']['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['results']['fake_b'])
assert outputs['results']['fake_b'].size() == (1, 3, 256, 256)
assert synthesizer.step_counter == i + 1
# test without pixel loss
model_cfg_ = copy.deepcopy(model_cfg)
model_cfg_.pop('pixel_loss')
synthesizer = build_model(model_cfg_, train_cfg=None, test_cfg=None)
optimizer = {
'generator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generator').parameters())),
'discriminator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminator').parameters()))
}
data_batch['img_a'] = inputs.cpu()
data_batch['img_b'] = targets.cpu()
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
assert outputs['log_vars'].get('loss_pixel') is None
for v in ['loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g']:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_a'], data_batch['img_a'])
assert torch.equal(outputs['results']['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['results']['fake_b'])
assert outputs['results']['fake_b'].size() == (1, 3, 256, 256)
# test b2a translation
data_batch['img_a'] = inputs.cpu()
data_batch['img_b'] = targets.cpu()
train_cfg = dict(direction='b2a')
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
optimizer = {
'generator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generator').parameters())),
'discriminator':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminator').parameters()))
}
assert synthesizer.step_counter == 0
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in [
'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_pixel'
]:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_a'], data_batch['img_b'])
assert torch.equal(outputs['results']['real_b'], data_batch['img_a'])
assert torch.is_tensor(outputs['results']['fake_b'])
assert outputs['results']['fake_b'].size() == (1, 3, 256, 256)
assert synthesizer.step_counter == 1
# test save image
# show input
train_cfg = None
test_cfg = dict(show_input=True)
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
with patch.object(mmcv, 'imwrite', return_value=True):
# test save path not None Assertion
with pytest.raises(AssertionError):
with torch.no_grad():
_ = synthesizer(inputs,
targets, [img_meta],
test_mode=True,
save_image=True)
# iteration is None
with torch.no_grad():
outputs = synthesizer(inputs,
targets, [img_meta],
test_mode=True,
save_image=True,
save_path='save_path')
assert torch.equal(outputs['real_a'], data_batch['img_a'])
assert torch.equal(outputs['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
assert outputs['saved_flag']
# iteration is not None
with torch.no_grad():
outputs = synthesizer(inputs,
targets, [img_meta],
test_mode=True,
save_image=True,
save_path='save_path',
iteration=1000)
assert torch.equal(outputs['real_a'], data_batch['img_a'])
assert torch.equal(outputs['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
assert outputs['saved_flag']
# not show input
train_cfg = None
test_cfg = dict(show_input=False)
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
with patch.object(mmcv, 'imwrite', return_value=True):
# test save path not None Assertion
with pytest.raises(AssertionError):
with torch.no_grad():
_ = synthesizer(inputs,
targets, [img_meta],
test_mode=True,
save_image=True)
# iteration is None
with torch.no_grad():
outputs = synthesizer(inputs,
targets, [img_meta],
test_mode=True,
save_image=True,
save_path='save_path')
assert torch.equal(outputs['real_a'], data_batch['img_a'])
assert torch.equal(outputs['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
assert outputs['saved_flag']
# iteration is not None
with torch.no_grad():
outputs = synthesizer(inputs,
targets, [img_meta],
test_mode=True,
save_image=True,
save_path='save_path',
iteration=1000)
assert torch.equal(outputs['real_a'], data_batch['img_a'])
assert torch.equal(outputs['real_b'], data_batch['img_b'])
assert torch.is_tensor(outputs['fake_b'])
assert outputs['fake_b'].size() == (1, 3, 256, 256)
assert outputs['saved_flag']
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
if args.gpu_ids is not None:
cfg.gpu_ids = args.gpu_ids[0:1]
warnings.warn('`--gpu-ids` is deprecated, please use `--gpu-id`. '
'Because we only support single GPU mode in '
'non-distributed testing. Use the first GPU '
'in `gpu_ids` now.')
else:
cfg.gpu_ids = [args.gpu_id]
# init distributed env first, since logger depends on the dist info.
if args.launcher == 'none':
distributed = False
rank = 0
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
rank, world_size = get_dist_info()
cfg.gpu_ids = range(world_size)
assert args.online or world_size == 1, (
'We only support online mode for distrbuted evaluation.')
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 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]
# check metrics for dist evaluation
if distributed and metrics:
for metric in metrics:
assert metric.name in _distributed_metrics, (
f'We only support {_distributed_metrics} for multi gpu '
f'evaluation, but receive {args.eval}.')
_ = load_checkpoint(model, args.checkpoint, map_location='cpu')
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.')
# The default loader config
loader_cfg = dict(samples_per_gpu=args.batch_size,
workers_per_gpu=cfg.data.get(
'val_workers_per_gpu', cfg.data.workers_per_gpu),
num_gpus=len(cfg.gpu_ids),
dist=distributed,
shuffle=True)
# The overall dataloader settings
loader_cfg.update({
k: v
for k, v in cfg.data.items() if k not in [
'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
'test_dataloader'
]
})
# specific config for test loader
test_loader_cfg = {**loader_cfg, **cfg.data.get('test_dataloader', {})}
data_loader = build_dataloader(dataset, **test_loader_cfg)
if args.sample_cfg is None:
args.sample_cfg = dict()
if not distributed:
model = MMDataParallel(model, device_ids=[0])
else:
find_unused_parameters = cfg.get('find_unused_parameters', False)
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
# online mode will not save samples
if args.online and len(metrics) > 0:
online_evaluation(model, data_loader, metrics, logger,
basic_table_info, args.batch_size, **args.sample_cfg)
else:
offline_evaluation(model, data_loader, metrics, logger,
basic_table_info, args.batch_size,
args.samples_path, **args.sample_cfg)
def test_default_dcgan_model_cpu(self):
sngan = build_model(self.default_config)
assert isinstance(sngan, BasicConditionalGAN)
assert not sngan.with_disc_auxiliary_loss
assert sngan.with_disc
# test forward train
with pytest.raises(NotImplementedError):
_ = sngan(None, return_loss=True)
# test forward test
imgs = sngan(None, return_loss=False, mode='sampling', num_batches=2)
assert imgs.shape == (2, 3, 32, 32)
# test train step
data = torch.randn((2, 3, 32, 32))
label = torch.randint(0, 10, (2, ))
data_input = dict(img=data, gt_label=label)
optimizer_g = torch.optim.SGD(sngan.generator.parameters(), lr=0.01)
optimizer_d = torch.optim.SGD(sngan.discriminator.parameters(),
lr=0.01)
optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
model_outputs = sngan.train_step(data_input, optim_dict)
assert 'results' in model_outputs
assert 'log_vars' in model_outputs
assert model_outputs['num_samples'] == 2
# more tests for different configs with heavy computation
# test disc_steps
config_ = deepcopy(self.default_config)
config_['train_cfg'] = dict(disc_steps=2)
sngan = build_model(config_)
model_outputs = sngan.train_step(data_input, optim_dict)
assert 'loss_disc_fake' in model_outputs['log_vars']
assert 'loss_disc_fake_g' not in model_outputs['log_vars']
assert sngan.disc_steps == 2
model_outputs = sngan.train_step(data_input,
optim_dict,
running_status=dict(iteration=1))
assert 'loss_disc_fake' in model_outputs['log_vars']
assert 'loss_disc_fake_g' in model_outputs['log_vars']
# test customized config
sagan = BasicConditionalGAN(
self.generator_cfg,
self.disc_cfg,
self.gan_loss,
self.disc_auxiliary_loss,
)
# test train step
data = torch.randn((2, 3, 32, 32))
data_input = dict(img=data, gt_label=label)
optimizer_g = torch.optim.SGD(sngan.generator.parameters(), lr=0.01)
optimizer_d = torch.optim.SGD(sngan.discriminator.parameters(),
lr=0.01)
optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
model_outputs = sagan.train_step(data_input, optim_dict)
assert 'results' in model_outputs
assert 'log_vars' in model_outputs
assert model_outputs['num_samples'] == 2
sagan = BasicConditionalGAN(
self.generator_cfg, self.disc_cfg, self.gan_loss,
dict(type='DiscShiftLoss',
loss_weight=0.5,
data_info=dict(pred='disc_pred_fake')),
dict(type='GeneratorPathRegularizer'))
assert isinstance(sagan.disc_auxiliary_losses, nn.ModuleList)
assert isinstance(sagan.gen_auxiliary_losses, nn.ModuleList)
sagan = BasicConditionalGAN(
self.generator_cfg, self.disc_cfg, self.gan_loss,
dict(type='DiscShiftLoss',
loss_weight=0.5,
data_info=dict(pred='disc_pred_fake')),
[dict(type='GeneratorPathRegularizer')])
assert isinstance(sagan.disc_auxiliary_losses, nn.ModuleList)
assert isinstance(sagan.gen_auxiliary_losses, nn.ModuleList)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
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:
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)
assert isinstance(model, _supported_model)
# 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()
_ = 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)
# build the dataloader
if len(metrics) == 0:
basic_table_info['num_samples'] = args.num_samples
data_loader = None
else:
basic_table_info['num_samples'] = -1
if cfg.data.get('test', None):
dataset = build_dataset(cfg.data.test)
else:
dataset = build_dataset(cfg.data.train)
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=False,
shuffle=True)
# decide samples path
samples_path = args.samples_path
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)
# select key to fetch fake images
fake_key = 'fake_b'
if isinstance(model.module, CycleGAN):
fake_key = 'fake_b' if model.module.test_direction == 'a2b' else \
'fake_a'
# if no images, `num_exist` should be zero
for begin in range(num_exist, num_needed, args.batch_size):
end = min(begin + args.batch_size, max_num_images)
# for translation model, we feed them images from dataloader
data_loader_iter = iter(data_loader)
data_batch = next(data_loader_iter)
output_dict = model(test_mode=True, **data_batch)
fakes = output_dict[fake_key]
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, args.batch_size)
# select key to fetch real images
if isinstance(model.module, CycleGAN):
real_key = 'img_b' if model.module.test_direction == 'a2b' else 'img_a'
if model.module.direction == 'b2a':
real_key = 'img_a' if real_key == 'img_b' else 'img_b'
if isinstance(model.module, Pix2Pix):
real_key = 'img_b' if model.module.direction == 'a2b' else 'img_a'
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_key]
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 test_pix2pix():
# model settings
model_cfg = dict(
type='Pix2Pix',
generator=dict(
type='UnetGenerator',
in_channels=3,
out_channels=3,
num_down=8,
base_channels=64,
norm_cfg=dict(type='BN'),
use_dropout=True,
init_cfg=dict(type='normal', gain=0.02)),
discriminator=dict(
type='PatchDiscriminator',
in_channels=6,
base_channels=64,
num_conv=3,
norm_cfg=dict(type='BN'),
init_cfg=dict(type='normal', gain=0.02)),
gan_loss=dict(
type='GANLoss',
gan_type='vanilla',
real_label_val=1.0,
fake_label_val=0.0,
loss_weight=1.0),
default_domain='photo',
reachable_domains=['photo'],
related_domains=['photo', 'mask'],
gen_auxiliary_loss=dict(
type='L1Loss',
loss_weight=100.0,
data_info=dict(pred='fake_photo', target='real_photo'),
reduction='mean'))
train_cfg = None
test_cfg = None
# build synthesizer
synthesizer = build_model(
model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)
# test attributes
assert synthesizer.__class__.__name__ == 'Pix2Pix'
assert isinstance(synthesizer.generators['photo'], UnetGenerator)
assert isinstance(synthesizer.discriminators['photo'], PatchDiscriminator)
assert isinstance(synthesizer.gan_loss, GANLoss)
assert isinstance(synthesizer.gen_auxiliary_losses[0], L1Loss)
assert synthesizer.test_cfg is None
# prepare data
img_mask = torch.rand(1, 3, 256, 256)
img_photo = torch.rand(1, 3, 256, 256)
data_batch = {'img_mask': img_mask, 'img_photo': img_photo}
# prepare optimizer
optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.5, 0.999))
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminators').parameters()))
}
# test forward_test
domain = 'photo'
with torch.no_grad():
outputs = synthesizer(img_mask, target_domain=domain, test_mode=True)
assert torch.equal(outputs['source'], data_batch['img_mask'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 256, 256)
# test forward_train
outputs = synthesizer(img_mask, target_domain=domain, test_mode=False)
assert torch.equal(outputs['source'], data_batch['img_mask'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 256, 256)
# test train_step
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in ['loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_l1']:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_mask'], data_batch['img_mask'])
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
# test cuda
if torch.cuda.is_available():
synthesizer = synthesizer.cuda()
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(
params=getattr(synthesizer,
'discriminators').parameters()))
}
data_batch_cuda = copy.deepcopy(data_batch)
data_batch_cuda['img_mask'] = img_mask.cuda()
data_batch_cuda['img_photo'] = img_photo.cuda()
# forward_test
with torch.no_grad():
outputs = synthesizer(
data_batch_cuda['img_mask'],
target_domain=domain,
test_mode=True)
assert torch.equal(outputs['source'],
data_batch_cuda['img_mask'].cpu())
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 256, 256)
# test forward_train
outputs = synthesizer(
data_batch_cuda['img_mask'], target_domain=domain, test_mode=False)
assert torch.equal(outputs['source'], data_batch_cuda['img_mask'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 256, 256)
# train_step
outputs = synthesizer.train_step(data_batch_cuda, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in [
'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_l1'
]:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_mask'],
data_batch_cuda['img_mask'].cpu())
assert torch.equal(outputs['results']['real_photo'],
data_batch_cuda['img_photo'].cpu())
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
# test disc_steps and disc_init_steps
data_batch['img_mask'] = img_mask.cpu()
data_batch['img_photo'] = img_photo.cpu()
train_cfg = dict(disc_steps=2, disc_init_steps=2)
synthesizer = build_model(
model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminators').parameters()))
}
# iter 0, 1
for i in range(2):
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
assert outputs['log_vars'].get('loss_gan_g') is None
assert outputs['log_vars'].get('loss_l1') is None
for v in ['loss_gan_d_fake', 'loss_gan_d_real']:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_mask'],
data_batch['img_mask'])
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
assert synthesizer.iteration == i + 1
# iter 2, 3, 4, 5
for i in range(2, 6):
assert synthesizer.iteration == i
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
log_check_list = [
'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_l1'
]
if i % 2 == 1:
assert outputs['log_vars'].get('loss_gan_g') is None
assert outputs['log_vars'].get('loss_pixel') is None
log_check_list.remove('loss_gan_g')
log_check_list.remove('loss_l1')
for v in log_check_list:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_mask'],
data_batch['img_mask'])
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
assert synthesizer.iteration == i + 1
# test without pixel loss
model_cfg_ = copy.deepcopy(model_cfg)
model_cfg_.pop('gen_auxiliary_loss')
synthesizer = build_model(model_cfg_, train_cfg=None, test_cfg=None)
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminators').parameters()))
}
data_batch['img_mask'] = img_mask.cpu()
data_batch['img_photo'] = img_photo.cpu()
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
assert outputs['log_vars'].get('loss_pixel') is None
for v in ['loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g']:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_mask'], data_batch['img_mask'])
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.cfg_options is not None:
cfg.merge_from_dict(args.cfg_options)
# import modules from string list.
if cfg.get('custom_imports', None):
from mmcv.utils import import_modules_from_strings
import_modules_from_strings(**cfg['custom_imports'])
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# work_dir is determined in this priority: CLI > segment in file > filename
if args.work_dir is not None:
# update configs according to CLI args if args.work_dir is not None
cfg.work_dir = args.work_dir
elif cfg.get('work_dir', None) is None:
# use config filename as default work_dir if cfg.work_dir is None
cfg.work_dir = osp.join('./work_dirs',
osp.splitext(osp.basename(args.config))[0])
if args.resume_from is not None:
cfg.resume_from = args.resume_from
if args.gpu_ids is not None:
cfg.gpu_ids = args.gpu_ids
else:
cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus)
# 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)
# re-set gpu_ids with distributed training mode
_, world_size = get_dist_info()
cfg.gpu_ids = range(world_size)
# create work_dir
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
# dump config
cfg.dump(osp.join(cfg.work_dir, osp.basename(args.config)))
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(cfg.work_dir, f'{timestamp}.log')
logger = get_root_logger(log_file=log_file, log_level=cfg.log_level)
# init the meta dict to record some important information such as
# environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
meta['env_info'] = env_info
meta['config'] = cfg.pretty_text
# log some basic info
logger.info(f'Distributed training: {distributed}')
logger.info(f'Config:\n{cfg.pretty_text}')
# set random seeds
if args.seed is not None:
logger.info(f'Set random seed to {args.seed}, '
f'deterministic: {args.deterministic}')
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
meta['exp_name'] = osp.basename(args.config)
model = build_model(cfg.model,
train_cfg=cfg.train_cfg,
test_cfg=cfg.test_cfg)
datasets = [build_dataset(cfg.data.train)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_dataset(val_dataset))
if cfg.checkpoint_config is not None:
# save mmgen version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(mmgen_version=__version__ +
get_git_hash()[:7])
train_model(model,
datasets,
cfg,
distributed=distributed,
validate=(not args.no_validate),
timestamp=timestamp,
meta=meta)
def test_ada_stylegan2_model_cpu(self):
synthesis_cfg = {
'type': 'SynthesisNetwork',
'channel_base': 1024,
'channel_max': 16,
'magnitude_ema_beta': 0.999
}
aug_kwargs = {
'xflip': 1,
'rotate90': 1,
'xint': 1,
'scale': 1,
'rotate': 1,
'aniso': 1,
'xfrac': 1,
'brightness': 1,
'contrast': 1,
'lumaflip': 1,
'hue': 1,
'saturation': 1
}
default_config = dict(type='StaticUnconditionalGAN',
generator=dict(type='StyleGANv3Generator',
out_size=8,
style_channels=8,
img_channels=3,
rgb2bgr=True,
synthesis_cfg=synthesis_cfg),
discriminator=dict(
type='ADAStyleGAN2Discriminator',
in_size=8,
input_bgr2rgb=True,
data_aug=dict(type='ADAAug',
update_interval=2,
aug_pipeline=aug_kwargs,
ada_kimg=100)),
gan_loss=dict(type='GANLoss',
gan_type='wgan-logistic-ns'))
s3gan = build_model(default_config)
assert isinstance(s3gan, StaticUnconditionalGAN)
assert not s3gan.with_disc_auxiliary_loss
assert s3gan.with_disc
# test forward train
with pytest.raises(NotImplementedError):
_ = s3gan(None, return_loss=True)
# test forward test
imgs = s3gan(None, return_loss=False, mode='sampling', num_batches=2)
assert imgs.shape == (2, 3, 8, 8)
# test train step
data = torch.randn((2, 3, 8, 8))
data_input = dict(real_img=data)
optimizer_g = torch.optim.SGD(s3gan.generator.parameters(), lr=0.01)
optimizer_d = torch.optim.SGD(s3gan.discriminator.parameters(),
lr=0.01)
optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
_ = s3gan.train_step(data_input,
optim_dict,
running_status=dict(iteration=1))
s3gan.discriminator.ada_aug.aug_pipeline.p.dtype == torch.float32
def test_cyclegan():
model_cfg = dict(type='CycleGAN',
default_domain='photo',
reachable_domains=['photo', 'mask'],
related_domains=['photo', 'mask'],
generator=dict(type='ResnetGenerator',
in_channels=3,
out_channels=3,
base_channels=64,
norm_cfg=dict(type='IN'),
use_dropout=False,
num_blocks=9,
padding_mode='reflect',
init_cfg=dict(type='normal', gain=0.02)),
discriminator=dict(type='PatchDiscriminator',
in_channels=3,
base_channels=64,
num_conv=3,
norm_cfg=dict(type='IN'),
init_cfg=dict(type='normal',
gain=0.02)),
gan_loss=dict(type='GANLoss',
gan_type='lsgan',
real_label_val=1.0,
fake_label_val=0.0,
loss_weight=1.0),
gen_auxiliary_loss=[
dict(type='L1Loss',
loss_weight=10.0,
data_info=dict(pred='cycle_photo',
target='real_photo'),
reduction='mean'),
dict(type='L1Loss',
loss_weight=10.0,
data_info=dict(
pred='cycle_mask',
target='real_mask',
),
reduction='mean'),
dict(type='L1Loss',
loss_weight=0.5,
data_info=dict(pred='identity_photo',
target='real_photo'),
reduction='mean'),
dict(type='L1Loss',
loss_weight=0.5,
data_info=dict(pred='identity_mask',
target='real_mask'),
reduction='mean')
])
train_cfg = None
test_cfg = None
# build synthesizer
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
# test attributes
assert synthesizer.__class__.__name__ == 'CycleGAN'
assert isinstance(synthesizer.generators['photo'], ResnetGenerator)
assert isinstance(synthesizer.generators['mask'], ResnetGenerator)
assert isinstance(synthesizer.discriminators['photo'], PatchDiscriminator)
assert isinstance(synthesizer.discriminators['mask'], PatchDiscriminator)
assert isinstance(synthesizer.gan_loss, GANLoss)
for loss_module in synthesizer.gen_auxiliary_losses:
assert isinstance(loss_module, L1Loss)
# prepare data
inputs = torch.rand(1, 3, 64, 64)
targets = torch.rand(1, 3, 64, 64)
data_batch = {'img_mask': inputs, 'img_photo': targets}
# prepare optimizer
optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.5, 0.999))
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminators').parameters()))
}
# test forward_test
with torch.no_grad():
outputs = synthesizer(inputs, target_domain='photo', test_mode=True)
assert torch.equal(outputs['source'], data_batch['img_mask'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
with torch.no_grad():
outputs = synthesizer(targets, target_domain='mask', test_mode=True)
assert torch.equal(outputs['source'], data_batch['img_photo'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
# test forward_train
with torch.no_grad():
outputs = synthesizer(inputs, target_domain='photo', test_mode=True)
assert torch.equal(outputs['source'], data_batch['img_mask'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
with torch.no_grad():
outputs = synthesizer(targets, target_domain='mask', test_mode=True)
assert torch.equal(outputs['source'], data_batch['img_photo'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
# test train_step
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in [
'loss_gan_d_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
'loss_gan_g_photo', 'loss_l1'
]:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.equal(outputs['results']['real_mask'], data_batch['img_mask'])
assert torch.is_tensor(outputs['results']['fake_mask'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
# test train_step and forward_test (gpu)
if torch.cuda.is_available():
synthesizer = synthesizer.cuda()
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer,
'discriminators').parameters()))
}
data_batch_cuda = copy.deepcopy(data_batch)
data_batch_cuda['img_mask'] = inputs.cuda()
data_batch_cuda['img_photo'] = targets.cuda()
# forward_test
with torch.no_grad():
outputs = synthesizer(data_batch_cuda['img_mask'],
target_domain='photo',
test_mode=True)
assert torch.equal(outputs['source'],
data_batch_cuda['img_mask'].cpu())
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
with torch.no_grad():
outputs = synthesizer(data_batch_cuda['img_photo'],
target_domain='mask',
test_mode=True)
assert torch.equal(outputs['source'],
data_batch_cuda['img_photo'].cpu())
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
# test forward_train
with torch.no_grad():
outputs = synthesizer(data_batch_cuda['img_mask'],
target_domain='photo',
test_mode=False)
assert torch.equal(outputs['source'], data_batch_cuda['img_mask'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
with torch.no_grad():
outputs = synthesizer(data_batch_cuda['img_photo'],
target_domain='mask',
test_mode=False)
assert torch.equal(outputs['source'], data_batch_cuda['img_photo'])
assert torch.is_tensor(outputs['target'])
assert outputs['target'].size() == (1, 3, 64, 64)
# train_step
outputs = synthesizer.train_step(data_batch_cuda, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
print(outputs['log_vars'].keys())
assert isinstance(outputs['results'], dict)
for v in [
'loss_gan_d_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
'loss_gan_g_photo', 'loss_l1'
]:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_photo'],
data_batch_cuda['img_photo'].cpu())
assert torch.equal(outputs['results']['real_mask'],
data_batch_cuda['img_mask'].cpu())
assert torch.is_tensor(outputs['results']['fake_mask'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
# test disc_steps and disc_init_steps
data_batch['img_mask'] = inputs.cpu()
data_batch['img_photo'] = targets.cpu()
train_cfg = dict(disc_steps=2, disc_init_steps=2)
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminators').parameters()))
}
# iter 0, 1
for i in range(2):
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in ['loss_gan_g_mask', 'loss_gan_g_photo', 'loss_l1']:
assert outputs['log_vars'].get(v) is None
assert isinstance(outputs['log_vars']['loss_gan_d_mask'], float)
assert isinstance(outputs['log_vars']['loss_gan_d_photo'], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.equal(outputs['results']['real_mask'],
data_batch['img_mask'])
assert torch.is_tensor(outputs['results']['fake_mask'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
assert synthesizer.iteration == i + 1
# iter 2, 3, 4, 5
for i in range(2, 6):
assert synthesizer.iteration == i
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
log_check_list = [
'loss_gan_d_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
'loss_gan_g_photo', 'loss_l1'
]
if i % 2 == 1:
log_None_list = ['loss_gan_g_mask', 'loss_gan_g_photo', 'loss_l1']
for v in log_None_list:
assert outputs['log_vars'].get(v) is None
log_check_list.remove(v)
for v in log_check_list:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_mask'],
data_batch['img_mask'])
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.is_tensor(outputs['results']['fake_mask'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
assert synthesizer.iteration == i + 1
# test GAN image buffer size = 0
data_batch['img_mask'] = inputs.cpu()
data_batch['img_photo'] = targets.cpu()
train_cfg = dict(buffer_size=0)
synthesizer = build_model(model_cfg,
train_cfg=train_cfg,
test_cfg=test_cfg)
optimizer = {
'generators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'generators').parameters())),
'discriminators':
obj_from_dict(
optim_cfg, torch.optim,
dict(params=getattr(synthesizer, 'discriminators').parameters()))
}
outputs = synthesizer.train_step(data_batch, optimizer)
assert isinstance(outputs, dict)
assert isinstance(outputs['log_vars'], dict)
assert isinstance(outputs['results'], dict)
for v in [
'loss_gan_d_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
'loss_gan_g_photo', 'loss_l1'
]:
assert isinstance(outputs['log_vars'][v], float)
assert outputs['num_samples'] == 1
assert torch.equal(outputs['results']['real_mask'], data_batch['img_mask'])
assert torch.equal(outputs['results']['real_photo'],
data_batch['img_photo'])
assert torch.is_tensor(outputs['results']['fake_mask'])
assert torch.is_tensor(outputs['results']['fake_photo'])
assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
assert synthesizer.iteration == 1
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)
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)
