def load_StudioGAN_ckpts(ckpt_dir, load_best, Gen, Dis, g_optimizer, d_optimizer, run_name, apply_g_ema, Gen_ema, ema,
is_train, RUN, logger, global_rank, device, cfg_file):
when = "best" if load_best is True else "current"
Gen_ckpt_path = glob.glob(join(ckpt_dir, "model=G-{when}-weights-step*.pth".format(when=when)))[0]
Dis_ckpt_path = glob.glob(join(ckpt_dir, "model=D-{when}-weights-step*.pth".format(when=when)))[0]
prev_run_name = torch.load(Dis_ckpt_path, map_location=lambda storage, loc: storage)["run_name"]
is_freezeD = True if RUN.freezeD > -1 else False
load_ckpt(model=Gen,
optimizer=g_optimizer,
ckpt_path=Gen_ckpt_path,
load_model=True,
load_opt=False if prev_run_name in blacklist or is_freezeD or not is_train else True,
load_misc=False,
is_freezeD=is_freezeD)
seed, prev_run_name, step, epoch, topk, aa_p, best_step, best_fid, best_ckpt_path =\
load_ckpt(model=Dis,
optimizer=d_optimizer,
ckpt_path=Dis_ckpt_path,
load_model=True,
load_opt=False if prev_run_name in blacklist or is_freezeD or not is_train else True,
load_misc=True,
is_freezeD=is_freezeD)
if not is_train:
prev_run_name = cfg_file[cfg_file.rindex("/")+1:cfg_file.index(".yaml")]+prev_run_name[prev_run_name.index("-train"):]
if apply_g_ema:
Gen_ema_ckpt_path = glob.glob(join(ckpt_dir, "model=G_ema-{when}-weights-step*.pth".format(when=when)))[0]
load_ckpt(model=Gen_ema,
optimizer=None,
ckpt_path=Gen_ema_ckpt_path,
load_model=True,
load_opt=False,
load_misc=False,
is_freezeD=is_freezeD)
ema.source, ema.target = Gen, Gen_ema
if is_train and RUN.seed != seed:
RUN.seed = seed + global_rank
misc.fix_seed(RUN.seed)
if device == 0:
if not is_freezeD:
logger = log.make_logger(RUN.save_dir, prev_run_name, None)
logger.info("Generator checkpoint is {}".format(Gen_ckpt_path))
if apply_g_ema:
logger.info("EMA_Generator checkpoint is {}".format(Gen_ema_ckpt_path))
logger.info("Discriminator checkpoint is {}".format(Dis_ckpt_path))
if is_freezeD:
prev_run_name, step, epoch, topk, aa_p, best_step, best_fid, best_ckpt_path =\
run_name, 0, 0, "initialize", None, 0, None, None
return prev_run_name, step, epoch, topk, aa_p, best_step, best_fid, best_ckpt_path, logger
def prepare_train_eval(rank, world_size, run_name, train_config, model_config,
hdf5_path_train):
cfgs = dict2clsattr(train_config, model_config)
prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path, mu, sigma, inception_model = None, 0, 0, None, None, None, None, None
if cfgs.distributed_data_parallel:
print("Use GPU: {} for training.".format(rank))
setup(rank, world_size)
torch.cuda.set_device(rank)
writer = SummaryWriter(
log_dir=join('./logs', run_name)) if rank == 0 else None
if rank == 0:
logger = make_logger(run_name, None)
logger.info('Run name : {run_name}'.format(run_name=run_name))
logger.info(train_config)
logger.info(model_config)
else:
logger = None
##### load dataset #####
if rank == 0: logger.info('Load train datasets...')
train_dataset = LoadDataset(cfgs.dataset_name,
cfgs.data_path,
train=True,
download=True,
resize_size=cfgs.img_size,
hdf5_path=hdf5_path_train,
random_flip=cfgs.random_flip_preprocessing)
if cfgs.reduce_train_dataset < 1.0:
num_train = int(cfgs.reduce_train_dataset * len(train_dataset))
train_dataset, _ = torch.utils.data.random_split(
train_dataset,
[num_train, len(train_dataset) - num_train])
if rank == 0:
logger.info('Train dataset size : {dataset_size}'.format(
dataset_size=len(train_dataset)))
if rank == 0:
logger.info('Load {mode} datasets...'.format(mode=cfgs.eval_type))
eval_mode = True if cfgs.eval_type == 'train' else False
eval_dataset = LoadDataset(cfgs.dataset_name,
cfgs.data_path,
train=eval_mode,
download=True,
resize_size=cfgs.img_size,
hdf5_path=None,
random_flip=False)
if rank == 0:
logger.info('Eval dataset size : {dataset_size}'.format(
dataset_size=len(eval_dataset)))
if cfgs.distributed_data_parallel:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
cfgs.batch_size = cfgs.batch_size // world_size
else:
train_sampler = None
train_dataloader = DataLoader(train_dataset,
batch_size=cfgs.batch_size,
shuffle=(train_sampler is None),
pin_memory=True,
num_workers=cfgs.num_workers,
sampler=train_sampler,
drop_last=True)
eval_dataloader = DataLoader(eval_dataset,
batch_size=cfgs.batch_size,
shuffle=False,
pin_memory=True,
num_workers=cfgs.num_workers,
drop_last=False)
##### build model #####
if rank == 0: logger.info('Build model...')
module = __import__(
'models.{architecture}'.format(architecture=cfgs.architecture),
fromlist=['something'])
if rank == 0:
logger.info('Modules are located on models.{architecture}.'.format(
architecture=cfgs.architecture))
Gen = module.Generator(cfgs.z_dim, cfgs.shared_dim, cfgs.img_size,
cfgs.g_conv_dim, cfgs.g_spectral_norm,
cfgs.attention, cfgs.attention_after_nth_gen_block,
cfgs.activation_fn, cfgs.conditional_strategy,
cfgs.num_classes, cfgs.g_init, cfgs.G_depth,
cfgs.mixed_precision).to(rank)
Dis = module.Discriminator(
cfgs.img_size, cfgs.d_conv_dim, cfgs.d_spectral_norm, cfgs.attention,
cfgs.attention_after_nth_dis_block, cfgs.activation_fn,
cfgs.conditional_strategy, cfgs.hypersphere_dim, cfgs.num_classes,
cfgs.nonlinear_embed, cfgs.normalize_embed, cfgs.d_init, cfgs.D_depth,
cfgs.mixed_precision).to(rank)
if cfgs.ema:
if rank == 0:
logger.info('Prepare EMA for G with decay of {}.'.format(
cfgs.ema_decay))
Gen_copy = module.Generator(
cfgs.z_dim,
cfgs.shared_dim,
cfgs.img_size,
cfgs.g_conv_dim,
cfgs.g_spectral_norm,
cfgs.attention,
cfgs.attention_after_nth_gen_block,
cfgs.activation_fn,
cfgs.conditional_strategy,
cfgs.num_classes,
initialize=False,
G_depth=cfgs.G_depth,
mixed_precision=cfgs.mixed_precision).to(rank)
Gen_ema = ema(Gen, Gen_copy, cfgs.ema_decay, cfgs.ema_start)
else:
Gen_copy, Gen_ema = None, None
if rank == 0: logger.info(count_parameters(Gen))
if rank == 0: logger.info(Gen)
if rank == 0: logger.info(count_parameters(Dis))
if rank == 0: logger.info(Dis)
### define loss functions and optimizers
G_loss = {
'vanilla': loss_dcgan_gen,
'least_square': loss_lsgan_gen,
'hinge': loss_hinge_gen,
'wasserstein': loss_wgan_gen
}
D_loss = {
'vanilla': loss_dcgan_dis,
'least_square': loss_lsgan_dis,
'hinge': loss_hinge_dis,
'wasserstein': loss_wgan_dis
}
if cfgs.optimizer == "SGD":
G_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
Gen.parameters()),
cfgs.g_lr,
momentum=cfgs.momentum,
nesterov=cfgs.nesterov)
D_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
Dis.parameters()),
cfgs.d_lr,
momentum=cfgs.momentum,
nesterov=cfgs.nesterov)
elif cfgs.optimizer == "RMSprop":
G_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
Gen.parameters()),
cfgs.g_lr,
momentum=cfgs.momentum,
alpha=cfgs.alpha)
D_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
Dis.parameters()),
cfgs.d_lr,
momentum=cfgs.momentum,
alpha=cfgs.alpha)
elif cfgs.optimizer == "Adam":
G_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
Gen.parameters()),
cfgs.g_lr, [cfgs.beta1, cfgs.beta2],
eps=1e-6)
D_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
Dis.parameters()),
cfgs.d_lr, [cfgs.beta1, cfgs.beta2],
eps=1e-6)
else:
raise NotImplementedError
if cfgs.LARS_optimizer:
G_optimizer = LARS(optimizer=G_optimizer, eps=1e-8, trust_coef=0.001)
D_optimizer = LARS(optimizer=D_optimizer, eps=1e-8, trust_coef=0.001)
##### load checkpoints if needed #####
if cfgs.checkpoint_folder is None:
checkpoint_dir = make_checkpoint_dir(cfgs.checkpoint_folder, run_name)
else:
when = "current" if cfgs.load_current is True else "best"
if not exists(abspath(cfgs.checkpoint_folder)):
raise NotADirectoryError
checkpoint_dir = make_checkpoint_dir(cfgs.checkpoint_folder, run_name)
g_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=G-{when}-weights-step*.pth".format(when=when)))[0]
d_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=D-{when}-weights-step*.pth".format(when=when)))[0]
Gen, G_optimizer, trained_seed, run_name, step, prev_ada_p = load_checkpoint(
Gen, G_optimizer, g_checkpoint_dir)
Dis, D_optimizer, trained_seed, run_name, step, prev_ada_p, best_step, best_fid, best_fid_checkpoint_path =\
load_checkpoint(Dis, D_optimizer, d_checkpoint_dir, metric=True)
if rank == 0: logger = make_logger(run_name, None)
if cfgs.ema:
g_ema_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=G_ema-{when}-weights-step*.pth".format(
when=when)))[0]
Gen_copy = load_checkpoint(Gen_copy,
None,
g_ema_checkpoint_dir,
ema=True)
Gen_ema.source, Gen_ema.target = Gen, Gen_copy
writer = SummaryWriter(
log_dir=join('./logs', run_name)) if rank == 0 else None
if cfgs.train_configs['train']:
assert cfgs.seed == trained_seed, "Seed for sampling random numbers should be same!"
if rank == 0:
logger.info('Generator checkpoint is {}'.format(g_checkpoint_dir))
if rank == 0:
logger.info(
'Discriminator checkpoint is {}'.format(d_checkpoint_dir))
if cfgs.freeze_layers > -1:
prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path = None, 0, 0, None, None
##### wrap models with DP and convert BN to Sync BN #####
if world_size > 1:
if cfgs.distributed_data_parallel:
if cfgs.synchronized_bn:
process_group = torch.distributed.new_group(
[w for w in range(world_size)])
Gen = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
Gen, process_group)
Dis = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
Dis, process_group)
if cfgs.ema:
Gen_copy = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
Gen_copy, process_group)
Gen = DDP(Gen,
device_ids=[rank],
broadcast_buffers=False,
find_unused_parameters=True)
Dis = DDP(Dis,
device_ids=[rank],
broadcast_buffers=False,
find_unused_parameters=True)
if cfgs.ema:
Gen_copy = DDP(Gen_copy,
device_ids=[rank],
broadcast_buffers=False,
find_unused_parameters=True)
else:
Gen = DataParallel(Gen, output_device=rank)
Dis = DataParallel(Dis, output_device=rank)
if cfgs.ema:
Gen_copy = DataParallel(Gen_copy, output_device=rank)
if cfgs.synchronized_bn:
Gen = convert_model(Gen).to(rank)
Dis = convert_model(Dis).to(rank)
if cfgs.ema:
Gen_copy = convert_model(Gen_copy).to(rank)
##### load the inception network and prepare first/secend moments for calculating FID #####
if cfgs.eval:
inception_model = InceptionV3().to(rank)
if world_size > 1 and cfgs.distributed_data_parallel:
toggle_grad(inception_model, on=True)
inception_model = DDP(inception_model,
device_ids=[rank],
broadcast_buffers=False,
find_unused_parameters=True)
elif world_size > 1 and cfgs.distributed_data_parallel is False:
inception_model = DataParallel(inception_model, output_device=rank)
else:
pass
mu, sigma = prepare_inception_moments(dataloader=eval_dataloader,
generator=Gen,
eval_mode=cfgs.eval_type,
inception_model=inception_model,
splits=1,
run_name=run_name,
logger=logger,
device=rank)
worker = make_worker(
cfgs=cfgs,
run_name=run_name,
best_step=best_step,
logger=logger,
writer=writer,
n_gpus=world_size,
gen_model=Gen,
dis_model=Dis,
inception_model=inception_model,
Gen_copy=Gen_copy,
Gen_ema=Gen_ema,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
train_dataloader=train_dataloader,
eval_dataloader=eval_dataloader,
G_optimizer=G_optimizer,
D_optimizer=D_optimizer,
G_loss=G_loss[cfgs.adv_loss],
D_loss=D_loss[cfgs.adv_loss],
prev_ada_p=prev_ada_p,
rank=rank,
checkpoint_dir=checkpoint_dir,
mu=mu,
sigma=sigma,
best_fid=best_fid,
best_fid_checkpoint_path=best_fid_checkpoint_path,
)
if cfgs.train_configs['train']:
step = worker.train(current_step=step, total_step=cfgs.total_step)
if cfgs.eval:
is_save = worker.evaluation(
step=step,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
if cfgs.save_images:
worker.save_images(is_generate=True,
png=True,
npz=True,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
if cfgs.image_visualization:
worker.run_image_visualization(
nrow=cfgs.nrow,
ncol=cfgs.ncol,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
if cfgs.k_nearest_neighbor:
worker.run_nearest_neighbor(
nrow=cfgs.nrow,
ncol=cfgs.ncol,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
if cfgs.interpolation:
assert cfgs.architecture in [
"big_resnet", "biggan_deep"
], "StudioGAN does not support interpolation analysis except for biggan and biggan_deep."
worker.run_linear_interpolation(
nrow=cfgs.nrow,
ncol=cfgs.ncol,
fix_z=True,
fix_y=False,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
worker.run_linear_interpolation(
nrow=cfgs.nrow,
ncol=cfgs.ncol,
fix_z=False,
fix_y=True,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
if cfgs.frequency_analysis:
worker.run_frequency_analysis(
num_images=len(train_dataset) // cfgs.num_classes,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
if cfgs.tsne_analysis:
worker.run_tsne(dataloader=eval_dataloader,
standing_statistics=cfgs.standing_statistics,
standing_step=cfgs.standing_step)
def load_frameowrk(
seed, disable_debugging_API, num_workers, config_path,
checkpoint_folder, reduce_train_dataset, standing_statistics,
standing_step, freeze_layers, load_current, eval_type, dataset_name,
num_classes, img_size, data_path, architecture, conditional_strategy,
hypersphere_dim, nonlinear_embed, normalize_embed, g_spectral_norm,
d_spectral_norm, activation_fn, attention,
attention_after_nth_gen_block, attention_after_nth_dis_block, z_dim,
shared_dim, g_conv_dim, d_conv_dim, G_depth, D_depth, optimizer,
batch_size, d_lr, g_lr, momentum, nesterov, alpha, beta1, beta2,
total_step, adv_loss, cr, g_init, d_init, random_flip_preprocessing,
prior, truncated_factor, ema, ema_decay, ema_start, synchronized_bn,
mixed_precision, hdf5_path_train, train_config, model_config, **_):
if seed == 0:
cudnn.benchmark = True
cudnn.deterministic = False
else:
fix_all_seed(seed)
cudnn.benchmark = False
cudnn.deterministic = True
if disable_debugging_API:
torch.autograd.set_detect_anomaly(False)
n_gpus = torch.cuda.device_count()
default_device = torch.cuda.current_device()
check_flag_0(batch_size, n_gpus, standing_statistics, ema, freeze_layers,
checkpoint_folder)
assert batch_size % n_gpus == 0, "batch_size should be divided by the number of gpus "
if n_gpus == 1:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path = None, 0, 0, None, None
standing_step = standing_step if standing_statistics is True else batch_size
run_name = make_run_name(RUN_NAME_FORMAT,
framework=config_path.split('/')[-1][:-5],
phase='train')
logger = make_logger(run_name, None)
writer = SummaryWriter(log_dir=join('./logs', run_name))
logger.info('Run name : {run_name}'.format(run_name=run_name))
logger.info(train_config)
logger.info(model_config)
logger.info('Loading train datasets...')
train_dataset = LoadDataset(dataset_name,
data_path,
train=True,
download=True,
resize_size=img_size,
hdf5_path=hdf5_path_train,
random_flip=random_flip_preprocessing)
if reduce_train_dataset < 1.0:
num_train = int(reduce_train_dataset * len(train_dataset))
train_dataset, _ = torch.utils.data.random_split(
train_dataset,
[num_train, len(train_dataset) - num_train])
logger.info('Train dataset size : {dataset_size}'.format(
dataset_size=len(train_dataset)))
logger.info('Loading {mode} datasets...'.format(mode=eval_type))
eval_mode = True if eval_type == 'train' else False
eval_dataset = LoadDataset(dataset_name,
data_path,
train=eval_mode,
download=True,
resize_size=img_size,
hdf5_path=None,
random_flip=False)
logger.info('Eval dataset size : {dataset_size}'.format(
dataset_size=len(eval_dataset)))
logger.info('Building model...')
if architecture == "dcgan":
assert img_size == 32, "Sry, StudioGAN does not support dcgan models for generation of images larger than 32 resolution."
module = __import__(
'models.{architecture}'.format(architecture=architecture),
fromlist=['something'])
logger.info('Modules are located on models.{architecture}'.format(
architecture=architecture))
Gen = module.Generator(z_dim, shared_dim, img_size, g_conv_dim,
g_spectral_norm, attention,
attention_after_nth_gen_block, activation_fn,
conditional_strategy, num_classes, g_init, G_depth,
mixed_precision).to(default_device)
Dis = module.Discriminator(img_size, d_conv_dim, d_spectral_norm,
attention, attention_after_nth_dis_block,
activation_fn, conditional_strategy,
hypersphere_dim, num_classes, nonlinear_embed,
normalize_embed, d_init, D_depth,
mixed_precision).to(default_device)
if ema:
print('Preparing EMA for G with decay of {}'.format(ema_decay))
Gen_copy = module.Generator(
z_dim,
shared_dim,
img_size,
g_conv_dim,
g_spectral_norm,
attention,
attention_after_nth_gen_block,
activation_fn,
conditional_strategy,
num_classes,
initialize=False,
G_depth=G_depth,
mixed_precision=mixed_precision).to(default_device)
Gen_ema = ema_(Gen, Gen_copy, ema_decay, ema_start)
else:
Gen_copy, Gen_ema = None, None
logger.info(count_parameters(Gen))
logger.info(Gen)
logger.info(count_parameters(Dis))
logger.info(Dis)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=num_workers,
drop_last=True)
eval_dataloader = DataLoader(eval_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=num_workers,
drop_last=False)
G_loss = {
'vanilla': loss_dcgan_gen,
'least_square': loss_lsgan_gen,
'hinge': loss_hinge_gen,
'wasserstein': loss_wgan_gen
}
D_loss = {
'vanilla': loss_dcgan_dis,
'least_square': loss_lsgan_dis,
'hinge': loss_hinge_dis,
'wasserstein': loss_wgan_dis
}
if optimizer == "SGD":
G_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
Gen.parameters()),
g_lr,
momentum=momentum,
nesterov=nesterov)
D_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
Dis.parameters()),
d_lr,
momentum=momentum,
nesterov=nesterov)
elif optimizer == "RMSprop":
G_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
Gen.parameters()),
g_lr,
momentum=momentum,
alpha=alpha)
D_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
Dis.parameters()),
d_lr,
momentum=momentum,
alpha=alpha)
elif optimizer == "Adam":
G_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
Gen.parameters()),
g_lr, [beta1, beta2],
eps=1e-6)
D_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
Dis.parameters()),
d_lr, [beta1, beta2],
eps=1e-6)
elif optimizer == "AdaBelief":
G_optimizer = AdaBelief(filter(lambda p: p.requires_grad,
Gen.parameters()),
g_lr, [beta1, beta2],
eps=1e-12,
rectify=False)
D_optimizer = AdaBelief(filter(lambda p: p.requires_grad,
Dis.parameters()),
d_lr, [beta1, beta2],
eps=1e-12,
rectify=False)
else:
raise NotImplementedError
if checkpoint_folder is not None:
when = "current" if load_current is True else "best"
if not exists(abspath(checkpoint_folder)):
raise NotADirectoryError
checkpoint_dir = make_checkpoint_dir(checkpoint_folder, run_name)
g_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=G-{when}-weights-step*.pth".format(when=when)))[0]
d_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=D-{when}-weights-step*.pth".format(when=when)))[0]
Gen, G_optimizer, trained_seed, run_name, step, prev_ada_p = load_checkpoint(
Gen, G_optimizer, g_checkpoint_dir)
Dis, D_optimizer, trained_seed, run_name, step, prev_ada_p, best_step, best_fid, best_fid_checkpoint_path =\
load_checkpoint(Dis, D_optimizer, d_checkpoint_dir, metric=True)
logger = make_logger(run_name, None)
if ema:
g_ema_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=G_ema-{when}-weights-step*.pth".format(
when=when)))[0]
Gen_copy = load_checkpoint(Gen_copy,
None,
g_ema_checkpoint_dir,
ema=True)
Gen_ema.source, Gen_ema.target = Gen, Gen_copy
writer = SummaryWriter(log_dir=join('./logs', run_name))
if train_config['train']:
assert seed == trained_seed, "seed for sampling random numbers should be same!"
logger.info('Generator checkpoint is {}'.format(g_checkpoint_dir))
logger.info('Discriminator checkpoint is {}'.format(d_checkpoint_dir))
if freeze_layers > -1:
prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path = None, 0, 0, None, None
else:
checkpoint_dir = make_checkpoint_dir(checkpoint_folder, run_name)
if n_gpus > 1:
Gen = DataParallel(Gen, output_device=default_device)
Dis = DataParallel(Dis, output_device=default_device)
if ema:
Gen_copy = DataParallel(Gen_copy, output_device=default_device)
if synchronized_bn:
Gen = convert_model(Gen).to(default_device)
Dis = convert_model(Dis).to(default_device)
if ema:
Gen_copy = convert_model(Gen_copy).to(default_device)
if train_config['eval']:
inception_model = InceptionV3().to(default_device)
if n_gpus > 1:
inception_model = DataParallel(inception_model,
output_device=default_device)
mu, sigma = prepare_inception_moments(dataloader=eval_dataloader,
generator=Gen,
eval_mode=eval_type,
inception_model=inception_model,
splits=1,
run_name=run_name,
logger=logger,
device=default_device)
else:
mu, sigma, inception_model = None, None, None
train_eval = Train_Eval(
run_name=run_name,
best_step=best_step,
dataset_name=dataset_name,
eval_type=eval_type,
logger=logger,
writer=writer,
n_gpus=n_gpus,
gen_model=Gen,
dis_model=Dis,
inception_model=inception_model,
Gen_copy=Gen_copy,
Gen_ema=Gen_ema,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
train_dataloader=train_dataloader,
eval_dataloader=eval_dataloader,
freeze_layers=freeze_layers,
conditional_strategy=conditional_strategy,
pos_collected_numerator=model_config['model']
['pos_collected_numerator'],
z_dim=z_dim,
num_classes=num_classes,
hypersphere_dim=hypersphere_dim,
d_spectral_norm=d_spectral_norm,
g_spectral_norm=g_spectral_norm,
G_optimizer=G_optimizer,
D_optimizer=D_optimizer,
batch_size=batch_size,
g_steps_per_iter=model_config['optimization']['g_steps_per_iter'],
d_steps_per_iter=model_config['optimization']['d_steps_per_iter'],
accumulation_steps=model_config['optimization']['accumulation_steps'],
total_step=total_step,
G_loss=G_loss[adv_loss],
D_loss=D_loss[adv_loss],
contrastive_lambda=model_config['loss_function']['contrastive_lambda'],
margin=model_config['loss_function']['margin'],
tempering_type=model_config['loss_function']['tempering_type'],
tempering_step=model_config['loss_function']['tempering_step'],
start_temperature=model_config['loss_function']['start_temperature'],
end_temperature=model_config['loss_function']['end_temperature'],
weight_clipping_for_dis=model_config['loss_function']
['weight_clipping_for_dis'],
weight_clipping_bound=model_config['loss_function']
['weight_clipping_bound'],
gradient_penalty_for_dis=model_config['loss_function']
['gradient_penalty_for_dis'],
gradient_penalty_lambda=model_config['loss_function']
['gradient_penalty_lambda'],
deep_regret_analysis_for_dis=model_config['loss_function']
['deep_regret_analysis_for_dis'],
regret_penalty_lambda=model_config['loss_function']
['regret_penalty_lambda'],
cr=cr,
cr_lambda=model_config['loss_function']['cr_lambda'],
bcr=model_config['loss_function']['bcr'],
real_lambda=model_config['loss_function']['real_lambda'],
fake_lambda=model_config['loss_function']['fake_lambda'],
zcr=model_config['loss_function']['zcr'],
gen_lambda=model_config['loss_function']['gen_lambda'],
dis_lambda=model_config['loss_function']['dis_lambda'],
sigma_noise=model_config['loss_function']['sigma_noise'],
diff_aug=model_config['training_and_sampling_setting']['diff_aug'],
ada=model_config['training_and_sampling_setting']['ada'],
prev_ada_p=prev_ada_p,
ada_target=model_config['training_and_sampling_setting']['ada_target'],
ada_length=model_config['training_and_sampling_setting']['ada_length'],
prior=prior,
truncated_factor=truncated_factor,
ema=ema,
latent_op=model_config['training_and_sampling_setting']['latent_op'],
latent_op_rate=model_config['training_and_sampling_setting']
['latent_op_rate'],
latent_op_step=model_config['training_and_sampling_setting']
['latent_op_step'],
latent_op_step4eval=model_config['training_and_sampling_setting']
['latent_op_step4eval'],
latent_op_alpha=model_config['training_and_sampling_setting']
['latent_op_alpha'],
latent_op_beta=model_config['training_and_sampling_setting']
['latent_op_beta'],
latent_norm_reg_weight=model_config['training_and_sampling_setting']
['latent_norm_reg_weight'],
default_device=default_device,
print_every=train_config['print_every'],
save_every=train_config['save_every'],
checkpoint_dir=checkpoint_dir,
evaluate=train_config['eval'],
mu=mu,
sigma=sigma,
best_fid=best_fid,
best_fid_checkpoint_path=best_fid_checkpoint_path,
mixed_precision=mixed_precision,
train_config=train_config,
model_config=model_config,
)
if train_config['train']:
step = train_eval.train(current_step=step, total_step=total_step)
if train_config['eval']:
is_save = train_eval.evaluation(
step=step,
standing_statistics=standing_statistics,
standing_step=standing_step)
if train_config['save_images']:
train_eval.save_images(is_generate=True,
png=True,
npz=True,
standing_statistics=standing_statistics,
standing_step=standing_step)
if train_config['image_visualization']:
train_eval.run_image_visualization(
nrow=train_config['nrow'],
ncol=train_config['ncol'],
standing_statistics=standing_statistics,
standing_step=standing_step)
if train_config['k_nearest_neighbor']:
train_eval.run_nearest_neighbor(
nrow=train_config['nrow'],
ncol=train_config['ncol'],
standing_statistics=standing_statistics,
standing_step=standing_step)
if train_config['interpolation']:
assert architecture in [
"big_resnet", "biggan_deep"
], "Not supported except for biggan and biggan_deep."
train_eval.run_linear_interpolation(
nrow=train_config['nrow'],
ncol=train_config['ncol'],
fix_z=True,
fix_y=False,
standing_statistics=standing_statistics,
standing_step=standing_step)
train_eval.run_linear_interpolation(
nrow=train_config['nrow'],
ncol=train_config['ncol'],
fix_z=False,
fix_y=True,
standing_statistics=standing_statistics,
standing_step=standing_step)
if train_config['frequency_analysis']:
train_eval.run_frequency_analysis(
num_images=len(train_dataset) // num_classes,
standing_statistics=standing_statistics,
standing_step=standing_step)
def prepare_train_eval(cfgs, hdf5_path_train, **_):
if cfgs.seed == -1:
cudnn.benchmark, cudnn.deterministic = True, False
else:
fix_all_seed(cfgs.seed)
cudnn.benchmark, cudnn.deterministic = False, True
n_gpus, default_device = torch.cuda.device_count(), torch.cuda.current_device()
if n_gpus ==1: warnings.warn('You have chosen a specific GPU. This will completely disable data parallelism.')
if cfgs.disable_debugging_API: torch.autograd.set_detect_anomaly(False)
check_flag_0(cfgs.batch_size, n_gpus, cfgs.freeze_layers, cfgs.checkpoint_folder, cfgs.architecture, cfgs.img_size)
run_name = make_run_name(RUN_NAME_FORMAT, framework=cfgs.config_path.split('/')[3][:-5], phase='train')
prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path, mu, sigma, inception_model = None, 0, 0, None, None, None, None, None
logger = make_logger(run_name, None)
writer = SummaryWriter(log_dir=join('./logs', run_name))
logger.info('Run name : {run_name}'.format(run_name=run_name))
logger.info(cfgs.train_configs)
logger.info(cfgs.model_configs)
##### load dataset #####
logger.info('Loading train datasets...')
train_dataset = LoadDataset(cfgs.dataset_name, cfgs.data_path, train=True, download=True, resize_size=cfgs.img_size,
hdf5_path=hdf5_path_train, random_flip=cfgs.random_flip_preprocessing)
if cfgs.reduce_train_dataset < 1.0:
num_train = int(cfgs.reduce_train_dataset*len(train_dataset))
train_dataset, _ = torch.utils.data.random_split(train_dataset, [num_train, len(train_dataset) - num_train])
logger.info('Train dataset size : {dataset_size}'.format(dataset_size=len(train_dataset)))
logger.info('Loading {mode} datasets...'.format(mode=cfgs.eval_type))
eval_mode = True if cfgs.eval_type == 'train' else False
eval_dataset = LoadDataset(cfgs.dataset_name, cfgs.data_path, train=eval_mode, download=True, resize_size=cfgs.img_size,
hdf5_path=None, random_flip=False)
logger.info('Eval dataset size : {dataset_size}'.format(dataset_size=len(eval_dataset)))
train_dataloader = DataLoader(train_dataset, batch_size=cfgs.batch_size, shuffle=True, pin_memory=True, num_workers=cfgs.num_workers, drop_last=True)
eval_dataloader = DataLoader(eval_dataset, batch_size=cfgs.batch_size, shuffle=True, pin_memory=True, num_workers=cfgs.num_workers, drop_last=False)
##### build model #####
logger.info('Building model...')
module = __import__('models.{architecture}'.format(architecture=cfgs.architecture), fromlist=['something'])
logger.info('Modules are located on models.{architecture}'.format(architecture=cfgs.architecture))
Gen = module.Generator(cfgs.z_dim, cfgs.shared_dim, cfgs.img_size, cfgs.g_conv_dim, cfgs.g_spectral_norm, cfgs.attention,
cfgs.attention_after_nth_gen_block, cfgs.activation_fn, cfgs.conditional_strategy, cfgs.num_classes,
cfgs.g_init, cfgs.G_depth, cfgs.mixed_precision).to(default_device)
Dis = module.Discriminator(cfgs.img_size, cfgs.d_conv_dim, cfgs.d_spectral_norm, cfgs.attention, cfgs.attention_after_nth_dis_block,
cfgs.activation_fn, cfgs.conditional_strategy, cfgs.hypersphere_dim, cfgs.num_classes, cfgs.nonlinear_embed,
cfgs.normalize_embed, cfgs.d_init, cfgs.D_depth, cfgs.mixed_precision).to(default_device)
if cfgs.ema:
print('Preparing EMA for G with decay of {}'.format(cfgs.ema_decay))
Gen_copy = module.Generator(cfgs.z_dim, cfgs.shared_dim, cfgs.img_size, cfgs.g_conv_dim, cfgs.g_spectral_norm, cfgs.attention,
cfgs.attention_after_nth_gen_block, cfgs.activation_fn, cfgs.conditional_strategy, cfgs.num_classes,
initialize=False, G_depth=cfgs.G_depth, mixed_precision=cfgs.mixed_precision).to(default_device)
Gen_ema = ema(Gen, Gen_copy, cfgs.ema_decay, cfgs.ema_start)
else:
Gen_copy, Gen_ema = None, None
logger.info(count_parameters(Gen))
logger.info(Gen)
logger.info(count_parameters(Dis))
logger.info(Dis)
### define loss functions and optimizers
G_loss = {'vanilla': loss_dcgan_gen, 'least_square': loss_lsgan_gen, 'hinge': loss_hinge_gen, 'wasserstein': loss_wgan_gen}
D_loss = {'vanilla': loss_dcgan_dis, 'least_square': loss_lsgan_dis, 'hinge': loss_hinge_dis, 'wasserstein': loss_wgan_dis}
if cfgs.optimizer == "SGD":
G_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, Gen.parameters()), cfgs.g_lr, momentum=cfgs.momentum, nesterov=cfgs.nesterov)
D_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, Dis.parameters()), cfgs.d_lr, momentum=cfgs.momentum, nesterov=cfgs.nesterov)
elif cfgs.optimizer == "RMSprop":
G_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad, Gen.parameters()), cfgs.g_lr, momentum=cfgs.momentum, alpha=cfgs.alpha)
D_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad, Dis.parameters()), cfgs.d_lr, momentum=cfgs.momentum, alpha=cfgs.alpha)
elif cfgs.optimizer == "Adam":
G_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, Gen.parameters()), cfgs.g_lr, [cfgs.beta1, cfgs.beta2], eps=1e-6)
D_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, Dis.parameters()), cfgs.d_lr, [cfgs.beta1, cfgs.beta2], eps=1e-6)
else:
raise NotImplementedError
##### load checkpoints if needed #####
if cfgs.checkpoint_folder is None:
checkpoint_dir = make_checkpoint_dir(cfgs.checkpoint_folder, run_name)
else:
when = "current" if cfgs.load_current is True else "best"
if not exists(abspath(cfgs.checkpoint_folder)):
raise NotADirectoryError
checkpoint_dir = make_checkpoint_dir(cfgs.checkpoint_folder, run_name)
g_checkpoint_dir = glob.glob(join(checkpoint_dir,"model=G-{when}-weights-step*.pth".format(when=when)))[0]
d_checkpoint_dir = glob.glob(join(checkpoint_dir,"model=D-{when}-weights-step*.pth".format(when=when)))[0]
Gen, G_optimizer, trained_seed, run_name, step, prev_ada_p = load_checkpoint(Gen, G_optimizer, g_checkpoint_dir)
Dis, D_optimizer, trained_seed, run_name, step, prev_ada_p, best_step, best_fid, best_fid_checkpoint_path =\
load_checkpoint(Dis, D_optimizer, d_checkpoint_dir, metric=True)
logger = make_logger(run_name, None)
if cfgs.ema:
g_ema_checkpoint_dir = glob.glob(join(checkpoint_dir, "model=G_ema-{when}-weights-step*.pth".format(when=when)))[0]
Gen_copy = load_checkpoint(Gen_copy, None, g_ema_checkpoint_dir, ema=True)
Gen_ema.source, Gen_ema.target = Gen, Gen_copy
writer = SummaryWriter(log_dir=join('./logs', run_name))
if cfgs.train_configs['train']:
assert cfgs.seed == trained_seed, "seed for sampling random numbers should be same!"
logger.info('Generator checkpoint is {}'.format(g_checkpoint_dir))
logger.info('Discriminator checkpoint is {}'.format(d_checkpoint_dir))
if cfgs.freeze_layers > -1 :
prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path = None, 0, 0, None, None
##### wrap models with DP and convert BN to Sync BN #####
if n_gpus > 1:
Gen = DataParallel(Gen, output_device=default_device)
Dis = DataParallel(Dis, output_device=default_device)
if cfgs.ema:
Gen_copy = DataParallel(Gen_copy, output_device=default_device)
if cfgs.synchronized_bn:
Gen = convert_model(Gen).to(default_device)
Dis = convert_model(Dis).to(default_device)
if cfgs.ema:
Gen_copy = convert_model(Gen_copy).to(default_device)
##### load the inception network and prepare first/secend moments for calculating FID #####
if cfgs.eval:
inception_model = InceptionV3().to(default_device)
if n_gpus > 1: inception_model = DataParallel(inception_model, output_device=default_device)
mu, sigma = prepare_inception_moments(dataloader=eval_dataloader,
generator=Gen,
eval_mode=cfgs.eval_type,
inception_model=inception_model,
splits=1,
run_name=run_name,
logger=logger,
device=default_device)
worker = make_worker(
cfgs=cfgs,
run_name=run_name,
best_step=best_step,
logger=logger,
writer=writer,
n_gpus=n_gpus,
gen_model=Gen,
dis_model=Dis,
inception_model=inception_model,
Gen_copy=Gen_copy,
Gen_ema=Gen_ema,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
train_dataloader=train_dataloader,
eval_dataloader=eval_dataloader,
G_optimizer=G_optimizer,
D_optimizer=D_optimizer,
G_loss=G_loss[cfgs.adv_loss],
D_loss=D_loss[cfgs.adv_loss],
prev_ada_p=prev_ada_p,
default_device=default_device,
checkpoint_dir=checkpoint_dir,
mu=mu,
sigma=sigma,
best_fid=best_fid,
best_fid_checkpoint_path=best_fid_checkpoint_path,
)
if cfgs.train_configs['train']:
step = worker.train(current_step=step, total_step=cfgs.total_step)
if cfgs.eval:
is_save = worker.evaluation(step=step, standing_statistics=cfgs.standing_statistics, standing_step=cfgs.standing_step)
if cfgs.save_images:
worker.save_images(is_generate=True, png=True, npz=True, standing_statistics=cfgs.standing_statistics, standing_step=cfgs.standing_step)
if cfgs.image_visualization:
worker.run_image_visualization(nrow=cfgs.nrow, ncol=cfgs.ncol, standing_statistics=cfgs.standing_statistics, standing_step=cfgs.standing_step)
if cfgs.k_nearest_neighbor:
worker.run_nearest_neighbor(nrow=cfgs.nrow, ncol=cfgs.ncol, standing_statistics=cfgs.standing_statistics, standing_step=cfgs.standing_step)
if cfgs.interpolation:
assert cfgs.architecture in ["big_resnet", "biggan_deep"], "Not supported except for biggan and biggan_deep."
worker.run_linear_interpolation(nrow=cfgs.nrow, ncol=cfgs.ncol, fix_z=True, fix_y=False,
standing_statistics=cfgs.standing_statistics, standing_step=cfgs.standing_step)
worker.run_linear_interpolation(nrow=cfgs.nrow, ncol=cfgs.ncol, fix_z=False, fix_y=True,
standing_statistics=cfgs.standing_statistics, standing_step=cfgs.standing_step)
if cfgs.frequency_analysis:
worker.run_frequency_analysis(num_images=len(train_dataset)//cfgs.num_classes,
standing_statistics=cfgs.standing_statistics, standing_step=cfgs.standing_step)
def train_framework(
seed, num_workers, config_path, reduce_train_dataset, load_current,
type4eval_dataset, dataset_name, num_classes, img_size, data_path,
architecture, conditional_strategy, hypersphere_dim, nonlinear_embed,
normalize_embed, g_spectral_norm, d_spectral_norm, activation_fn,
attention, attention_after_nth_gen_block,
attention_after_nth_dis_block, z_dim, shared_dim, g_conv_dim,
d_conv_dim, G_depth, D_depth, optimizer, batch_size, d_lr, g_lr,
momentum, nesterov, alpha, beta1, beta2, total_step, adv_loss,
consistency_reg, g_init, d_init, random_flip_preprocessing, prior,
truncated_factor, latent_op, ema, ema_decay, ema_start,
synchronized_bn, hdf5_path_train, train_config, model_config, **_):
fix_all_seed(seed)
cudnn.benchmark = False # Not good Generator for undetermined input size
cudnn.deterministic = True
n_gpus = torch.cuda.device_count()
default_device = torch.cuda.current_device()
second_device = default_device if n_gpus == 1 else default_device + 1
assert batch_size % n_gpus == 0, "batch_size should be divided by the number of gpus "
if n_gpus == 1:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
start_step, best_step, best_fid, best_fid_checkpoint_path = 0, 0, None, None
run_name = make_run_name(RUN_NAME_FORMAT,
framework=config_path.split('/')[3][:-5],
phase='train')
logger = make_logger(run_name, None)
writer = SummaryWriter(log_dir=join('./logs', run_name))
logger.info('Run name : {run_name}'.format(run_name=run_name))
logger.info(train_config)
logger.info(model_config)
logger.info('Loading train datasets...')
train_dataset = LoadDataset(dataset_name,
data_path,
train=True,
download=True,
resize_size=img_size,
hdf5_path=hdf5_path_train,
consistency_reg=consistency_reg,
random_flip=random_flip_preprocessing)
if reduce_train_dataset < 1.0:
num_train = int(reduce_train_dataset * len(train_dataset))
train_dataset, _ = torch.utils.data.random_split(
train_dataset,
[num_train, len(train_dataset) - num_train])
logger.info('Train dataset size : {dataset_size}'.format(
dataset_size=len(train_dataset)))
logger.info('Loading {mode} datasets...'.format(mode=type4eval_dataset))
eval_mode = True if type4eval_dataset == 'train' else False
eval_dataset = LoadDataset(dataset_name,
data_path,
train=eval_mode,
download=True,
resize_size=img_size,
hdf5_path=None,
random_flip=False)
logger.info('Eval dataset size : {dataset_size}'.format(
dataset_size=len(eval_dataset)))
logger.info('Building model...')
if architecture == "dcgan":
assert img_size == 32, "Sry, StudioGAN does not support dcgan models for generation of images larger than 32 resolution."
module = __import__(
'models.{architecture}'.format(architecture=architecture),
fromlist=['something'])
logger.info('Modules are located on models.{architecture}'.format(
architecture=architecture))
Gen = module.Generator(z_dim, shared_dim, img_size, g_conv_dim,
g_spectral_norm, attention,
attention_after_nth_gen_block, activation_fn,
conditional_strategy, num_classes, synchronized_bn,
g_init, G_depth).to(default_device)
Dis = module.Discriminator(img_size, d_conv_dim, d_spectral_norm,
attention, attention_after_nth_dis_block,
activation_fn, conditional_strategy,
hypersphere_dim, num_classes, nonlinear_embed,
normalize_embed, synchronized_bn, d_init,
D_depth).to(default_device)
if ema:
print('Preparing EMA for G with decay of {}'.format(ema_decay))
Gen_copy = module.Generator(z_dim,
shared_dim,
img_size,
g_conv_dim,
g_spectral_norm,
attention,
attention_after_nth_gen_block,
activation_fn,
conditional_strategy,
num_classes,
synchronized_bn=False,
initialize=False,
G_depth=G_depth).to(default_device)
Gen_ema = ema_(Gen, Gen_copy, ema_decay, ema_start)
else:
Gen_copy, Gen_ema = None, None
if n_gpus > 1:
Gen = DataParallel(Gen, output_device=second_device)
Dis = DataParallel(Dis, output_device=second_device)
if ema:
Gen_copy = DataParallel(Gen_copy, output_device=second_device)
if synchronized_bn:
patch_replication_callback(Gen)
patch_replication_callback(Dis)
logger.info(count_parameters(Gen))
logger.info(Gen)
logger.info(count_parameters(Dis))
logger.info(Dis)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=num_workers,
drop_last=True)
eval_dataloader = DataLoader(eval_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=num_workers,
drop_last=False)
G_loss = {
'vanilla': loss_dcgan_gen,
'hinge': loss_hinge_gen,
'wasserstein': loss_wgan_gen
}
D_loss = {
'vanilla': loss_dcgan_dis,
'hinge': loss_hinge_dis,
'wasserstein': loss_wgan_dis
}
if optimizer == "SGD":
G_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
Gen.parameters()),
g_lr,
momentum=momentum,
nesterov=nesterov)
D_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
Dis.parameters()),
d_lr,
momentum=momentum,
nesterov=nesterov)
elif optimizer == "RMSprop":
G_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
Gen.parameters()),
g_lr,
momentum=momentum,
alpha=alpha)
D_optimizer = torch.optim.RMSprop(filter(lambda p: p.requires_grad,
Dis.parameters()),
d_lr,
momentum=momentum,
alpha=alpha)
elif optimizer == "Adam":
G_optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, Gen.parameters()), g_lr,
[beta1, beta2])
D_optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, Dis.parameters()), d_lr,
[beta1, beta2])
elif optimizer == "AdamP":
G_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
Gen.parameters()),
g_lr,
betas=(beta1, beta2))
D_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
Dis.parameters()),
d_lr,
betas=(beta1, beta2))
else:
raise NotImplementedError
checkpoint_dir = make_checkpoint_dir(train_config['checkpoint_folder'],
run_name)
if train_config['checkpoint_folder'] is not None:
when = "current" if load_current is True else "best"
g_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=G-{when}-weights-step*.pth".format(when=when)))[0]
d_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=D-{when}-weights-step*.pth".format(when=when)))[0]
Gen, G_optimizer, trained_seed, run_name, start_step, best_step = load_checkpoint(
Gen, G_optimizer, g_checkpoint_dir)
Dis, D_optimizer, trained_seed, run_name, start_step, best_step, best_fid, best_fid_checkpoint_path = load_checkpoint(
Dis, D_optimizer, d_checkpoint_dir, metric=True)
logger = make_logger(run_name, None)
if ema:
g_ema_checkpoint_dir = glob.glob(
join(checkpoint_dir,
"model=G_ema-{when}-weights-step*.pth".format(
when=when)))[0]
Gen_copy = load_checkpoint(Gen_copy,
None,
g_ema_checkpoint_dir,
ema=True)
Gen_ema.source, Gen_ema.target = Gen, Gen_copy
writer = SummaryWriter(log_dir=join('./logs', run_name))
assert seed == trained_seed, "seed for sampling random numbers should be same!"
logger.info('Generator checkpoint is {}'.format(g_checkpoint_dir))
logger.info('Discriminator checkpoint is {}'.format(d_checkpoint_dir))
if train_config['eval']:
inception_model = InceptionV3().to(default_device)
inception_model = DataParallel(inception_model,
output_device=second_device)
mu, sigma, is_score, is_std = prepare_inception_moments_eval_dataset(
dataloader=eval_dataloader,
generator=Gen,
eval_mode=type4eval_dataset,
inception_model=inception_model,
splits=10,
run_name=run_name,
logger=logger,
device=second_device)
else:
mu, sigma, inception_model = None, None, None
logger.info('Start training...')
trainer = Trainer(
run_name=run_name,
best_step=best_step,
dataset_name=dataset_name,
type4eval_dataset=type4eval_dataset,
logger=logger,
writer=writer,
n_gpus=n_gpus,
gen_model=Gen,
dis_model=Dis,
inception_model=inception_model,
Gen_copy=Gen_copy,
Gen_ema=Gen_ema,
train_dataloader=train_dataloader,
eval_dataloader=eval_dataloader,
conditional_strategy=conditional_strategy,
z_dim=z_dim,
num_classes=num_classes,
hypersphere_dim=hypersphere_dim,
d_spectral_norm=d_spectral_norm,
g_spectral_norm=g_spectral_norm,
G_optimizer=G_optimizer,
D_optimizer=D_optimizer,
batch_size=batch_size,
g_steps_per_iter=model_config['optimization']['g_steps_per_iter'],
d_steps_per_iter=model_config['optimization']['d_steps_per_iter'],
accumulation_steps=model_config['optimization']['accumulation_steps'],
total_step=total_step,
G_loss=G_loss[adv_loss],
D_loss=D_loss[adv_loss],
contrastive_lambda=model_config['loss_function']['contrastive_lambda'],
tempering_type=model_config['loss_function']['tempering_type'],
tempering_step=model_config['loss_function']['tempering_step'],
start_temperature=model_config['loss_function']['start_temperature'],
end_temperature=model_config['loss_function']['end_temperature'],
gradient_penalty_for_dis=model_config['loss_function']
['gradient_penalty_for_dis'],
gradient_penelty_lambda=model_config['loss_function']
['gradient_penelty_lambda'],
weight_clipping_for_dis=model_config['loss_function']
['weight_clipping_for_dis'],
weight_clipping_bound=model_config['loss_function']
['weight_clipping_bound'],
consistency_reg=consistency_reg,
consistency_lambda=model_config['loss_function']['consistency_lambda'],
diff_aug=model_config['training_and_sampling_setting']['diff_aug'],
prior=prior,
truncated_factor=truncated_factor,
ema=ema,
latent_op=latent_op,
latent_op_rate=model_config['training_and_sampling_setting']
['latent_op_rate'],
latent_op_step=model_config['training_and_sampling_setting']
['latent_op_step'],
latent_op_step4eval=model_config['training_and_sampling_setting']
['latent_op_step4eval'],
latent_op_alpha=model_config['training_and_sampling_setting']
['latent_op_alpha'],
latent_op_beta=model_config['training_and_sampling_setting']
['latent_op_beta'],
latent_norm_reg_weight=model_config['training_and_sampling_setting']
['latent_norm_reg_weight'],
default_device=default_device,
second_device=second_device,
print_every=train_config['print_every'],
save_every=train_config['save_every'],
checkpoint_dir=checkpoint_dir,
evaluate=train_config['eval'],
mu=mu,
sigma=sigma,
best_fid=best_fid,
best_fid_checkpoint_path=best_fid_checkpoint_path,
train_config=train_config,
model_config=model_config,
)
if conditional_strategy == 'ContraGAN' and train_config['train']:
trainer.run_ours(current_step=start_step, total_step=total_step)
elif train_config['train']:
trainer.run(current_step=start_step, total_step=total_step)
elif train_config['eval']:
is_save = trainer.evaluation(step=start_step)
if train_config['k_nearest_neighbor'] > 0:
trainer.K_Nearest_Neighbor(
train_config['criterion_4_k_nearest_neighbor'],
train_config['number_of_nearest_samples'],
random.randrange(num_classes))
def pretrain(data_dir, train_path, val_path, dictionary_path,
dataset_limit, vocabulary_size, batch_size, max_len, epochs, clip_grads, device,
layers_count, hidden_size, heads_count, d_ff, dropout_prob,
log_output, checkpoint_dir, print_every, save_every, config, run_name=None, **_):
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
train_path = train_path if data_dir is None else join(data_dir, train_path)
val_path = val_path if data_dir is None else join(data_dir, val_path)
dictionary_path = dictionary_path if data_dir is None else join(data_dir, dictionary_path)
run_name = run_name if run_name is not None else make_run_name(RUN_NAME_FORMAT, phase='pretrain', config=config)
logger = make_logger(run_name, log_output)
logger.info('Run name : {run_name}'.format(run_name=run_name))
logger.info(config)
logger.info('Constructing dictionaries...')
dictionary = IndexDictionary.load(dictionary_path=dictionary_path,
vocabulary_size=vocabulary_size)
vocabulary_size = len(dictionary)
logger.info(f'dictionary vocabulary : {vocabulary_size} tokens')
logger.info('Loading datasets...')
train_dataset = PairedDataset(data_path=train_path, dictionary=dictionary, dataset_limit=dataset_limit)
val_dataset = PairedDataset(data_path=val_path, dictionary=dictionary, dataset_limit=dataset_limit)
logger.info('Train dataset size : {dataset_size}'.format(dataset_size=len(train_dataset)))
logger.info('Building model...')
model = build_model(layers_count, hidden_size, heads_count, d_ff, dropout_prob, max_len, vocabulary_size)
logger.info(model)
logger.info('{parameters_count} parameters'.format(
parameters_count=sum([p.nelement() for p in model.parameters()])))
loss_model = MLMNSPLossModel(model)
if torch.cuda.device_count() > 1:
loss_model = DataParallel(loss_model, output_device=1)
metric_functions = [mlm_accuracy, nsp_accuracy]
train_dataloader = DataLoader(
train_dataset,
batch_size=batch_size,
collate_fn=pretraining_collate_function)
val_dataloader = DataLoader(
val_dataset,
batch_size=batch_size,
collate_fn=pretraining_collate_function)
optimizer = NoamOptimizer(model.parameters(),
d_model=hidden_size, factor=2, warmup_steps=10000, betas=(0.9, 0.999), weight_decay=0.01)
checkpoint_dir = make_checkpoint_dir(checkpoint_dir, run_name, config)
logger.info('Start training...')
trainer = Trainer(
loss_model=loss_model,
train_dataloader=train_dataloader,
val_dataloader=val_dataloader,
metric_functions=metric_functions,
optimizer=optimizer,
clip_grads=clip_grads,
logger=logger,
checkpoint_dir=checkpoint_dir,
print_every=print_every,
save_every=save_every,
device=device
)
trainer.run(epochs=epochs)
return trainer
def finetune(pretrained_checkpoint,
data_dir, train_path, val_path, dictionary_path,
vocabulary_size, batch_size, max_len, epochs, lr, clip_grads, device,
layers_count, hidden_size, heads_count, d_ff, dropout_prob,
log_output, checkpoint_dir, print_every, save_every, config, run_name=None, **_):
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
train_path = train_path if data_dir is None else join(data_dir, train_path)
val_path = val_path if data_dir is None else join(data_dir, val_path)
dictionary_path = dictionary_path if data_dir is None else join(data_dir, dictionary_path)
run_name = run_name if run_name is not None else make_run_name(RUN_NAME_FORMAT, phase='finetune', config=config)
logger = make_logger(run_name, log_output)
logger.info('Run name : {run_name}'.format(run_name=run_name))
logger.info(config)
logger.info('Constructing dictionaries...')
dictionary = IndexDictionary.load(dictionary_path=dictionary_path,
vocabulary_size=vocabulary_size)
vocabulary_size = len(dictionary)
logger.info(f'dictionary vocabulary : {vocabulary_size} tokens')
logger.info('Loading datasets...')
train_dataset = SST2IndexedDataset(data_path=train_path, dictionary=dictionary)
val_dataset = SST2IndexedDataset(data_path=val_path, dictionary=dictionary)
logger.info('Train dataset size : {dataset_size}'.format(dataset_size=len(train_dataset)))
logger.info('Building model...')
pretrained_model = build_model(layers_count, hidden_size, heads_count, d_ff, dropout_prob, max_len, vocabulary_size)
pretrained_model.load_state_dict(torch.load(pretrained_checkpoint, map_location='cpu')['state_dict'])
model = FineTuneModel(pretrained_model, hidden_size, num_classes=2)
logger.info(model)
logger.info('{parameters_count} parameters'.format(
parameters_count=sum([p.nelement() for p in model.parameters()])))
loss_model = ClassificationLossModel(model)
metric_functions = [classification_accuracy]
train_dataloader = DataLoader(
train_dataset,
batch_size=batch_size,
collate_fn=classification_collate_function)
val_dataloader = DataLoader(
val_dataset,
batch_size=batch_size,
collate_fn=classification_collate_function)
optimizer = Adam(model.parameters(), lr=lr)
checkpoint_dir = make_checkpoint_dir(checkpoint_dir, run_name, config)
logger.info('Start training...')
trainer = Trainer(
loss_model=loss_model,
train_dataloader=train_dataloader,
val_dataloader=val_dataloader,
metric_functions=metric_functions,
optimizer=optimizer,
clip_grads=clip_grads,
logger=logger,
checkpoint_dir=checkpoint_dir,
print_every=print_every,
save_every=save_every,
device=device
)
trainer.run(epochs=epochs)
return trainer
def load_worker(local_rank, cfgs, gpus_per_node, run_name, hdf5_path):
# -----------------------------------------------------------------------------
# define default variables for loading ckpt or evaluating the trained GAN model.
# -----------------------------------------------------------------------------
step, epoch, topk, best_step, best_fid, best_ckpt_path, is_best = \
0, 0, cfgs.OPTIMIZATION.batch_size, 0, None, None, False
aa_p = cfgs.AUG.ada_initial_augment_p if cfgs.AUG.ada_initial_augment_p != "N/A" else cfgs.AUG.apa_initial_augment_p
mu, sigma, eval_model, num_rows, num_cols = None, None, None, 10, 8
loss_list_dict = {"gen_loss": [], "dis_loss": [], "cls_loss": []}
metric_dict_during_train = {}
if "none" in cfgs.RUN.eval_metrics:
cfgs.RUN.eval_metrics = []
if "is" in cfgs.RUN.eval_metrics:
metric_dict_during_train.update({
"IS": [],
"Top1_acc": [],
"Top5_acc": []
})
if "fid" in cfgs.RUN.eval_metrics:
metric_dict_during_train.update({"FID": []})
if "prdc" in cfgs.RUN.eval_metrics:
metric_dict_during_train.update({
"Improved_Precision": [],
"Improved_Recall": [],
"Density": [],
"Coverage": []
})
# -----------------------------------------------------------------------------
# determine cuda, cudnn, and backends settings.
# -----------------------------------------------------------------------------
if cfgs.RUN.fix_seed:
cudnn.benchmark, cudnn.deterministic = False, True
else:
cudnn.benchmark, cudnn.deterministic = True, False
if cfgs.MODEL.backbone in ["stylegan2", "stylegan3"]:
# Improves training speed
conv2d_gradfix.enabled = True
# Avoids errors with the augmentation pipe
grid_sample_gradfix.enabled = True
if cfgs.RUN.mixed_precision:
# Allow PyTorch to internally use tf32 for matmul
torch.backends.cuda.matmul.allow_tf32 = False
# Allow PyTorch to internally use tf32 for convolutions
torch.backends.cudnn.allow_tf32 = False
# -----------------------------------------------------------------------------
# initialize all processes and fix seed of each process
# -----------------------------------------------------------------------------
if cfgs.RUN.distributed_data_parallel:
global_rank = cfgs.RUN.current_node * (gpus_per_node) + local_rank
print("Use GPU: {global_rank} for training.".format(
global_rank=global_rank))
misc.setup(global_rank, cfgs.OPTIMIZATION.world_size, cfgs.RUN.backend)
torch.cuda.set_device(local_rank)
else:
global_rank = local_rank
misc.fix_seed(cfgs.RUN.seed + global_rank)
# -----------------------------------------------------------------------------
# Intialize python logger.
# -----------------------------------------------------------------------------
if local_rank == 0:
logger = log.make_logger(cfgs.RUN.save_dir, run_name, None)
if cfgs.RUN.ckpt_dir is not None and cfgs.RUN.freezeD == -1:
folder_hier = cfgs.RUN.ckpt_dir.split("/")
if folder_hier[-1] == "":
folder_hier.pop()
logger.info(
"Run name : {run_name}".format(run_name=folder_hier.pop()))
else:
logger.info("Run name : {run_name}".format(run_name=run_name))
for k, v in cfgs.super_cfgs.items():
logger.info("cfgs." + k + " =")
logger.info(json.dumps(vars(v), indent=2))
else:
logger = None
# -----------------------------------------------------------------------------
# load train and evaluation datasets.
# -----------------------------------------------------------------------------
if cfgs.RUN.train or cfgs.RUN.intra_class_fid or cfgs.RUN.GAN_train or cfgs.RUN.GAN_test:
if local_rank == 0:
logger.info(
"Load {name} train dataset.".format(name=cfgs.DATA.name))
train_dataset = Dataset_(
data_name=cfgs.DATA.name,
data_dir=cfgs.RUN.data_dir,
train=True,
crop_long_edge=cfgs.PRE.crop_long_edge,
resize_size=cfgs.PRE.resize_size,
random_flip=cfgs.PRE.apply_rflip,
normalize=True,
hdf5_path=hdf5_path,
load_data_in_memory=cfgs.RUN.load_data_in_memory)
if local_rank == 0:
logger.info("Train dataset size: {dataset_size}".format(
dataset_size=len(train_dataset)))
else:
train_dataset = None
if len(cfgs.RUN.eval_metrics) + +cfgs.RUN.save_real_images + cfgs.RUN.k_nearest_neighbor + \
cfgs.RUN.frequency_analysis + cfgs.RUN.tsne_analysis:
if local_rank == 0:
logger.info("Load {name} {ref} dataset.".format(
name=cfgs.DATA.name, ref=cfgs.RUN.ref_dataset))
eval_dataset = Dataset_(
data_name=cfgs.DATA.name,
data_dir=cfgs.RUN.data_dir,
train=True if cfgs.RUN.ref_dataset == "train" else False,
crop_long_edge=False
if cfgs.DATA.name in cfgs.MISC.no_proc_data else True,
resize_size=None if cfgs.DATA.name in cfgs.MISC.no_proc_data else
cfgs.DATA.img_size,
random_flip=False,
hdf5_path=None,
normalize=True,
load_data_in_memory=False)
if local_rank == 0:
logger.info("Eval dataset size: {dataset_size}".format(
dataset_size=len(eval_dataset)))
else:
eval_dataset = None
# -----------------------------------------------------------------------------
# define a distributed sampler for DDP train and evaluation.
# define dataloaders for train and evaluation.
# -----------------------------------------------------------------------------
if cfgs.RUN.distributed_data_parallel:
cfgs.OPTIMIZATION.batch_size = cfgs.OPTIMIZATION.batch_size // cfgs.OPTIMIZATION.world_size
if cfgs.RUN.train and cfgs.RUN.distributed_data_parallel:
train_sampler = DistributedSampler(
train_dataset,
num_replicas=cfgs.OPTIMIZATION.world_size,
rank=local_rank,
shuffle=True,
drop_last=True)
topk = cfgs.OPTIMIZATION.batch_size
else:
train_sampler = None
cfgs.OPTIMIZATION.basket_size = cfgs.OPTIMIZATION.batch_size * cfgs.OPTIMIZATION.acml_steps * cfgs.OPTIMIZATION.d_updates_per_step
if cfgs.RUN.train or cfgs.RUN.intra_class_fid or cfgs.RUN.GAN_train or cfgs.RUN.GAN_test:
train_dataloader = DataLoader(dataset=train_dataset,
batch_size=cfgs.OPTIMIZATION.basket_size,
shuffle=(train_sampler is None),
pin_memory=True,
num_workers=cfgs.RUN.num_workers,
sampler=train_sampler,
drop_last=True,
persistent_workers=True)
else:
train_dataloader = None
if len(cfgs.RUN.eval_metrics) + +cfgs.RUN.save_real_images + cfgs.RUN.k_nearest_neighbor + \
cfgs.RUN.frequency_analysis + cfgs.RUN.tsne_analysis:
if cfgs.RUN.distributed_data_parallel:
eval_sampler = DistributedSampler(
eval_dataset,
num_replicas=cfgs.OPTIMIZATION.world_size,
rank=local_rank,
shuffle=False,
drop_last=False)
else:
eval_sampler = None
eval_dataloader = DataLoader(dataset=eval_dataset,
batch_size=cfgs.OPTIMIZATION.batch_size,
shuffle=False,
pin_memory=True,
num_workers=cfgs.RUN.num_workers,
sampler=eval_sampler,
drop_last=False)
else:
eval_dataloader = None
# -----------------------------------------------------------------------------
# load a generator and a discriminator
# if cfgs.MODEL.apply_g_ema is True, load an exponential moving average generator (Gen_copy).
# -----------------------------------------------------------------------------
Gen, Gen_mapping, Gen_synthesis, Dis, Gen_ema, Gen_ema_mapping, Gen_ema_synthesis, ema =\
model.load_generator_discriminator(DATA=cfgs.DATA,
OPTIMIZATION=cfgs.OPTIMIZATION,
MODEL=cfgs.MODEL,
STYLEGAN=cfgs.STYLEGAN,
MODULES=cfgs.MODULES,
RUN=cfgs.RUN,
device=local_rank,
logger=logger)
if local_rank != 0:
custom_ops.verbosity = "none"
# -----------------------------------------------------------------------------
# define optimizers for adversarial training
# -----------------------------------------------------------------------------
cfgs.define_optimizer(Gen, Dis)
# -----------------------------------------------------------------------------
# load the generator and the discriminator from a checkpoint if possible
# -----------------------------------------------------------------------------
if cfgs.RUN.ckpt_dir is not None:
if local_rank == 0:
os.remove(join(cfgs.RUN.save_dir, "logs", run_name + ".log"))
run_name, step, epoch, topk, aa_p, best_step, best_fid, best_ckpt_path, logger =\
ckpt.load_StudioGAN_ckpts(ckpt_dir=cfgs.RUN.ckpt_dir,
load_best=cfgs.RUN.load_best,
Gen=Gen,
Dis=Dis,
g_optimizer=cfgs.OPTIMIZATION.g_optimizer,
d_optimizer=cfgs.OPTIMIZATION.d_optimizer,
run_name=run_name,
apply_g_ema=cfgs.MODEL.apply_g_ema,
Gen_ema=Gen_ema,
ema=ema,
is_train=cfgs.RUN.train,
RUN=cfgs.RUN,
logger=logger,
global_rank=global_rank,
device=local_rank,
cfg_file=cfgs.RUN.cfg_file)
if topk == "initialize":
topk == cfgs.OPTIMIZATION.batch_size
if cfgs.MODEL.backbone in ["stylegan2", "stylegan3"]:
ema.ema_rampup = "N/A" # disable EMA rampup
if cfgs.MODEL.backbone == "stylegan3" and cfgs.STYLEGAN.stylegan3_cfg == "stylegan3-r":
cfgs.STYLEGAN.blur_init_sigma = "N/A" # disable blur rampup
if cfgs.AUG.apply_ada:
cfgs.AUG.ada_kimg = 100 # make ADA react faster at the beginning
if cfgs.RUN.ckpt_dir is None or cfgs.RUN.freezeD != -1:
if local_rank == 0:
cfgs.RUN.ckpt_dir = ckpt.make_ckpt_dir(
join(cfgs.RUN.save_dir, "checkpoints", run_name))
dict_dir = join(cfgs.RUN.save_dir, "statistics", run_name)
loss_list_dict = misc.load_log_dicts(directory=dict_dir,
file_name="losses.npy",
ph=loss_list_dict)
metric_dict_during_train = misc.load_log_dicts(
directory=dict_dir,
file_name="metrics.npy",
ph=metric_dict_during_train)
# -----------------------------------------------------------------------------
# prepare parallel training
# -----------------------------------------------------------------------------
if cfgs.OPTIMIZATION.world_size > 1:
Gen, Gen_mapping, Gen_synthesis, Dis, Gen_ema, Gen_ema_mapping, Gen_ema_synthesis =\
model.prepare_parallel_training(Gen=Gen,
Gen_mapping=Gen_mapping,
Gen_synthesis=Gen_synthesis,
Dis=Dis,
Gen_ema=Gen_ema,
Gen_ema_mapping=Gen_ema_mapping,
Gen_ema_synthesis=Gen_ema_synthesis,
MODEL=cfgs.MODEL,
world_size=cfgs.OPTIMIZATION.world_size,
distributed_data_parallel=cfgs.RUN.distributed_data_parallel,
synchronized_bn=cfgs.RUN.synchronized_bn,
apply_g_ema=cfgs.MODEL.apply_g_ema,
device=local_rank)
# -----------------------------------------------------------------------------
# load a pre-trained network (InceptionV3 or ResNet50 trained using SwAV)
# -----------------------------------------------------------------------------
if len(cfgs.RUN.eval_metrics) or cfgs.RUN.intra_class_fid:
eval_model = pp.LoadEvalModel(
eval_backbone=cfgs.RUN.eval_backbone,
resize_fn=cfgs.RUN.resize_fn,
world_size=cfgs.OPTIMIZATION.world_size,
distributed_data_parallel=cfgs.RUN.distributed_data_parallel,
device=local_rank)
if "fid" in cfgs.RUN.eval_metrics:
mu, sigma = pp.prepare_moments(data_loader=eval_dataloader,
eval_model=eval_model,
quantize=True,
cfgs=cfgs,
logger=logger,
device=local_rank)
if cfgs.RUN.is_ref_dataset:
pp.calculate_ins(data_loader=eval_dataloader,
eval_model=eval_model,
quantize=True,
splits=1,
cfgs=cfgs,
logger=logger,
device=local_rank)
# -----------------------------------------------------------------------------
# initialize WORKER for training and evaluating GAN
# -----------------------------------------------------------------------------
worker = WORKER(
cfgs=cfgs,
run_name=run_name,
Gen=Gen,
Gen_mapping=Gen_mapping,
Gen_synthesis=Gen_synthesis,
Dis=Dis,
Gen_ema=Gen_ema,
Gen_ema_mapping=Gen_ema_mapping,
Gen_ema_synthesis=Gen_ema_synthesis,
ema=ema,
eval_model=eval_model,
train_dataloader=train_dataloader,
eval_dataloader=eval_dataloader,
global_rank=global_rank,
local_rank=local_rank,
mu=mu,
sigma=sigma,
logger=logger,
aa_p=aa_p,
best_step=best_step,
best_fid=best_fid,
best_ckpt_path=best_ckpt_path,
loss_list_dict=loss_list_dict,
metric_dict_during_train=metric_dict_during_train,
)
# -----------------------------------------------------------------------------
# train GAN until "total_steps" generator updates
# -----------------------------------------------------------------------------
if cfgs.RUN.train:
if global_rank == 0:
logger.info("Start training!")
worker.training, worker.topk = True, topk
worker.prepare_train_iter(epoch_counter=epoch)
while step <= cfgs.OPTIMIZATION.total_steps:
if cfgs.OPTIMIZATION.d_first:
real_cond_loss, dis_acml_loss = worker.train_discriminator(
current_step=step)
gen_acml_loss = worker.train_generator(current_step=step)
else:
gen_acml_loss = worker.train_generator(current_step=step)
real_cond_loss, dis_acml_loss = worker.train_discriminator(
current_step=step)
if global_rank == 0 and (step + 1) % cfgs.RUN.print_every == 0:
worker.log_train_statistics(current_step=step,
real_cond_loss=real_cond_loss,
gen_acml_loss=gen_acml_loss,
dis_acml_loss=dis_acml_loss)
step += 1
if cfgs.LOSS.apply_topk:
if (epoch + 1) == worker.epoch_counter:
epoch += 1
worker.topk = losses.adjust_k(
current_k=worker.topk,
topk_gamma=cfgs.LOSS.topk_gamma,
sup_k=int(cfgs.OPTIMIZATION.batch_size *
cfgs.LOSS.topk_nu))
if step % cfgs.RUN.save_every == 0:
# visuailize fake images
if global_rank == 0:
worker.visualize_fake_images(num_cols=num_cols,
current_step=step)
# evaluate GAN for monitoring purpose
if len(cfgs.RUN.eval_metrics):
is_best = worker.evaluate(step=step,
metrics=cfgs.RUN.eval_metrics,
writing=True,
training=True)
# save GAN in "./checkpoints/RUN_NAME/*"
if global_rank == 0:
worker.save(step=step, is_best=is_best)
# stop processes until all processes arrive
if cfgs.RUN.distributed_data_parallel:
dist.barrier(worker.group)
if global_rank == 0:
logger.info("End of training!")
# -----------------------------------------------------------------------------
# re-evaluate the best GAN and conduct ordered analyses
# -----------------------------------------------------------------------------
print("")
worker.training, worker.epoch_counter = False, epoch
worker.gen_ctlr.standing_statistics = cfgs.RUN.standing_statistics
worker.gen_ctlr.standing_max_batch = cfgs.RUN.standing_max_batch
worker.gen_ctlr.standing_step = cfgs.RUN.standing_step
if global_rank == 0:
best_step = ckpt.load_best_model(ckpt_dir=cfgs.RUN.ckpt_dir,
Gen=Gen,
Dis=Dis,
apply_g_ema=cfgs.MODEL.apply_g_ema,
Gen_ema=Gen_ema,
ema=ema)
if len(cfgs.RUN.eval_metrics):
for e in range(cfgs.RUN.num_eval):
if global_rank == 0:
print(""), logger.info("-" * 80)
_ = worker.evaluate(step=best_step,
metrics=cfgs.RUN.eval_metrics,
writing=False,
training=False)
if cfgs.RUN.save_real_images:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.save_real_images()
if cfgs.RUN.save_fake_images:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.save_fake_images(num_images=cfgs.RUN.save_fake_images_num)
if cfgs.RUN.vis_fake_images:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.visualize_fake_images(num_cols=num_cols, current_step=best_step)
if cfgs.RUN.k_nearest_neighbor:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.run_k_nearest_neighbor(dataset=eval_dataset,
num_rows=num_rows,
num_cols=num_cols)
if cfgs.RUN.interpolation:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.run_linear_interpolation(num_rows=num_rows,
num_cols=num_cols,
fix_z=True,
fix_y=False)
worker.run_linear_interpolation(num_rows=num_rows,
num_cols=num_cols,
fix_z=False,
fix_y=True)
if cfgs.RUN.frequency_analysis:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.run_frequency_analysis(dataloader=eval_dataloader)
if cfgs.RUN.tsne_analysis:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.run_tsne(dataloader=eval_dataloader)
if cfgs.RUN.intra_class_fid:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.calculate_intra_class_fid(dataset=train_dataset)
if cfgs.RUN.semantic_factorization:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.run_semantic_factorization(
num_rows=cfgs.RUN.num_semantic_axis,
num_cols=num_cols,
maximum_variations=cfgs.RUN.maximum_variations)
if cfgs.RUN.GAN_train:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.compute_GAN_train_or_test_classifier_accuracy_score(
GAN_train=True, GAN_test=False)
if cfgs.RUN.GAN_test:
if global_rank == 0:
print(""), logger.info("-" * 80)
worker.compute_GAN_train_or_test_classifier_accuracy_score(
GAN_train=False, GAN_test=True)
if global_rank == 0:
wandb.finish()
def train_framework(dataset_name, architecture, num_classes, img_size, data_path, eval_dataset, hdf5_path_train, hdf5_path_valid, train_rate, auxiliary_classifier,
projection_discriminator, contrastive_training, hyper_dim, nonlinear_embed, normalize_embed, g_spectral_norm, d_spectral_norm, attention, reduce_class,
at_after_th_gen_block, at_after_th_dis_block, leaky_relu, g_init, d_init, latent_op, consistency_reg, make_positive_aug, synchronized_bn, ema,
ema_decay, ema_start, adv_loss, z_dim, shared_dim, g_conv_dim, d_conv_dim, batch_size, total_step, truncated_factor, prior, d_lr, g_lr,
beta1, beta2, batch4metrics, config, **_):
fix_all_seed(config['seed'])
cudnn.benchmark = True # Not good Generator for undetermined input size
cudnn.deterministic = False
n_gpus = torch.cuda.device_count()
default_device = torch.cuda.current_device()
second_device = default_device if n_gpus == 1 else default_device+1
assert batch_size % n_gpus == 0, "batch_size should be divided by the number of gpus "
if n_gpus == 1:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
start_step = 0
best_val_fid, best_checkpoint_fid_path, best_val_is, best_checkpoint_is_path = None, None, None, None
run_name = make_run_name(RUN_NAME_FORMAT,
framework=config['config_path'].split('/')[3][:-5],
phase='train',
config=config)
logger = make_logger(run_name, None)
writer = SummaryWriter(log_dir=join('./logs', run_name))
logger.info('Run name : {run_name}'.format(run_name=run_name))
logger.info(config)
logger.info('Loading train datasets...')
train_dataset = LoadDataset(dataset_name, data_path, train=True, download=True, resize_size=img_size, hdf5_path=hdf5_path_train,
consistency_reg=consistency_reg, make_positive_aug=make_positive_aug)
if train_rate < 1.0:
num_train = int(train_rate*len(train_dataset))
train_dataset, _ = torch.utils.data.random_split(train_dataset, [num_train, len(train_dataset) - num_train])
logger.info('Train dataset size : {dataset_size}'.format(dataset_size=len(train_dataset)))
logger.info('Loading valid datasets...')
valid_dataset = LoadDataset(dataset_name, data_path, train=False, download=True, resize_size=img_size, hdf5_path=hdf5_path_valid)
logger.info('Valid dataset size : {dataset_size}'.format(dataset_size=len(valid_dataset)))
logger.info('Building model...')
module = __import__('models.{architecture}'.format(architecture=architecture),fromlist=['something'])
logger.info('Modules are located on models.{architecture}'.format(architecture=architecture))
num_classes = int(reduce_class*num_classes)
Gen = module.Generator(z_dim, shared_dim, g_conv_dim, g_spectral_norm, attention, at_after_th_gen_block, leaky_relu, auxiliary_classifier,
projection_discriminator, num_classes, contrastive_training, synchronized_bn, g_init).to(default_device)
Dis = module.Discriminator(d_conv_dim, d_spectral_norm, attention, at_after_th_dis_block, leaky_relu, auxiliary_classifier,
projection_discriminator, hyper_dim, num_classes, contrastive_training, nonlinear_embed, normalize_embed,
synchronized_bn, d_init).to(default_device)
if ema:
print('Preparing EMA for G with decay of {}'.format(ema_decay))
Gen_copy = module.Generator(z_dim, shared_dim, g_conv_dim, g_spectral_norm, attention, at_after_th_gen_block, leaky_relu, auxiliary_classifier,
projection_discriminator, num_classes, contrastive_training, synchronized_bn=False, initialize=False).to(default_device)
Gen_ema = ema_(Gen, Gen_copy, ema_decay, ema_start)
else:
Gen_copy, Gen_ema = None, None
if n_gpus > 1:
Gen = DataParallel(Gen, output_device=second_device)
Dis = DataParallel(Dis, output_device=second_device)
if ema:
Gen_copy = DataParallel(Gen_copy, output_device=second_device)
if config['synchronized_bn']:
patch_replication_callback(Gen)
patch_replication_callback(Dis)
logger.info(count_parameters(Gen))
logger.info(Gen)
logger.info(count_parameters(Dis))
logger.info(Dis)
if reduce_class != 1.0:
assert dataset_name == "TINY_ILSVRC2012" or "ILSVRC2012", "reduce_class mode can not be applied on the CIFAR10 dataset"
n_train = int(reduce_class*len(train_dataset))
n_valid = int(reduce_class*len(valid_dataset))
train_weights = [1.0]*n_train + [0.0]*(len(train_dataset) - n_train)
valid_weights = [1.0]*n_valid + [0.0]*(len(valid_dataset) - n_valid)
train_sampler = torch.utils.data.sampler.WeightedRandomSampler(train_weights, len(train_weights))
valid_sampler = torch.utils.data.sampler.WeightedRandomSampler(valid_weights, len(valid_weights))
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size, sampler=train_sampler, shuffle=False,
pin_memory=True, num_workers=config['num_workers'], drop_last=True)
evaluation_dataloader = DataLoader(valid_dataset,
sampler=valid_sampler, batch_size=batch4metrics, shuffle=False,
pin_memory=True, num_workers=config['num_workers'], drop_last=False)
else:
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size, shuffle=True, pin_memory=True,
num_workers=config['num_workers'], drop_last=True)
evaluation_dataloader = DataLoader(valid_dataset,
batch_size=batch4metrics, shuffle=True, pin_memory=True,
num_workers=config['num_workers'], drop_last=False)
G_loss = {'vanilla': loss_dcgan_gen, 'hinge': loss_hinge_gen, 'wasserstein': loss_wgan_gen}
D_loss = {'vanilla': loss_dcgan_dis, 'hinge': loss_hinge_dis, 'wasserstein': loss_wgan_dis}
G_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, Gen.parameters()), g_lr, [beta1, beta2])
D_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, Dis.parameters()), d_lr, [beta1, beta2])
checkpoint_dir = make_checkpoint_dir(config['checkpoint_folder'], run_name, config)
if config['checkpoint_folder'] is not None:
logger = make_logger(run_name, config['log_output_path'])
g_checkpoint_dir = glob.glob(os.path.join(checkpoint_dir,"model=G-step=" + str(config['step']) + "*.pth"))[0]
d_checkpoint_dir = glob.glob(os.path.join(checkpoint_dir,"model=D-step=" + str(config['step']) + "*.pth"))[0]
Gen, G_optimizer, seed, run_name, start_step = load_checkpoint(Gen, G_optimizer, g_checkpoint_dir)
Dis, D_optimizer, seed, run_name, start_step, best_val_fid, best_checkpoint_fid_path,\
best_val_is, best_checkpoint_is_path = load_checkpoint(Dis, D_optimizer, d_checkpoint_dir, metric=True)
if ema:
g_ema_checkpoint_dir = glob.glob(os.path.join(checkpoint_dir, "model=G_ema-step=" + str(config['step']) + "*.pth"))[0]
Gen_copy = load_checkpoint(Gen_copy, None, g_ema_checkpoint_dir, ema=ema)
Gen_ema.source, Gen_ema.target = Gen, Gen_copy
writer = SummaryWriter(log_dir=join('./logs', run_name))
assert config['seed'] == seed, "seed for sampling random numbers should be same!"
logger.info('Generator checkpoint is {}'.format(g_checkpoint_dir))
logger.info('Discriminator checkpoint is {}'.format(d_checkpoint_dir))
if config['eval']:
inception_model = InceptionV3().to(default_device)
inception_model = DataParallel(inception_model, output_device=second_device)
mu, sigma, is_score, is_std = prepare_inception_moments_eval_dataset(dataloader=evaluation_dataloader,
inception_model=inception_model,
reduce_class=reduce_class,
splits=10,
logger=logger,
device=second_device,
eval_dataset=eval_dataset)
else:
mu, sigma, inception_model = None, None, None
logger.info('Start training...')
trainer = Trainer(
run_name=run_name,
logger=logger,
writer=writer,
n_gpus=n_gpus,
gen_model=Gen,
dis_model=Dis,
inception_model=inception_model,
Gen_copy=Gen_copy,
Gen_ema=Gen_ema,
train_dataloader=train_dataloader,
evaluation_dataloader=evaluation_dataloader,
G_loss=G_loss[adv_loss],
D_loss=D_loss[adv_loss],
auxiliary_classifier=auxiliary_classifier,
contrastive_training=contrastive_training,
contrastive_lambda=config['contrastive_lambda'],
softmax_posterior=config['softmax_posterior'],
contrastive_softmax=config['contrastive_softmax'],
hyper_dim=config['hyper_dim'],
tempering=config['tempering'],
discrete_tempering=config['discrete_tempering'],
tempering_times=config['tempering_times'],
start_temperature=config['start_temperature'],
end_temperature=config['end_temperature'],
gradient_penalty_for_dis=config['gradient_penalty_for_dis'],
lambda4lp=config['lambda4lp'],
lambda4gp=config['lambda4gp'],
weight_clipping_for_dis=config['weight_clipping_for_dis'],
weight_clipping_bound=config['weight_clipping_bound'],
latent_op=latent_op,
latent_op_rate=config['latent_op_rate'],
latent_op_step=config['latent_op_step'],
latent_op_step4eval=config['latent_op_step4eval'],
latent_op_alpha=config['latent_op_alpha'],
latent_op_beta=config['latent_op_beta'],
latent_norm_reg_weight=config['latent_norm_reg_weight'],
consistency_reg=consistency_reg,
consistency_lambda=config['consistency_lambda'],
make_positive_aug=make_positive_aug,
G_optimizer=G_optimizer,
D_optimizer=D_optimizer,
default_device=default_device,
second_device=second_device,
batch_size=batch_size,
z_dim=z_dim,
num_classes=num_classes,
truncated_factor=truncated_factor,
prior=prior,
g_steps_per_iter=config['g_steps_per_iter'],
d_steps_per_iter=config['d_steps_per_iter'],
accumulation_steps=config['accumulation_steps'],
lambda4ortho=config['lambda4ortho'],
print_every=config['print_every'],
save_every=config['save_every'],
checkpoint_dir=checkpoint_dir,
evaluate=config['eval'],
mu=mu,
sigma=sigma,
best_val_fid=best_val_fid,
best_checkpoint_fid_path=best_checkpoint_fid_path,
best_val_is=best_val_is,
best_checkpoint_is_path=best_checkpoint_is_path,
config=config,
)
if contrastive_training:
trainer.run_ours(current_step=start_step, total_step=total_step)
else:
trainer.run(current_step=start_step, total_step=total_step)