def func(trainer):
c = 0
mixing = mix > random()
while c < number:
n = min(number - c, batch)
z = generator.generate_latent(n)
mix_z = generator.generate_latent(n) if mixing else None
y = generator(z, stage, trainer.updater.alpha, mix_z)
z.to_cpu()
y.to_cpu()
for i in range(n):
path = filepath(
directory,
f"{stage}_{trainer.updater.iteration}_{trainer.updater.alpha:.3f}_{c + i + 1}",
"png")
path = path if force else altfilepath(path)
save_image(y.array[i], path)
if save_latent:
path = filepath(
directory,
f"{stage}_{trainer.updater.iteration}_{trainer.updater.alpha:.3f}_{c + i + 1}",
"npy")
path = path if force else altfilepath(path)
save_array(z.array[i], path)
c += n
def func(trainer):
path = filepath(directory, f"mopt_{stage}_{trainer.updater.iteration}",
"hdf5")
path = path if force else altfilepath(path)
serializers.save_hdf5(path, mapper_optimizer)
path = filepath(directory, f"gopt_{stage}_{trainer.updater.iteration}",
"hdf5")
path = path if force else altfilepath(path)
serializers.save_hdf5(path, generator_optimizer)
path = filepath(directory, f"dopt_{stage}_{trainer.updater.iteration}",
"hdf5")
path = path if force else altfilepath(path)
serializers.save_hdf5(path, discriminator_optimizer)
def func(trainer):
generator.to_cpu()
averaged_generator.to_cpu()
discriminator.to_cpu()
path = filepath(directory, f"gen_{stage}_{trainer.updater.iteration}",
"hdf5")
path = path if force else altfilepath(path)
serializers.save_hdf5(path, generator)
path = filepath(directory,
f"avgen_{stage}_{trainer.updater.iteration}", "hdf5")
path = path if force else altfilepath(path)
serializers.save_hdf5(path, averaged_generator)
path = filepath(directory, f"dis_{stage}_{trainer.updater.iteration}",
"hdf5")
path = path if force else altfilepath(path)
serializers.save_hdf5(path, discriminator)
if device >= 0:
generator.to_gpu(device)
averaged_generator.to_gpu(device)
discriminator.to_gpu(device)
generator, averaged_generator, discriminator, iterator, {
"mapper": mapper_optimizer,
"generator": generator_optimizer,
"discriminator": discriminator_optimizer
}, args.device, args.stage, args.mix, args.alpha, args.delta, args.gamma,
args.decay, args.lsgan)
# Init result directory
print("Initializing destination directory...")
if args.wipe:
rmtree(args.result, ignore_errors=True)
mkdirp(args.result)
# Dump command-line options
path = filepath(args.result, "args_quit" if args.quit else "args", "json")
path = path if args.force else altfilepath(path)
with open(path, mode="w", encoding="utf-8") as fp:
dump(vars(args), fp, indent=2, sort_keys=True)
# Define extension to output images in progress
def save_middle_images(generator,
stage,
directory,
number,
batch,
mix,
force=True,
save_latent=True):
@make_extension()
def func(trainer):
if args.center is not None:
print("Loading latent")
center = generator.wrap_latent(load_array(args.center))
else:
center = None
# Init destination folder
print("Initializing destination directory")
if args.wipe:
rmtree(args.directory, ignore_errors=True)
mkdirp(args.directory)
# Dump command-line options
if args.dump_json:
path = filepath(args.directory, "args_quit" if args.quit else "args", "json")
path = path if args.force else altfilepath(path)
with open(path, mode="w", encoding="utf-8") as fp:
dump(vars(args), fp, indent=2, sort_keys=True)
# Quit mode
if args.quit:
print("Finished (Quit mode)")
exit(0)
# Generate images
c = 0
mean_w = None if args.psi is None else generator.calculate_mean_w()
while c < args.number:
n = min(args.number - c, args.batch)
z = generator.generate_latent(n, center=center, sd=args.sd)
y = generator(z, args.stage, alpha=args.alpha, psi=args.psi, mean_w=mean_w)