def gen_images(gen, n=50000, batchsize=100):
ims = []
xp = gen.xp
# start_time = time.time()
# print('Start!')
for i in range(0, n, batchsize):
if i % 2500 == 2500 - batchsize:
print(str(i) + " generated")
config = yaml_utils.Config(yaml.load(open(args.config_path)))
is_conditional = config.updater['args']['conditional']
if is_conditional:
y = sample_categorical(gen.n_classes, batchsize, xp=gen.xp)
else:
y = None
if args.sampling_space == 'pixel':
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
x = gen(batchsize, y=y)
x = sampler.langevin(x, y, dis)
elif args.sampling_space == 'latent':
x, _ = latent_sampler.langevin(batchsize,
gen,
dis,
y_fake=y,
eval=True)
x = chainer.cuda.to_cpu(x.data)
x = np.asarray(np.clip(x * 127.5 + 127.5, 0.0, 255.0), dtype=np.uint8)
ims.append(x)
ims = np.asarray(ims)
_, _, _, h, w = ims.shape
ims = ims.reshape((n, 3, h, w))
# stop_time = time.time()
# print('Stop! Time: '+str(stop_time-start_time))
return ims
def main():
from argparser import args
config = yaml_utils.Config(yaml.load(open(args.config_path)))
chainer.cuda.get_device_from_id(args.gpu).use()
out = args.eval_results_dir
from model_loader import gen, dis
np.random.seed(666)
classes = tuple(args.classes) if args.classes is not None else np.arange(
0, gen.n_classes, dtype=np.int32)
for c in classes:
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', True):
# x = gen_images_with_condition(gen, c=c, n=args.rows * args.columns, batchsize=args.rows * args.columns)
x = gen_eval_images(gen,
n=args.rows,
batchsize=args.rows,
langevin_steps=args.columns,
seeds=c)
_, _, h, w = x.shape
if args.sampling_space == 'latent':
args.columns += 1
x = x.reshape((args.rows, args.columns, 3, h, w))
x = x.transpose(0, 3, 1, 4, 2)
x = x.reshape((args.rows * h, args.columns * w, 3))
save_path = os.path.join(out, '{}.png'.format(str(c)))
if not os.path.exists(out):
os.makedirs(out)
Image.fromarray(x).save(save_path)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--rows', type=int, default=1)
parser.add_argument('--columns', type=int, default=1)
parser.add_argument('--classes', type=int, nargs="*", default=None)
parser.add_argument('--seed', type=int, default=1234)
parser.add_argument('--num_pngs', type=int, default=100)
args = parser.parse_args()
chainer.cuda.get_device_from_id(args.gpu).use()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
# Model
gen = load_models(config)
gen.to_gpu(args.gpu)
out = args.results_dir
chainer.serializers.load_npz(args.snapshot, gen)
np.random.seed(args.seed)
classes = tuple(args.classes) if args.classes is not None else np.arange(0, gen.n_classes, dtype=np.int32)
for c in classes:
for png_idx in range(args.num_pngs):
print('Generating png %d / %d...' % (png_idx, args.num_pngs))
with chainer.using_config('train', False), chainer.using_config('enable_backprop', False):
x = gen_images_with_condition(gen, c=c, n=args.rows * args.columns, batchsize=args.rows * args.columns)
_, _, h, w = x.shape
x = x.reshape((args.rows, args.columns, 3, h, w))
x = x.transpose(0, 3, 1, 4, 2)
x = x.reshape((args.rows * h, args.columns * w, 3))
save_path = os.path.join(out, 'SNGAN_%08d.png' % png_idx)
if not os.path.exists(out):
os.makedirs(out)
Image.fromarray(x).save(save_path)
def gen_eval_images(gen, n=50000, batchsize=100, seeds=1234, langevin_steps=5):
'''
langevin_steps: column
'''
ims = []
xp = gen.xp
xp.random.seed(seeds)
for i in range(0, n, batchsize):
print(i)
config = yaml_utils.Config(yaml.load(open(args.config_path)))
is_conditional = config.updater['args']['conditional']
if is_conditional:
y = sample_categorical(gen.n_classes, batchsize, xp=gen.xp)
else:
y = None
if args.sampling_space == 'pixel':
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
x = gen(batchsize, y=y)
for j in range(langevin_steps):
x = sampler.langevin(x, y, dis)
nx = chainer.cuda.to_cpu(x.data)
nx = np.asarray(np.clip(nx * 127.5 + 127.5, 0.0, 255.0),
dtype=np.uint8)
ims.append(nx)
elif args.sampling_space == 'latent':
z = Variable(
sample_continuous(gen.dim_z,
batchsize,
distribution=gen.distribution,
xp=gen.xp))
x = gen(batchsize, y=y, z=z)
nx = chainer.cuda.to_cpu(x.data)
nx = np.asarray(np.clip(nx * 127.5 + 127.5, 0.0, 255.0),
dtype=np.uint8)
ims.append(nx)
for j in range(langevin_steps):
x, z = latent_sampler.langevin(batchsize,
gen,
dis,
y_fake=y,
eval=True,
given_z=z)
nx = chainer.cuda.to_cpu(x.data)
nx = np.asarray(np.clip(nx * 127.5 + 127.5, 0.0, 255.0),
dtype=np.uint8)
ims.append(nx)
ims = list(map(list, zip(*ims)))
ims = np.asarray(ims)
_, _, _, h, w = ims.shape
if args.sampling_space == 'latent':
langevin_steps += 1
ims = ims.reshape((n * langevin_steps, 3, h, w))
return ims
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--n_intp', type=int, default=5)
parser.add_argument('--n_zs', type=int, default=5)
parser.add_argument('--classes', type=int, nargs="*", default=None)
parser.add_argument('--seed', type=int, default=1236)
args = parser.parse_args()
chainer.cuda.get_device(args.gpu).use()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
gen = load_models(config)
gen.to_gpu()
chainer.serializers.load_npz(args.snapshot, gen)
out = args.results_dir
xp = gen.xp
imgs = []
classes = tuple(args.classes) if args.classes is not None else [
np.random.randint(config.models['generator']['args']['n_classes']),
np.random.randint(config.models['generator']['args']['n_classes'])
]
for _ in range(args.n_zs):
z = xp.array([np.random.normal(size=(128, ))] * args.n_intp,
xp.float32)
ys = xp.array([[classes[0], classes[1]]] * args.n_intp, dtype=xp.int32)
ws_y = xp.array([
np.linspace(0, 1, args.n_intp)[::-1],
np.linspace(0, 1, args.n_intp)
],
dtype=xp.float32).T
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
x = gen(z=z, y=ys, weights=ws_y)
x = chainer.cuda.to_cpu(x.data)
x = np.asarray(np.clip(x * 127.5 + 127.5, 0.0, 255.0), dtype=np.uint8)
imgs.append(x)
img = np.stack(imgs)
_, _, _, h, w = img.shape
img = img.transpose(0, 3, 1, 4, 2)
img = img.reshape((args.n_zs * h, args.n_intp * w, 3))
save_path = os.path.join(
out, 'interpolated_images_{}-{}.png'.format(classes[0], classes[1]))
if not os.path.exists(out):
os.makedirs(out)
Image.fromarray(img).save(save_path)
def test(args):
config = yaml_utils.Config(yaml.load(open(args.config_path)))
chainer.cuda.get_device_from_id(0).use()
gen = load_gen(config)
chainer.serializers.load_npz(args.gen_model, gen)
gen.to_gpu()
xp = gen.xp
_, test = yaml_utils.load_dataset(config)
test_iter = chainer.iterators.SerialIterator(test,
config.batchsize_test,
repeat=False,
shuffle=False)
results_dir = args.results_dir
images_dir = os.path.join(results_dir, 'images')
if not os.path.exists(images_dir):
os.makedirs(images_dir)
n = 0
psnr = []
ssim = []
while True:
x, t, batchsize = get_batch(test_iter, xp)
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
y = gen(x)
x = x.array.get() * 127.5 + 127.5
y = np.clip(y.array.get() * 127.5 + 127.5, 0., 255.)
t = t.array.get() * 127.5 + 127.5
_psnr, _ssim = save_images(x, y, t, images_dir, current_n=n)
psnr += _psnr
ssim += _ssim
n += len(x)
if test_iter.is_new_epoch:
test_iter.reset()
break
print('psnr: {}'.format(np.mean(psnr)))
print('ssim: {}'.format(np.mean(ssim)))
psnr = list(map(str, psnr))
ssim = list(map(str, ssim))
with open(os.path.join(results_dir, 'psnr.txt'), 'w') as f:
f.write('\n'.join(psnr))
with open(os.path.join(results_dir, 'ssim.txt'), 'w') as f:
f.write('\n'.join(ssim))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--rows', type=int, default=5)
parser.add_argument('--columns', type=int, default=5)
parser.add_argument('--classes', type=int, nargs="*", default=None)
parser.add_argument('--dump_z', action='store_true', default=False)
args = parser.parse_args()
np.random.seed(1234)
chainer.cuda.get_device_from_id(args.gpu).use()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
gen = load_models(config)
gen.to_gpu(args.gpu)
chainer.serializers.load_npz(args.snapshot, gen)
out = args.results_dir
classes = tuple(args.classes) if args.classes is not None else np.arange(
0, gen.n_classes, dtype=np.int32)
for c in classes:
z = sample_continuous(128, args.rows * args.columns, xp=gen.xp)
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
x = gen_images_with_condition(gen,
z=z,
c=c,
n=args.rows * args.columns,
batchsize=args.rows * args.columns)
_, _, h, w = x.shape
x = x.reshape((args.rows, args.columns, 3, h, w))
x = x.transpose(0, 3, 1, 4, 2)
x = x.reshape((args.rows * h, args.columns * w, 3))
save_path = os.path.join(out, '{}.png'.format(str(c)))
if not os.path.exists(out):
os.makedirs(out)
Image.fromarray(x).save(save_path)
# dump latent variables
if args.dump_z:
for i in range(args.rows * args.columns):
np.save(os.path.join(out, 'z_{}-{}.npy'.format(c, i)),
z.get()[i])
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--stat_dir_path', type=str, default='')
parser.add_argument('--inception_model_path', type=str, default='')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--n_classes', type=int, default=1000)
parser.add_argument('--class_start_from', type=int, default=0)
parser.add_argument('--tf', action='store_true', default=False)
args = parser.parse_args()
chainer.cuda.get_device_from_id(args.gpu).use()
if not os.path.exists(args.results_dir):
os.makedirs(args.results_dir)
config = yaml_utils.Config(yaml.load(open(args.config_path)))
# Model
gen = load_models(config)
gen.to_gpu()
chainer.serializers.load_npz(args.snapshot, gen)
np.random.seed(1234)
if args.tf:
import source.inception.inception_score_tf
from source.inception.inception_score_tf import get_mean_and_cov as get_mean_cov
else:
from evaluation import get_mean_cov
model = load_inception_model(args.inception_model_path)
for c in range(args.class_start_from, args.n_classes):
print("class:{}".format(c))
stat = np.load(os.path.join(args.stat_dir_path, '{}.npz'.format(c)))
ims = gen_images_with_condition(gen, c, 5000,
batchsize=100).astype("f")
if args.tf:
mean, cov = get_mean_cov(ims)
else:
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
mean, cov = get_mean_cov(model, ims, batch_size=100)
fid = FID(stat["mean"], stat["cov"], mean, cov)
print(fid)
np.savetxt('{}/fid_{}.txt'.format(args.results_dir, c),
np.array([fid]))
def test(args):
config = yaml_utils.Config(yaml.load(open(args.config_path)))
if args.dir_nir is not None:
config.dataset['args']['args_test']['dir_nir'] = args.dir_nir
config.dataset['args']['args_test']['imlist_nir'] = args.imlist_nir
if args.dir_rgb is not None:
config.dataset['args']['args_test']['dir_rgb'] = args.dir_rgb
config.dataset['args']['args_test']['imlist_rgb'] = args.imlist_rgb
chainer.cuda.get_device_from_id(0).use()
gen = load_gen(config)
chainer.serializers.load_npz(args.gen_model, gen)
gen.to_gpu()
xp = gen.xp
_, test = yaml_utils.load_dataset(config)
test_iter = chainer.iterators.SerialIterator(test,
config.batchsize_test,
repeat=False,
shuffle=False)
results_dir = args.results_dir
if not os.path.exists(results_dir):
os.makedirs(results_dir)
n = 0
while True:
x, batchsize = get_batch(test_iter, xp)
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
out = gen(x)
out = np.clip(out.array.get() * 127.5 + 127.5, 0., 255.)
x = x.array.get() * 127.5 + 127.5
save_images(x, out, results_dir, current_n=n)
n += len(out)
if test_iter.is_new_epoch:
test_iter.reset()
break
return None
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--inception_model_path', type=str, default='')
parser.add_argument('--splits', type=int, default=10)
parser.add_argument('--tf', action='store_true', default=False)
args = parser.parse_args()
chainer.cuda.get_device_from_id(args.gpu).use()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
# Model
gen = load_models(config)
gen.to_gpu(args.gpu)
chainer.serializers.load_npz(args.snapshot, gen)
np.random.seed(1234)
xp = gen.xp
n = int(5000 * args.splits)
#for _ in range(50):
# gen(128)
print("Gen")
ims = gen_images(gen, n, batchsize=100).astype("f")
print(np.max(ims), np.min(ims))
if args.tf:
import source.inception.inception_score_tf as inception_score
mean, std = inception_score.get_inception_score(ims, args.splits)
print(mean, std)
else:
from evaluation import load_inception_model
from source.inception.inception_score import inception_score, Inception
model = load_inception_model(args.inception_model_path)
mean, std = inception_score(model, ims, splits=args.splits)
print(mean, std)
if not os.path.exists(args.results_dir):
os.makedirs(args.results_dir)
np.savetxt('{}/inception_score.txt'.format(args.results_dir),
np.array([mean, std]))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml', help='path to config file')
parser.add_argument('--results_dir', type=str, default='./results',
help='directory to save the results to')
parser.add_argument('--snapshot', type=str, default='',
help='path to the snapshot')
parser.add_argument('--loaderjob', type=int,
help='number of parallel data loading processes')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
device = 0
chainer.cuda.get_device_from_id(device).use()
print("init")
multiprocessing.set_start_method('forkserver')
# Model
gen, dis = load_models(config)
gen.to_gpu()
dis.to_gpu()
models = {"gen": gen, "dis": dis}
# Optimizer
opt_gen = make_optimizer(gen, alpha=config.adam['alpha'], beta1=config.adam['beta1'], beta2=config.adam['beta2'])
opt_dis = make_optimizer(dis, alpha=config.adam['alpha'], beta1=config.adam['beta1'], beta2=config.adam['beta2'])
opts = {"opt_gen": opt_gen, "opt_dis": opt_dis}
# Dataset
train, test = yaml_utils.load_dataset(config)
# Iterator
train_iter = chainer.iterators.MultiprocessIterator(train, config.batchsize, n_processes=args.loaderjob)
test_iter = chainer.iterators.SerialIterator(test, config.batchsize_test, repeat=False, shuffle=False)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': train_iter,
'optimizer': opts,
'device': device,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'), out=out)
report_keys = ["loss_dis", "loss_gen", "loss_l1", "psnr", "ssim"]
eval_func = yaml_utils.load_eval_func(config)
# Set up logging
trainer.extend(extensions.snapshot(), trigger=(config.snapshot_interval, 'iteration'))
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'), trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(extensions.LogReport(keys=report_keys,
trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys), trigger=(config.display_interval, 'iteration'))
trainer.extend(eval_func(test_iter, gen, dst=out), trigger=(config.evaluation_interval, 'iteration'), priority=chainer.training.extension.PRIORITY_WRITER)
for key in report_keys:
trainer.extend(extensions.PlotReport(key, trigger=(config.evaluation_interval, 'iteration'), file_name='{}.png'.format(key)))
trainer.extend(extensions.ProgressBar(update_interval=config.progressbar_interval))
ext_opt_gen = extensions.LinearShift('alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_gen)
ext_opt_dis = extensions.LinearShift('alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_dis)
trainer.extend(ext_opt_gen)
trainer.extend(ext_opt_dis)
if args.snapshot:
print("Resume training with snapshot:{}".format(args.snapshot))
chainer.serializers.load_npz(args.snapshot, trainer)
# Run the training
print("start training")
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path',
type=str,
default='configs/base.yml',
help='path to config file')
parser.add_argument('--gpu',
type=int,
default=0,
help='index of gpu to be used')
parser.add_argument('--data_dir', type=str, default='./data/imagenet')
parser.add_argument('--results_dir',
type=str,
default='./results/gans',
help='directory to save the results to')
parser.add_argument('--inception_model_path',
type=str,
default='./datasets/inception_model',
help='path to the inception model')
parser.add_argument('--snapshot',
type=str,
default='',
help='path to the snapshot')
parser.add_argument('--loaderjob',
type=int,
help='number of parallel data loading processes')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
chainer.cuda.get_device_from_id(args.gpu).use()
gen, dis = load_models(config)
gen.to_gpu(device=args.gpu)
dis.to_gpu(device=args.gpu)
seed_weights(gen)
seed_weights(dis)
print(np.sum(gen.block2.c2.b))
_ = input()
models = {"gen": gen, "dis": dis}
# Optimizer
opt_gen = make_optimizer(gen,
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opt_dis = make_optimizer(dis,
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opts = {"opt_gen": opt_gen, "opt_dis": opt_dis}
# Dataset
if config['dataset'][
'dataset_name'] != 'CIFAR10Dataset': # Cifar10 dataset handler does not take "root" as argument.
config['dataset']['args']['root'] = args.data_dir
dataset = yaml_utils.load_dataset(config)
# Iterator
iterator = chainer.iterators.MultiprocessIterator(
dataset, config.batchsize, n_processes=args.loaderjob)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'),
out=out)
report_keys = ["loss_dis", "loss_gen", "inception_mean", "inception_std"]
# Set up logging
trainer.extend(extensions.snapshot(),
trigger=(config.snapshot_interval, 'iteration'))
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'),
trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(config.display_interval, 'iteration'))
if gen.n_classes > 0:
trainer.extend(sample_generate_conditional(gen,
out,
n_classes=gen.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
else:
trainer.extend(sample_generate(gen, out),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(sample_generate_light(gen, out, rows=10, cols=10),
trigger=(config.evaluation_interval // 10, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(calc_inception(gen,
n_ims=5000,
splits=1,
path=args.inception_model_path),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(
extensions.ProgressBar(update_interval=config.progressbar_interval))
ext_opt_gen = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_gen)
ext_opt_dis = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_dis)
trainer.extend(ext_opt_gen)
trainer.extend(ext_opt_dis)
if args.snapshot:
print("Resume training with snapshot:{}".format(args.snapshot))
chainer.serializers.load_npz(args.snapshot, trainer)
# Run the training
print("start training")
trainer.run()
def main():
config = yaml_utils.Config(yaml.load(open(args.config_path)))
chainer.cuda.get_device_from_id(args.gpu).use()
from model_loader import gen, dis
models = {"gen": gen, "dis": dis}
# Optimizer
opt_gen = make_optimizer(
gen, alpha=config.adam['alpha'], beta1=config.adam['beta1'], beta2=config.adam['beta2'])
opt_dis = make_optimizer(
dis, alpha=config.adam['alpha'], beta1=config.adam['beta1'], beta2=config.adam['beta2'])
opts = {"opt_gen": opt_gen, "opt_dis": opt_dis}
# Dataset
# Cifar10 and STL10 dataset handler does not take "root" as argument.
if config['dataset']['dataset_name'] != 'CIFAR10Dataset' and \
config['dataset']['dataset_name'] != 'STL10Dataset':
config['dataset']['args']['root'] = args.data_dir
dataset = yaml_utils.load_dataset(config)
# Iterator
iterator = chainer.iterators.MultiprocessIterator(
dataset, config.batchsize, n_processes=args.loaderjob)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'), out=out)
report_keys = ["loss_dis", "loss_gen", "inception_mean", "inception_std", "FID_mean", "FID_std"]
# Set up logging
trainer.extend(extensions.snapshot(), trigger=(config.snapshot_interval, 'iteration'))
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'), trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(extensions.LogReport(keys=report_keys,
trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys), trigger=(config.display_interval, 'iteration'))
if gen.n_classes > 0:
trainer.extend(sample_generate_conditional(gen, out, n_classes=gen.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
else:
trainer.extend(sample_generate(gen, out),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(sample_generate_light(gen, out, rows=10, cols=10),
trigger=(config.evaluation_interval // 10, 'iteration'),
priority=extension.PRIORITY_WRITER)
# trainer.extend(calc_inception(gen, n_ims=50000, dst=out, splits=10, path=args.inception_model_path),
# trigger=(config.evaluation_interval, 'iteration'),
# priority=extension.PRIORITY_WRITER)
trainer.extend(calc_inception_and_FID(gen, n_ims=50000, dst=out, path=args.inception_model_path, splits=10, stat_file=args.FID_stat_file),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(extensions.ProgressBar(update_interval=config.progressbar_interval))
ext_opt_gen = extensions.LinearShift('alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_gen)
ext_opt_dis = extensions.LinearShift('alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_dis)
trainer.extend(ext_opt_gen)
trainer.extend(ext_opt_dis)
if args.snapshot:
print("Resume training with snapshot:{}".format(args.snapshot))
chainer.serializers.load_npz(args.snapshot, trainer)
# Run the training
print("start training")
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml', help='path to config file')
parser.add_argument('--gpu', type=int, default=0, help='index of gpu to be used')
parser.add_argument('--data_dir', type=str, default='./dataset/imagenet')
parser.add_argument('--results_dir', type=str, default='./results/temp', help='directory to save the results to')
parser.add_argument('--loaderjob', type=int, help='number of parallel data loading processes')
parser.add_argument('--layer', type=int, default=0, help='freeze discriminator layer')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path), Loader=yaml.FullLoader))
chainer.cuda.get_device_from_id(args.gpu).use()
# Fix randomness
random.seed(config['seed'])
np.random.seed(config['seed'])
cupy.random.seed(config['seed'])
# Model
G_src, D_src = load_models(config, mode='source')
G_tgt, D_tgt = load_models(config, mode='target')
G_src.to_gpu(device=args.gpu)
D_src.to_gpu(device=args.gpu)
G_tgt.to_gpu(device=args.gpu)
D_tgt.to_gpu(device=args.gpu)
models = {"G_src": G_src, "D_src": D_src, "G_tgt": G_tgt, "D_tgt": D_tgt}
# Optimizer
opt_G_tgt = make_optimizer(G_tgt, alpha=config.adam['alpha'], beta1=config.adam['beta1'], beta2=config.adam['beta2'])
opt_D_tgt = make_optimizer(D_tgt, alpha=config.adam['alpha'], beta1=config.adam['beta1'], beta2=config.adam['beta2'])
for i in range(1, args.layer + 1): # freeze discriminator
getattr(D_tgt, f'block{i}').disable_update()
opts = {"opt_G_tgt": opt_G_tgt, "opt_D_tgt": opt_D_tgt}
# Dataset
config['dataset']['args']['root'] = args.data_dir
dataset = yaml_utils.load_dataset(config)
# Iterator
iterator = chainer.iterators.MultiprocessIterator(dataset, config.batchsize, n_processes=args.loaderjob)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'), out=out)
report_keys = ["loss_dis", "loss_gen", "loss_FM", "FID"]
# Set up logging
trainer.extend(extensions.snapshot(filename='snapshot_best'),
trigger=training.triggers.MinValueTrigger("FID", trigger=(config.evaluation_interval, 'iteration')))
for m in models.values():
trainer.extend(extensions.snapshot_object(m, m.__class__.__name__ + '_best.npz'),
trigger=training.triggers.MinValueTrigger("FID", trigger=(config.evaluation_interval, 'iteration')))
trainer.extend(extensions.LogReport(keys=report_keys, trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys), trigger=(config.display_interval, 'iteration'))
trainer.extend(sample_generate_conditional(G_tgt, out, n_classes=G_tgt.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(sample_generate_light(G_tgt, out, rows=10, cols=10),
trigger=(config.evaluation_interval // 10, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(calc_FID(G_tgt, n_ims=5000, stat_file=config['eval']['stat_file']),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(extensions.ProgressBar(update_interval=config.progressbar_interval))
ext_opt_gen = extensions.LinearShift('alpha', (config.adam['alpha'], 0.), (config.iteration_decay_start, config.iteration), opt_G_tgt)
ext_opt_dis = extensions.LinearShift('alpha', (config.adam['alpha'], 0.), (config.iteration_decay_start, config.iteration), opt_D_tgt)
trainer.extend(ext_opt_gen)
trainer.extend(ext_opt_dis)
# Load source networks
chainer.serializers.load_npz(config['pretrained']['gen'], trainer.updater.models['G_src'])
chainer.serializers.load_npz(config['pretrained']['dis'], trainer.updater.models['D_src'])
load_parameters(trainer.updater.models['G_src'], trainer.updater.models['G_tgt'])
load_parameters(trainer.updater.models['D_src'], trainer.updater.models['D_tgt'])
# Run the training
print("start training")
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--rows', type=int, default=5)
parser.add_argument('--columns', type=int, default=5)
parser.add_argument('--classes', type=int, nargs="*", default=None)
parser.add_argument('--seed', type=int, default=1234)
parser.add_argument('--class_mask', type=str, default=None)
parser.add_argument('--z_file', type=str, default=None)
args = parser.parse_args()
np.random.seed(args.seed)
chainer.cuda.get_device(args.gpu).use()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
gen = load_models(config)
gen.to_gpu()
chainer.serializers.load_npz(args.snapshot, gen)
out = args.results_dir
xp = gen.xp
imgs = []
for _ in range(args.rows):
if args.z_file is None:
_z = xp.array(np.random.normal(size=(128, )), dtype=xp.float32)
z = xp.broadcast_to(_z, (args.columns, 128))
else:
_z = xp.array(np.load(args.z_file), dtype=xp.float32)
z = xp.broadcast_to(_z, (args.columns, 128))
if args.class_mask is not None:
# two-class interpolation with the given mask
classes = tuple(args.classes) if args.classes is not None and len(args.classes) == 2\
else [np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes)]
img_mask = Image.open(args.class_mask).convert("L")
sizes = [4, 8, 8, 16, 16, 32, 32, 64, 64, 128, 128, 256]
ws = []
for i_size, size in enumerate(sizes):
resized_mask = xp.array(img_mask.resize((size, size)),
dtype=xp.float32) / 255
w = xp.zeros((args.columns, size, size, gen.n_classes),
dtype=xp.float32)
for i in range(args.columns):
weight = i / (args.columns - 1.0)
w[i, :, :, classes[0]] = 1.0 - resized_mask * weight
w[i, :, :, classes[1]] += resized_mask * weight
ws.append(chainer.Variable(w))
else:
# 1-to-4-class interpolation on the quadrants (if mask is not given)
classes = tuple(args.classes) if args.classes is not None and len(args.classes) == 5\
else [np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes)]
sizes = [4, 8, 8, 16, 16, 32, 32, 64, 64, 128, 128, 256]
ws = [
chainer.Variable(
xp.zeros((args.columns, size, size, gen.n_classes),
dtype=xp.float32)) for size in sizes
]
for i, size in enumerate(sizes):
ws[i].data[:, :, :, classes[0]] = 1.0
for j in range(args.columns):
ws[i].data[j, :, :,
classes[0]] = 1.0 - j / (args.columns - 1.0)
ws[i].data[j, :size / 2, :size / 2,
classes[1]] = j / (args.columns - 1.0)
ws[i].data[j, :size / 2, size / 2:,
classes[2]] = j / (args.columns - 1.0)
ws[i].data[j, size / 2:, :size / 2,
classes[3]] = j / (args.columns - 1.0)
ws[i].data[j, size / 2:, size / 2:,
classes[4]] = j / (args.columns - 1.0)
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
x = gen.spatial_interpolation(z, ws)
x = chainer.cuda.to_cpu(x.data)
x = np.asarray(np.clip(x * 127.5 + 127.5, 0.0, 255.0), dtype=np.uint8)
imgs.append(x)
img = np.stack(imgs)
_, _, _, h, w = img.shape
img = img.transpose(0, 3, 1, 4, 2)
img = img.reshape((args.rows * h, args.columns * w, 3))
save_path = os.path.join(out, 'interpolated_images.png')
if not os.path.exists(out):
os.makedirs(out)
Image.fromarray(img).save(save_path)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--n_ys', type=int, default=5)
parser.add_argument('--n_zs', type=int, default=5)
parser.add_argument('--classes', type=int, nargs="*", default=None)
parser.add_argument('--n_images', type=int, default=100)
parser.add_argument('--seed', type=int, default=1234)
args = parser.parse_args()
chainer.cuda.get_device(args.gpu).use()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
# Model
gen = load_models(config)
gen.to_gpu()
out = args.results_dir
chainer.serializers.load_npz(args.snapshot, gen)
np.random.seed(args.seed)
xp = gen.xp
n_images = args.n_zs * args.n_ys
for k in range(args.n_images):
print(k)
classes = tuple(args.classes) if args.classes is not None else [
np.random.randint(1000),
np.random.randint(1000)
]
classes = [
xp.array([classes[0]] * n_images, xp.int32),
xp.array([classes[1]] * n_images, xp.int32)
]
noises = [
np.random.normal(size=(128, )),
np.random.normal(size=(128, ))
]
noises = [
xp.array([noises[0]] * n_images, dtype=xp.float32),
xp.array([noises[1]] * n_images, dtype=xp.float32)
]
ws_z_0 = []
ws_z_1 = []
ws_y_0 = []
ws_y_1 = []
for i in np.linspace(0, 1, args.n_zs):
for j in np.linspace(0, 1, args.n_ys):
ws_z_0.append(1 - i)
ws_z_1.append(i)
ws_y_0.append(1 - j)
ws_y_1.append(j)
ws_z = [xp.array(ws_z_0, xp.float32), xp.array(ws_z_1, xp.float32)]
ws_y = [xp.array(ws_y_0, xp.float32), xp.array(ws_y_1, xp.float32)]
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
x = gen(zs=noises, ys=classes, ws_z=ws_z, ws_y=ws_y)
x = chainer.cuda.to_cpu(x.data)
x = np.asarray(np.clip(x * 127.5 + 127.5, 0.0, 255.0), dtype=np.uint8)
print(x.shape)
_, _, h, w = x.shape
x = x.reshape((args.n_zs, args.n_ys, 3, h, w))
x = x.transpose(0, 3, 1, 4, 2)
x = x.reshape((args.n_zs * h, args.n_ys * w, 3))
save_path = os.path.join(out, 'interpolated_images_{}.png'.format(k))
if not os.path.exists(out):
os.makedirs(out)
Image.fromarray(x).save(save_path)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path',
type=str,
default='configs/base.yml',
help='path to config file')
parser.add_argument('--gpu',
type=int,
default=0,
help='index of gpu to be used')
parser.add_argument('--enc_model',
type=str,
default=None,
help='path to the encoder .npz file')
parser.add_argument('--gen_model',
type=str,
default='',
help='path to the generator .npz file')
parser.add_argument('--dis_model',
type=str,
default='',
help='path to the discriminator .npz file')
parser.add_argument('--src_class', type=int, help="target class")
parser.add_argument('--input', type=str, help="input image")
parser.add_argument('--mode',
type=str,
default='noaux',
help='set "aux" if you want to use auxiliary network')
parser.add_argument('--noenc',
action='store_true',
help='specify if you do not want to use encoder')
parser.add_argument('--iter_opt', type=int, default=200)
parser.add_argument('--result_dir', type=str, default="opt_output")
args = parser.parse_args()
np.random.seed(1234)
config_path = args.config_path
gpu = args.gpu
enc_model = args.enc_model
gen_model = args.gen_model
dis_model = args.dis_model
src_class = args.src_class
optimize_iterations = args.iter_opt
lmd_pixel = 0.2
use_aux = args.mode == "aux"
no_enc = args.noenc
# optimizer setup
config = yaml_utils.Config(yaml.load(open(config_path)))
chainer.cuda.get_device_from_id(gpu).use()
if 'encoder' in config.models:
gen, dis, enc = load_models(config)
else:
gen, dis = load_models(config)
enc = None
chainer.serializers.load_npz(gen_model, gen, strict=False)
chainer.serializers.load_npz(dis_model, dis)
gen.to_gpu(device=gpu)
dis.to_gpu(device=gpu)
if not no_enc and enc is not None:
chainer.serializers.load_npz(enc_model, enc)
enc.to_gpu(device=gpu)
xp = gen.xp
# load image
img_size = config['dataset']['args']['size']
inp_size = img_size
if 'input_size' in config['updater']['args']:
inp_size = config['updater']['args']['input_size']
patch_inp = np.asarray(
Image.open(args.input).resize((inp_size, inp_size)))[:, :, :3]
patch_gt = np.asarray(Image.open(args.input).resize(
(img_size, img_size)))[:, :, :3]
x = patch_inp.astype(np.float32) / 128. - 1
x = Variable(
xp.asarray(x).transpose(2, 0, 1).reshape(1, 3, inp_size, inp_size))
gt = patch_gt.astype(np.float32) / 128. - 1
gt = Variable(
xp.asarray(gt).transpose(2, 0, 1).reshape(1, 3, img_size, img_size))
c = Variable(xp.array([src_class], dtype=xp.int32))
with chainer.using_config("train", False), chainer.using_config(
"enable_backprop", False):
if enc is not None and not no_enc:
z = enc(x, c)
else:
z = Variable(xp.array([np.random.normal(size=(128, ))],
xp.float32))
if use_aux:
z2 = gen.forward_A(z)
if use_aux:
optimizer = optimizers.Adam(alpha=0.015)
target = chainer.links.Parameter(z2.data)
else:
optimizer = optimizers.Adam(alpha=0.075)
target = chainer.links.Parameter(z.data)
optimizer.setup(target)
# optimization body
print("optimizing...")
for iteration in range(optimize_iterations):
target.cleargrads()
_target = target.W
with chainer.using_config('train', False):
if use_aux:
recon, _, _, z = gen.forward(zeta=_target, y=c, return_zs=True)
else:
recon = gen(z=_target, y=c)
l1 = reconstruction_loss(dis, recon, gt)
l2 = lmd_pixel * pixel_loss(recon, gt)
loss = l1 + l2
loss.backward()
target.W.grad = _target.grad
optimizer.update()
if not use_aux:
z = target.W
print(iteration, l1.data, l2.data, loss.data)
if iteration % 10 == 0:
img = recon.data.get()[0].transpose(1, 2, 0) * 127.5 + 127.5
patch_recon = np.asarray(np.clip(img, 0, 255), dtype=np.uint8)
if not os.path.exists(args.result_dir):
os.mkdir(args.result_dir)
Image.fromarray(patch_recon).save("{}/opt_{}.png".format(
args.result_dir, iteration))
with open("{}/opt_{}.npy".format(args.result_dir, iteration),
"wb") as f:
np.save(f, z.data.get())
print("done.")
img = recon.data.get()[0].transpose(1, 2, 0) * 127.5 + 127.5
patch_recon = np.asarray(np.clip(img, 0, 255), dtype=np.uint8)
# reconstruction preview
Image.fromarray(patch_recon).save("{}/opt_final.png".format(
args.result_dir))
with open("{}/opt_final.npy".format(args.result_dir), "wb") as f:
np.save(f, z.data.get())
# z save
with open("opt-z.npy", "wb") as f:
np.save(f, z.data.get())
default='configs/base.yml',
help='path to config file')
parser.add_argument('--gen_model',
type=str,
default='',
help='path to the generator .npz file')
parser.add_argument('--gpu', type=int, default=-1)
args = parser.parse_args()
app = Flask(__name__)
app.config["MAX_CONTENT_LENGTH"] = 10000000 # allow 10 MB post
config_path = args.config_path
snapshot_path = args.gen_model
config = yaml_utils.Config(yaml.load(open(args.config_path)))
def load_models(config):
gen_conf = config.models['generator']
gen = yaml_utils.load_model(gen_conf['fn'], gen_conf['name'],
gen_conf['args'])
return gen
gen = load_models(config)
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use()
gen.to_gpu()
chainer.serializers.load_npz(args.gen_model, gen)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path',
type=str,
default='configs/sr.yml',
help='path to config file')
parser.add_argument('--gpu',
type=int,
default=0,
help='index of gpu to be used')
parser.add_argument('--results_dir',
type=str,
default='./results',
help='directory to save the results to')
parser.add_argument(
'--inception_model_path',
type=str,
default='./datasets/inception_model/inception_score.model',
help='path to the inception model')
parser.add_argument(
'--stat_file',
type=str,
default='./datasets/inception_model/fid_stats_cifar10_train.npz',
help='path to the inception model')
parser.add_argument('--snapshot',
type=str,
default='',
help='path to the snapshot')
parser.add_argument('--loaderjob',
type=int,
help='number of parallel data loading processes')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
chainer.cuda.get_device_from_id(args.gpu).use()
# set up the model
devices = {'main': 0, 'second': 1, 'third': 2, 'fourth': 3}
gen, dis = load_models(config)
model = {"gen": gen, "dis": dis}
names = list(six.iterkeys(devices))
try:
names.remove('main')
except ValueError:
raise KeyError("devices must contain a 'main' key.")
models = {'main': model}
for name in names:
g = copy.deepcopy(model['gen'])
d = copy.deepcopy(model['dis'])
if devices[name] >= 0:
g.to_gpu(device=devices[name])
d.to_gpu(device=devices[name])
models[name] = {"gen": g, "dis": d}
if devices['main'] >= 0:
models['main']['gen'].to_gpu(device=devices['main'])
models['main']['dis'].to_gpu(device=devices['main'])
links = [[name, link] for name, link in sorted(dis.namedlinks())]
for name, link in links:
print(name)
links = [[name, link] for name, link in sorted(gen.namedlinks())]
for name, link in links:
print(name)
# Optimizer
opt_gen = make_optimizer(models['main']['gen'],
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opt_dis = make_optimizer(models['main']['dis'],
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opts = {"opt_gen": opt_gen, "opt_dis": opt_dis}
# Dataset
dataset = yaml_utils.load_dataset(config)
# Iterator
iterator = chainer.iterators.MultiprocessIterator(
dataset, config.batchsize, n_processes=args.loaderjob)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'devices': devices,
'models': models,
'iterator': iterator,
'optimizer': opts
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'),
out=out)
report_keys = ["loss_dis", "loss_gen", "inception_mean", "FID"]
# Set up logging
trainer.extend(extensions.snapshot(),
trigger=(config.snapshot_interval, 'iteration'))
for m in models['main'].values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'),
trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(config.display_interval, 'iteration'))
if gen.n_classes > 0:
trainer.extend(sample_generate_conditional(models['main']['gen'],
out,
n_classes=gen.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
else:
trainer.extend(sample_generate(models['main']['gen'], out),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(sample_generate_light(models['main']['gen'],
out,
rows=10,
cols=10),
trigger=(config.evaluation_interval // 10, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(calc_inception(models['main']['gen'],
n_ims=5000,
splits=1,
dst=args.results_dir,
path=args.inception_model_path),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(calc_FID(models['main']['gen'],
n_ims=5000,
dst=args.results_dir,
path=args.inception_model_path,
stat_file=args.stat_file),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(monitor_largest_singular_values(models['main']['dis'],
dst=args.results_dir),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(
extensions.ProgressBar(update_interval=config.progressbar_interval))
ext_opt_gen = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_gen)
ext_opt_dis = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_dis)
trainer.extend(ext_opt_gen)
trainer.extend(ext_opt_dis)
if args.snapshot:
print("Resume training with snapshot:{}".format(args.snapshot))
chainer.serializers.load_npz(args.snapshot, trainer)
# Run the training
print("start training")
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--gpu', '-g', type=int, default=0)
parser.add_argument('--results_dir', type=str, default='./results/gans')
parser.add_argument('--snapshot', type=str, default='')
parser.add_argument('--rows', type=int, default=5)
parser.add_argument('--columns', type=int, default=5)
parser.add_argument('--classes', type=int, nargs="*", default=None)
parser.add_argument('--seed', type=int, default=1234)
parser.add_argument('--z1', type=str, default="z1.npy")
parser.add_argument('--z2', type=str, default="z2.npy")
parser.add_argument('--class_mask', type=str, default=None)
args = parser.parse_args()
np.random.seed(args.seed)
chainer.cuda.get_device(args.gpu).use()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
gen = load_models(config)
gen.to_gpu()
chainer.serializers.load_npz(args.snapshot, gen)
out = args.results_dir
xp = gen.xp
imgs = []
z1 = xp.array(np.load(args.z1), dtype=xp.float32)
z1 = xp.broadcast_to(z1, (args.columns, 128))
z2 = xp.array(np.load(args.z2), dtype=xp.float32)
z2 = xp.broadcast_to(z2, (args.columns, 128))
classes = None
for i_row in range(args.rows):
if args.class_mask is not None:
if classes is None:
classes = tuple(args.classes) if args.classes is not None and len(args.classes) == 2\
else [np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes)]
img_mask = Image.open(args.class_mask).convert("L")
sizes = [4, 8, 8, 16, 16, 32, 32, 64, 64, 128, 128, 256]
ws = []
for i_size, size in enumerate(sizes):
resized_mask = xp.array(img_mask.resize((size, size)),
dtype=xp.float32) / 255
w = xp.zeros((args.columns, size, size, gen.n_classes),
dtype=xp.float32)
for i in range(args.columns):
weight = i / (args.columns - 1.0)
if i_size <= 0:
weight *= 0.0
w[i, :, :, classes[0]] = 1.0 - resized_mask * weight
w[i, :, :, classes[1]] += resized_mask * weight
ws.append(chainer.Variable(w))
else:
if classes is None:
classes = tuple(args.classes) if args.classes is not None and len(args.classes) == 5\
else [np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes),
np.random.randint(gen.n_classes)]
sizes = [4, 8, 8, 16, 16, 32, 32, 64, 64, 128, 128, 256]
ws = [
chainer.Variable(
xp.zeros((args.columns, size, size, gen.n_classes),
dtype=xp.float32)) for size in sizes
]
for i, size in enumerate(sizes):
ws[i].data[:, :, :, classes[0]] = 1.0
for j in range(args.columns):
ws[i].data[j, :, :,
classes[0]] = 1.0 - j / (args.columns - 1.0)
ws[i].data[j, :size / 2, :size / 2,
classes[1]] = j / (args.columns - 1.0)
ws[i].data[j, :size / 2, size / 2:,
classes[2]] = j / (args.columns - 1.0)
ws[i].data[j, size / 2:, :size / 2,
classes[3]] = j / (args.columns - 1.0)
ws[i].data[j, size / 2:, size / 2:,
classes[4]] = j / (args.columns - 1.0)
print(classes)
# dimension of the feature map for each layer
sizes_blend = [4, 8, 16, 32, 64, 128, 256]
# specifying feature blending weights:
# blends[LAYER_ID][DATA_ID, Y, X, IMAGE_ID]
# LATENT_ID: 0 = source, 1 = 1st reference, ..., -1 = image to be generated
blends = [
xp.zeros((args.rows, size, size, 3), dtype=xp.float32)
for size in sizes_blend
]
# example: applying feature blending to the center region of the 0-th feature map (4x4 resolution)
blends[0][:, :, :, 0] = 1.0
blends[0][:, 1:3, 1:3, 0] = 1.0 - i_row / (
args.rows - 1.0) # source weight: 1, 0.75, 0.5, ...
blends[0][:, 1:3, 1:3, 1] = i_row / (
args.rows - 1.0) # reference weight: 0, 0.25, 0.5, ...
# you can change the values of blends[LAYER_ID > 0][VERTICAL_RANGE, HORIZONTAL_RANGE, IMAGE_ID]
# to apply feature blending to later layers
# e.g., blends[1][2:6, 2:6, 1] = i_row / (args.rows - 1.0)
# adds more detailed features of the reference image to the center region of the generated image
# Note: feature blending will not be applied to the i-th layer if blends[i][:, :, -1] == 1.0
for i in range(len(blends)):
blends[i][:, :, :, -1] = 1.0
blends[i][:, :, :, -1] -= blends[i][:, :, :, 0]
blends[i][:, :, :, -1] -= blends[i][:, :, :, 1]
with chainer.using_config('train', False), chainer.using_config(
'enable_backprop', False):
x = gen.spatial_interpolation(zs=[z1, z2],
weights=ws,
blends=blends)
x = chainer.cuda.to_cpu(x.data)
x = np.asarray(np.clip(x * 127.5 + 127.5, 0.0, 255.0), dtype=np.uint8)
imgs.append(x)
img = np.stack(imgs)
_, _, _, h, w = img.shape
img = img.transpose(
0, 3, 1, 4, 2
) # class interpolation on horizontal axis, feature blending on vertiacl axis.
img = img.reshape((args.rows * h, args.columns * w, 3))
save_path = os.path.join(out, 'interpolated_images.png')
if not os.path.exists(out):
os.makedirs(out)
Image.fromarray(img).save(save_path)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--config_path',
type=str,
default='jobs/pinet/fashionmnist_cnn_prodpoly_linear.yml')
parser.add_argument('--n_devices', type=int)
parser.add_argument('--test', action='store_true', default=False)
parser.add_argument('--communicator',
type=str,
default='hierarchical',
help='Type of communicator')
parser.add_argument('--results_dir',
type=str,
default='results_polynomial')
parser.add_argument('--inception_model_path',
type=str,
default='/home/user/inception/inception.model')
parser.add_argument('--resume', type=str, default='')
parser.add_argument('--gen_snapshot',
type=str,
default=None,
help='path to the generator snapshot')
parser.add_argument('--dis_snapshot',
type=str,
default=None,
help='path to the discriminator snapshot')
parser.add_argument('--loaderjob',
type=int,
help='Number of parallel data loading processes')
parser.add_argument('--multiprocessing',
action='store_true',
default=False)
parser.add_argument('--label', type=str, default='synth')
parser.add_argument('--batch_val', type=int, default=1000)
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
# # ensure that the paths of the config are correct.
config = ensure_config_paths(config)
try:
comm = chainermn.create_communicator(args.communicator)
except:
comm = chainermn.create_communicator()
device = comm.intra_rank
chainer.cuda.get_device_from_id(device).use()
# # get the pc name, e.g. for chainerui.
pcname = gethostname()
print('Init on pc: {}.'.format(pcname))
if comm.rank == 0:
print('==========================================')
print('Using {} communicator'.format(args.communicator))
print('==========================================')
gen, dis = load_models(config)
gen.to_gpu()
dis.to_gpu()
mma = ModelMovingAverage(0.999, gen)
models = {"gen": gen, "dis": dis}
if args.gen_snapshot is not None:
print('Loading generator: {}.'.format(args.gen_snapshot))
chainer.serializers.load_npz(args.gen_snapshot, gen)
if args.dis_snapshot is not None:
print('Loading discriminator: {}.'.format(args.dis_snapshot))
chainer.serializers.load_npz(args.dis_snapshot, dis)
# Optimizer
# # convenience function for optimizer:
func_opt = lambda net, alpha, wdr0=0: make_optimizer(
net,
comm,
chmn=args.multiprocessing,
alpha=alpha,
beta1=config.adam['beta1'],
beta2=config.adam['beta2'],
weight_decay_rate=wdr0)
# Optimizer
wdr = 0 if 'weight_decay_rate' not in config.updater[
'args'] else config.updater['args']['weight_decay_rate_gener']
opt_gen = func_opt(gen, config.adam['alpha'], wdr0=wdr)
keydopt = 'alphad' if 'alphad' in config.adam.keys() else 'alpha'
opt_dis = func_opt(dis, config.adam[keydopt])
opts = {"opt_gen": opt_gen, "opt_dis": opt_dis}
if hasattr(dis, 'fix_last') and hasattr(dis, 'lin') and dis.fix_last:
# # This should be used with care. It fixes the linear layer that
# # makes the classification in the discriminator.
print('Fixing the linear layer of the discriminator!')
dis.disable_update()
# Dataset
if comm.rank == 0:
dataset = yaml_utils.load_dataset(config)
# # even though not new samples, use as proxy for iid validation ones.
if hasattr(dataset, 'n_concats') and dataset.n_concats == 1:
valid_samples = (dataset.base[:args.batch_val] + 1) * 127.5
else:
valid_samples = [(dataset.get_example(i)[0] + 1) * 127.5
for i in range(args.batch_val)]
# # convert the validation to an array as required by the kl script.
valid_samples = np.array(valid_samples, dtype=np.float32)
else:
_ = yaml_utils.load_dataset(
config) # Dummy, for adding path to the dataset module
dataset = None
dataset = chainermn.scatter_dataset(dataset, comm)
# Iterator
multiprocessing.set_start_method('forkserver')
if args.multiprocessing:
# # In minoas this might fail with the forkserver.py error.
iterator = chainer.iterators.MultiprocessIterator(
dataset, config.batchsize, n_processes=args.loaderjob)
else:
iterator = chainer.iterators.SerialIterator(dataset, config.batchsize)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
'device': device,
'mma': mma,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
if not args.test:
mainf = '{}_{}'.format(strftime('%Y_%m_%d__%H_%M_%S'), args.label)
out = os.path.join(args.results_dir, mainf, '')
else:
out = 'results/test'
if comm.rank == 0:
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'),
out=out)
report_keys = [
'loss_dis', 'loss_gen', 'kl', 'ndb', 'JS', 'dis_real', 'dis_fake'
]
if comm.rank == 0:
# Set up logging
# trainer.extend(extensions.snapshot(), trigger=(config.snapshot_interval, 'iteration'))
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'),
trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(
extensions.LogReport(trigger=(config.display_interval,
'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(config.display_interval, 'iteration'))
trainer.extend(
extensions.ProgressBar(update_interval=config.display_interval))
if gen.n_classes == 0:
trainer.extend(sample_generate(mma.avg_model, out),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
print('unconditional image generation extension added.')
else:
trainer.extend(sample_generate_conditional(
mma.avg_model, out, n_classes=gen.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(divergence_trainer(gen,
valid_samples,
metric=['kl', 'ndb'],
batch=args.batch_val),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
# # convenience function for linearshift in optimizer:
func_opt_shift = lambda optim1: extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), optim1)
# # define the actual extensions (for optimizer shift).
trainer.extend(func_opt_shift(opt_gen))
trainer.extend(func_opt_shift(opt_dis))
if args.resume:
print("Resume Trainer")
chainer.serializers.load_npz(args.resume, trainer)
m1 = 'Generator params: {}. Discriminator params: {}.'
print(m1.format(gen.count_params(), dis.count_params()))
# Run the training
print("start training")
trainer.run()
print('The output dir was {}.'.format(out))
plot_losses_log(out, savefig=True)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path',
type=str,
default='configs/base.yml',
help='path to config file')
parser.add_argument('--data_dir', type=str, default='./data/imagenet')
parser.add_argument('--results_dir',
type=str,
default='./results/gans',
help='directory to save the results to')
parser.add_argument('--inception_model_path',
type=str,
default='./datasets/inception_model',
help='path to the inception model')
parser.add_argument('--snapshot',
type=str,
default='',
help='path to the snapshot')
parser.add_argument('--loaderjob',
type=int,
help='number of parallel data loading processes')
parser.add_argument('--communicator',
type=str,
default='hierarchical',
help='Type of communicator')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
comm = chainermn.create_communicator(args.communicator)
device = comm.intra_rank
chainer.cuda.get_device_from_id(device).use()
print("init")
if comm.rank == 0:
print('==========================================')
print('Using {} communicator'.format(args.communicator))
print('==========================================')
# Model
gen, dis = load_models(config)
gen.to_gpu()
dis.to_gpu()
models = {"gen": gen, "dis": dis}
# Optimizer
opt_gen = make_optimizer(gen,
comm,
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opt_dis = make_optimizer(dis,
comm,
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opts = {"opt_gen": opt_gen, "opt_dis": opt_dis}
# Dataset
config['dataset']['args']['root'] = args.data_dir
if comm.rank == 0:
dataset = yaml_utils.load_dataset(config)
else:
_ = yaml_utils.load_dataset(
config) # Dummy, for adding path to the dataset module
dataset = None
dataset = chainermn.scatter_dataset(dataset, comm)
# Iterator
multiprocessing.set_start_method('forkserver')
iterator = chainer.iterators.MultiprocessIterator(
dataset, config.batchsize, n_processes=args.loaderjob)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
'device': device,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
if comm.rank == 0:
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'),
out=out)
report_keys = ["loss_dis", "loss_gen", "inception_mean", "inception_std"]
if comm.rank == 0:
# Set up logging
trainer.extend(extensions.snapshot(),
trigger=(config.snapshot_interval, 'iteration'))
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'),
trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(config.display_interval,
'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(config.display_interval, 'iteration'))
trainer.extend(sample_generate_conditional(gen,
out,
n_classes=gen.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(sample_generate_light(gen, out, rows=10, cols=10),
trigger=(config.evaluation_interval // 10, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(calc_inception(gen,
n_ims=5000,
splits=1,
path=args.inception_model_path),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(
extensions.ProgressBar(
update_interval=config.progressbar_interval))
ext_opt_gen = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_gen)
ext_opt_dis = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_dis)
trainer.extend(ext_opt_gen)
trainer.extend(ext_opt_dis)
if args.snapshot:
print("Resume training with snapshot:{}".format(args.snapshot))
chainer.serializers.load_npz(args.snapshot, trainer)
# Run the training
print("start training")
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path', type=str, default='configs/base.yml')
parser.add_argument('--n_devices', type=int)
parser.add_argument('--test', action='store_true', default=False)
parser.add_argument('--communicator', type=str,
default='hierarchical', help='Type of communicator')
parser.add_argument('--results_dir', type=str, default='results_rocgan')
parser.add_argument('--inception_model_path', type=str,
default='/home/user/inception/inception.model')
parser.add_argument('--resume', type=str, default='')
parser.add_argument('--enc_snapshot', type=str, default=None, help='path to the encoder snapshot')
parser.add_argument('--dec_snapshot', type=str, default=None, help='path to the decoder snapshot')
parser.add_argument('--dis_snapshot', type=str, default=None, help='path to the discriminator snapshot')
parser.add_argument('--loaderjob', type=int,
help='Number of parallel data loading processes')
parser.add_argument('--multiprocessing', action='store_true', default=False)
parser.add_argument('--validation', type=int, default=1)
parser.add_argument('--valid_fn', type=str, default='files_valid_4k.txt',
help='filename of the validation file')
parser.add_argument('--label', type=str, default='synth')
parser.add_argument('--stats_fid', type=str, default='', help='path for FID stats')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
# # ensure that the paths of the config are correct.
config = ensure_config_paths(config)
comm = chainermn.create_communicator(args.communicator)
device = comm.intra_rank
chainer.cuda.get_device_from_id(device).use()
# # get the pc name, e.g. for chainerui.
pcname = gethostname()
imperialpc = 'doc.ic.ac.uk' in pcname or pcname in ['ladybug', 'odysseus']
print('Init on pc: {}.'.format(pcname))
if comm.rank == 0:
print('==========================================')
print('Using {} communicator'.format(args.communicator))
print('==========================================')
enc, dec, dis = load_models_cgan(config)
if chainer.cuda.available:
enc.to_gpu()
dec.to_gpu()
dis.to_gpu()
else:
print('No GPU found!!!\n')
mma1 = ModelMovingAverage(0.999, enc)
mma2 = ModelMovingAverage(0.999, dec)
models = {'enc': enc, 'dec': dec, 'dis': dis}
if args.enc_snapshot is not None:
print('Loading encoder: {}.'.format(args.enc_snapshot))
chainer.serializers.load_npz(args.enc_snapshot, enc)
if args.dec_snapshot is not None:
print('Loading decoder: {}.'.format(args.dec_snapshot))
chainer.serializers.load_npz(args.dec_snapshot, dec)
if args.dis_snapshot is not None:
print('Loading discriminator: {}.'.format(args.dis_snapshot))
chainer.serializers.load_npz(args.dis_snapshot, dis)
# # convenience function for optimizer:
func_opt = lambda net: make_optimizer(net, comm, chmn=args.multiprocessing,
alpha=config.adam['alpha'], beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
# Optimizer
opt_enc = func_opt(enc)
opt_dec = func_opt(dec)
opt_dis = func_opt(dis)
opts = {'opt_enc': opt_enc, 'opt_dec': opt_dec, 'opt_dis': opt_dis}
# Dataset
if comm.rank == 0:
dataset = yaml_utils.load_dataset(config)
printtime('Length of dataset: {}.'.format(len(dataset)))
if args.validation:
# # add the validation db if we do perform validation.
db_valid = yaml_utils.load_dataset(config, validation=True, valid_path=args.valid_fn)
else:
_ = yaml_utils.load_dataset(config) # Dummy, for adding path to the dataset module
dataset = None
if args.validation:
_ = yaml_utils.load_dataset(config, validation=True, valid_path=args.valid_fn)
db_valid = None
dataset = chainermn.scatter_dataset(dataset, comm)
if args.validation:
db_valid = chainermn.scatter_dataset(db_valid, comm)
# Iterator
multiprocessing.set_start_method('forkserver')
if args.multiprocessing:
# # In minoas this might fail with the forkserver.py error.
iterator = chainer.iterators.MultiprocessIterator(dataset, config.batchsize,
n_processes=args.loaderjob)
if args.validation:
iter_val = chainer.iterators.MultiprocessIterator(db_valid, config.batchsize,
n_processes=args.loaderjob,
shuffle=False, repeat=False)
else:
iterator = chainer.iterators.SerialIterator(dataset, config.batchsize)
if args.validation:
iter_val = chainer.iterators.SerialIterator(db_valid, config.batchsize,
shuffle=False, repeat=False)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
'device': device,
'mma1': mma1,
'mma2': mma2,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
if not args.test:
mainf = '{}_{}'.format(strftime('%Y_%m_%d__%H_%M_%S'), args.label)
out = os.path.join(args.results_dir, mainf, '')
else:
out = 'results/test'
if comm.rank == 0:
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'), out=out)
# # abbreviations below: inc -> incpetion, gadv -> grad_adv, lgen -> loss gener,
# # {m, sd, b}[var] -> {mean, std, best position} [var],
report_keys = ['loss_dis', 'lgen_adv', 'dis_real', 'dis_fake', 'loss_l1',
'mssim', 'sdssim', 'mmae', 'loss_projl', 'FID']
if comm.rank == 0:
# Set up logging
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'), trigger=(config.snapshot_interval, 'iteration'))
# trainer.extend(extensions.snapshot_object(
# mma.avg_model, mma.avg_model.__class__.__name__ + '_avgmodel_{.updater.iteration}.npz'),
# trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(extensions.LogReport(trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys), trigger=(config.display_interval, 'iteration'))
if args.validation:
# # add the appropriate extension for validating the model.
models_mma = {'enc': mma1.avg_model, 'dec': mma2.avg_model, 'dis': dis}
trainer.extend(validation_trainer(models_mma, iter_val, n=len(db_valid), export_best=True,
pout=out, p_inc=args.inception_model_path, eval_fid=True,
sfile=args.stats_fid),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(extensions.ProgressBar(update_interval=config.display_interval))
if imperialpc:
# [ChainerUI] Observe learning rate
trainer.extend(extensions.observe_lr(optimizer_name='opt_dis'))
# [ChainerUI] enable to send commands from ChainerUI
trainer.extend(CommandsExtension())
# [ChainerUI] save 'args' to show experimental conditions
save_args(args, out)
# # convenience function for linearshift in optimizer:
func_opt_shift = lambda optim1: extensions.LinearShift('alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start,
config.iteration), optim1)
# # define the actual extensions (for optimizer shift).
trainer.extend(func_opt_shift(opt_enc))
trainer.extend(func_opt_shift(opt_dec))
trainer.extend(func_opt_shift(opt_dis))
if args.resume:
print('Resume Trainer')
chainer.serializers.load_npz(args.resume, trainer)
# Run the training
printtime('start training')
trainer.run()
plot_losses_log(out, savefig=True)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path',
type=str,
default='configs/base.yml',
help='path to config file')
parser.add_argument('--gpu',
type=int,
default=0,
help='index of gpu to be used')
parser.add_argument('--input_dir', type=str, default='./data/imagenet')
parser.add_argument('--truth_dir', type=str, default='./data/imagenet')
parser.add_argument('--results_dir',
type=str,
default='./results/gans',
help='directory to save the results to')
parser.add_argument('--snapshot',
type=str,
default='',
help='path to the snapshot file to use')
parser.add_argument('--enc_model',
type=str,
default='',
help='path to the generator .npz file')
parser.add_argument('--gen_model',
type=str,
default='',
help='path to the generator .npz file')
parser.add_argument('--dis_model',
type=str,
default='',
help='path to the discriminator .npz file')
parser.add_argument('--loaderjob',
type=int,
help='number of parallel data loading processes')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
chainer.cuda.get_device_from_id(args.gpu).use()
gen, dis, enc = load_models(config)
chainer.serializers.load_npz(args.gen_model, gen, strict=False)
chainer.serializers.load_npz(args.dis_model, dis)
chainer.serializers.load_npz(args.enc_model, enc)
gen.to_gpu(device=args.gpu)
dis.to_gpu(device=args.gpu)
enc.to_gpu(device=args.gpu)
models = {"gen": gen, "dis": dis, "enc": enc}
opt_gen = make_optimizer(gen,
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opt_gen.add_hook(chainer.optimizer.WeightDecay(config.weight_decay))
opt_gen.add_hook(chainer.optimizer.GradientClipping(config.grad_clip))
# disable update of pre-trained weights
layers_to_train = ['lA1', 'lA2', 'lB1', 'lB2', 'preluW', 'preluMiddleW']
for layer in gen.children():
if not layer.name in layers_to_train:
layer.disable_update()
lmd_pixel = 0.05
def fast_loss(out, gt):
l1 = reconstruction_loss(dis, out, gt)
l2 = lmd_pixel * pixel_loss(out, gt)
loss = l1 + l2
return loss
gen.set_fast_loss(fast_loss)
opts = {"opt_gen": opt_gen}
# Dataset
config['dataset']['args']['root_input'] = args.input_dir
config['dataset']['args']['root_truth'] = args.truth_dir
dataset = yaml_utils.load_dataset(config)
# Iterator
iterator = chainer.iterators.MultiprocessIterator(
dataset, config.batchsize, n_processes=args.loaderjob)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'),
out=out)
report_keys = [
"loss_noab", "loss1", "loss2", "loss3", "fast_alpha", "loss_ae",
"fast_benefit", "min_slope", "max_slope", "min_slope_middle",
"max_slope_middle"
]
# Set up logging
trainer.extend(extensions.snapshot(),
trigger=(config.snapshot_interval, 'iteration'))
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'),
trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.ParameterStatistics(gen),
trigger=(config.display_interval, 'iteration'))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(config.display_interval, 'iteration'))
trainer.extend(sample_reconstruction_auxab(enc,
gen,
out,
n_classes=gen.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(
extensions.ProgressBar(update_interval=config.progressbar_interval))
ext_opt_gen = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_gen)
trainer.extend(ext_opt_gen)
if args.snapshot:
print("Resume training with snapshot:{}".format(args.snapshot))
chainer.serializers.load_npz(args.snapshot, trainer)
# Run the training
print("start training")
trainer.run()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config_path',
type=str,
default='configs/base.yml',
help='path to config file')
parser.add_argument('--gpu',
type=int,
default=0,
help='index of gpu to be used')
parser.add_argument('--input_dir', type=str, default='./data/imagenet')
parser.add_argument('--truth_dir', type=str, default='./data/imagenet')
parser.add_argument('--results_dir',
type=str,
default='./results/gans',
help='directory to save the results to')
parser.add_argument('--snapshot',
type=str,
default='',
help='path to the snapshot')
parser.add_argument('--gen_model',
type=str,
default='',
help='path to the generator .npz file')
parser.add_argument('--dis_model',
type=str,
default='',
help='path to the discriminator .npz file')
parser.add_argument('--loaderjob',
type=int,
help='number of parallel data loading processes')
args = parser.parse_args()
config = yaml_utils.Config(yaml.load(open(args.config_path)))
chainer.cuda.get_device_from_id(args.gpu).use()
gen, dis, enc = load_models(config)
chainer.serializers.load_npz(args.gen_model, gen)
chainer.serializers.load_npz(args.dis_model, dis)
gen.to_gpu(device=args.gpu)
dis.to_gpu(device=args.gpu)
enc.to_gpu(device=args.gpu)
models = {"gen": gen, "dis": dis, "enc": enc}
opt_enc = make_optimizer(enc,
alpha=config.adam['alpha'],
beta1=config.adam['beta1'],
beta2=config.adam['beta2'])
opts = {"opt_enc": opt_enc}
# Dataset
config['dataset']['args']['root_input'] = args.input_dir
config['dataset']['args']['root_truth'] = args.truth_dir
dataset = yaml_utils.load_dataset(config)
# Iterator
iterator = chainer.iterators.MultiprocessIterator(
dataset, config.batchsize, n_processes=args.loaderjob)
kwargs = config.updater['args'] if 'args' in config.updater else {}
kwargs.update({
'models': models,
'iterator': iterator,
'optimizer': opts,
})
updater = yaml_utils.load_updater_class(config)
updater = updater(**kwargs)
out = args.results_dir
create_result_dir(out, args.config_path, config)
trainer = training.Trainer(updater, (config.iteration, 'iteration'),
out=out)
report_keys = ["loss", "min_slope", "max_slope", "min_z", "max_z"]
# Set up logging
trainer.extend(extensions.snapshot(),
trigger=(config.snapshot_interval, 'iteration'))
for m in models.values():
trainer.extend(extensions.snapshot_object(
m, m.__class__.__name__ + '_{.updater.iteration}.npz'),
trigger=(config.snapshot_interval, 'iteration'))
trainer.extend(
extensions.LogReport(keys=report_keys,
trigger=(config.display_interval, 'iteration')))
trainer.extend(extensions.PrintReport(report_keys),
trigger=(config.display_interval, 'iteration'))
trainer.extend(sample_reconstruction(enc,
gen,
out,
n_classes=gen.n_classes),
trigger=(config.evaluation_interval, 'iteration'),
priority=extension.PRIORITY_WRITER)
trainer.extend(
extensions.ProgressBar(update_interval=config.progressbar_interval))
ext_opt_enc = extensions.LinearShift(
'alpha', (config.adam['alpha'], 0.),
(config.iteration_decay_start, config.iteration), opt_enc)
trainer.extend(ext_opt_enc)
if args.snapshot:
print("Resume training with snapshot:{}".format(args.snapshot))
chainer.serializers.load_npz(args.snapshot, trainer)
# Run the training
print("start training")
trainer.run()
