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()