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', 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(): 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/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(): 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()