def main(): parser = argparse.ArgumentParser(description='chainer implementation of pix2pix') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--model', '-m', default='', help='model snapshot') parser.add_argument('--enc', '-e', type=str, default='enc_iter_60000.npz', help='encoder snapshot') parser.add_argument('--dec', '-d', type=str, default='dec_iter_60000.npz', help='decoder snapshot') parser.add_argument('--out', '-o', type=str, default='out', help='output dir') parser.add_argument('--input', '-i', default='sample.jpg', help='input jpg', required=True) parser.add_argument('--contour', '-c', action='store_true', help='from contour image or not') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) enc = Encoder(in_ch=3) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=3, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer if args.model: opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) # Set up a trainer updater = FacadeUpdater( models=(enc, dec, dis), iterator={}, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis}, device=args.gpu) trainer = training.Trainer(updater, (200, 'epoch'), out='generate/') chainer.serializers.load_npz(args.model, trainer) elif args.enc and args.dec: chainer.serializers.load_npz(args.enc, enc) chainer.serializers.load_npz(args.dec, dec) if not args.contour: from make_contour import get_contour_image get_contour_image(args.input) generate_image_from_contour(args.input, enc, dec, args.out)
def discriminate(): dis = Discriminator() chainer.backends.cuda.get_device_from_id(0).use() serializers.load_npz("./result/dis_iter_25000.npz", dis) dis.to_gpu() # Copy the model to the GPU paths = ["./data/70.png", "./data/35.png"] for path in paths: img = read_image_as_array(path, dtype=np.float32) img = img.reshape(-1, img.shape[0], img.shape[1], img.shape[2]).transpose(0, 3, 1, 2) img = Variable(chainer.backends.cuda.to_gpu(img)) with chainer.using_config('train', False): x = dis(img) x = chainer.backends.cuda.to_cpu(x.data) print(x)
def train(args): # setting for logging if not os.path.exists(args.log): os.mkdir(args.log) logger = logging.getLogger() logging.basicConfig(level=logging.INFO) log_path = os.path.join(args.log, 'log') file_handler = logging.FileHandler(log_path) fmt = logging.Formatter('%(asctime)s %(levelname)s %(message)s') file_handler.setFormatter(fmt) logger.addHandler(file_handler) logger.info('Arguments...') for arg, val in vars(args).items(): logger.info('{:>10} -----> {}'.format(arg, val)) x, y = gen_synthetic_data(DIM, DIM_EMB, NUM) train_x, test_x, train_y, test_y = train_test_split(x, y, test_size=0.2) valid_x, test_x, valid_y, test_y = train_test_split(test_x, test_y, test_size=0.5) gen = Generator(DIM_EMB) dis = Discriminator(DIM_EMB) gen_opt = optimizers.Adam() dis_opt = optimizers.Adam() gen_opt.setup(gen) dis_opt.setup(dis) trainer = GANTrainer((gen, dis), (gen_opt, dis_opt), logger, (valid_x, valid_y), args.epoch) trainer.fit(train_x, train_y)
def run_seq_gan(): config = Config() n_samples = config.get('n_samples') batch_size = config.get('batch_size') gen_embedding_dim = config.get('gen_embedding_dim') gen_hidden_dim = config.get('gen_hidden_dim') dis_embedding_dim = config.get('dis_embedding_dim') dis_hidden_dim = config.get('dis_hidden_dim') dataset_features = config.get('dataset_features') dataset_dtypes = config.get('dataset_dtypes') generated_features = config.get('generated_features') service_list = config.get('service_list') protocol_service_dict = config.get('protocol_service_dict') service_port_dict = config.get('service_port_dict') file_path = config.get('file_path') CUDA = torch.cuda.is_available() dataset = Traffic_Dataset(file_path, dataset_features, dataset_dtypes, generated_features, batch_size=batch_size, transform=build_input_indices) vocab_dim = dataset.vocabulary_length max_seq_len = dataset.max_seq_length train_epochs = 100 g = Generator(gen_embedding_dim, gen_hidden_dim, vocab_dim, max_seq_len, CUDA) d = Discriminator(dis_embedding_dim, dis_hidden_dim, vocab_dim, max_seq_len, CUDA) if CUDA: g.cuda() d.cuda() g_opt = optim.Adam(g.parameters()) d_opt = optim.Adagrad(d.parameters()) pre_training(g, d, g_opt, d_opt, dataset, n_samples, batch_size, CUDA) training(g, d, g_opt, d_opt, dataset, train_epochs, n_samples, batch_size, CUDA, service_list, protocol_service_dict, service_port_dict) visualize(dataset_features, generated_features)
def __init__(self, opt): self.opt = opt self.global_step = opt.load_iter self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # define net class instance self.G_net = Generator(opt).to(self.device) self.D_net = Discriminator(opt).to(self.device) if opt.load_model and opt.load_iter > 0: self._load_pre_model(self.G_net, 'G') self._load_pre_model(self.D_net, 'D') # define objectives and optimizers self.adv_loss = torch.mean # 这里的adv loss直接是真假结果的评分,真图越大越好,假图越小越好 # self.cls_loss = torch.nn.BCELoss() # ??????????? 有啥区别 self.cls_loss = F.binary_cross_entropy_with_logits self.rec_loss = torch.mean self.G_optimizer = torch.optim.Adam(self.G_net.parameters(), opt.G_lr, [opt.beta1, opt.beta2]) self.D_optimizer = torch.optim.Adam(self.D_net.parameters(), opt.D_lr, [opt.beta1, opt.beta2]) self.sample_gotten = False # 把它放在init里面,是因为它只随着类的调用初始化一次,是固定的sample self.writer = TBVisualizer(opt)
def test_discriminate_fakevideo(self): dis = Discriminator() gen = Generator() z = Variable(np.asarray(gen.make_noize(self.batchsize))) self.assertEqual((self.batchsize, 100), z.shape) x_fake = gen(z) self.assertEqual( (self.batchsize, 3, self.frame, self.height, self.width), x_fake.shape) y_fake = dis(x_fake) self.assertEqual((self.batchsize, 1), y_fake.shape)
def __init__(self, args): self.args = args self.writer = SummaryWriter(args.output_dir) self.iter_i = 1 # data transform_list = [ transforms.Resize(args.imsize), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ] train_dataset = datasets.CIFAR10( './data', train=True, transform=transforms.Compose(transform_list), download=True) self.train_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.n_workers, shuffle=True) # network self.G = Generator(args.nz, args.ngf, args.nc).to(args.device) self.D = Discriminator(args.nc, args.ndf).to(args.device) self.criterion = nn.BCELoss() self.optimizer_G = optim.Adam(self.G.parameters(), lr=args.lr, betas=(args.beta1, args.beta2)) self.optimizer_D = optim.Adam(self.D.parameters(), lr=args.lr, betas=(args.beta1, args.beta2)) self.real_label = 1 self.fake_label = 0 self.fixed_z = torch.randn((args.batch_size, args.nz, 1, 1), device=args.device)
def setup_discriminator(config): from net import Discriminator, BigBiGANDiscriminator num_z = 1 if config.generator_architecture == "dcgan" else 2 if hasattr(config, "bigan") and config.bigan: discriminator = BigBiGANDiscriminator(config.ch, config.ch * num_z, enable_blur=config.enable_blur, sn=config.sn, res=config.res_dis) else: discriminator = Discriminator(ch=config.ch, enable_blur=config.enable_blur, sn=config.sn, res=config.res_dis) return discriminator
def test_realvideo(self): dis = Discriminator() all_files = os.listdir(self.dataset) video_files = [f for f in all_files if ('mp4' in f)] train = PreprocessedDataset(paths=video_files, root=self.dataset) train_iter = chainer.iterators.SerialIterator(train, self.batchsize) batch = train_iter.next() x_real = Variable(convert.concat_examples(batch, self.gpu)) self.assertEqual( (self.batchsize, 3, self.frame, self.height, self.width), x_real.shape) y_real = dis(x_real) self.assertEqual((self.batchsize, 1), y_real.shape)
def main(): gpu_id = 0 batchsize = 10 report_keys = ["loss_dis", "loss_gen"] train_dataset = datasets.LabeledImageDataset(str(argv[1])) train_iter = iterators.SerialIterator(train_dataset, batchsize) models = [] generator = Generator() discriminator = Discriminator() opts = {} opts["opt_gen"] = setup_adam_optimizer(generator) opts["opt_dis"] = setup_adam_optimizer(discriminator) models = [generator, discriminator] chainer.cuda.get_device_from_id(gpu_id).use() print("use gpu {}".format(gpu_id)) for m in models: m.to_gpu() updater_args = { "iterator": { 'main': train_iter }, "device": gpu_id, "optimizer": opts, "models": models } output = 'result' display_interval = 20 evaluation_interval = 1000 max_iter = 10000 updater = Updater(**updater_args) trainer = training.Trainer(updater, (max_iter, 'iteration'), out=output) trainer.extend( extensions.LogReport(keys=report_keys, trigger=(display_interval, 'iteration'))) trainer.extend(extensions.PrintReport(report_keys), trigger=(display_interval, 'iteration')) trainer.extend(sample_generate(generator, output), trigger=(evaluation_interval, 'iteration'), priority=extension.PRIORITY_WRITER) trainer.run()
help='input model file path without extension') parser.add_argument('--output', '-o', required=True, type=str, help='output model file path without extension') parser.add_argument('--iter', default=100, type=int, help='number of iteration') parser.add_argument('--out_image_dir', default=None, type=str, help='output directory to output images') parser.add_argument('--dataset', '-d', default='dataset/etl9g.pkl', type=str, help='dataset file path') args = parser.parse_args() gen_model = Generator() optimizer_gen = optimizers.Adam(alpha=0.0002, beta1=0.5) optimizer_gen.setup(gen_model) optimizer_gen.add_hook(chainer.optimizer.WeightDecay(0.00001)) dis_model = Discriminator() optimizer_dis = optimizers.Adam(alpha=0.0002, beta1=0.5) optimizer_dis.setup(dis_model) optimizer_dis.add_hook(chainer.optimizer.WeightDecay(0.00001)) if args.input != None: serializers.load_hdf5(args.input + '.gen.model', gen_model) serializers.load_hdf5(args.input + '.gen.state', optimizer_gen) serializers.load_hdf5(args.input + '.dis.model', dis_model) serializers.load_hdf5(args.input + '.dis.state', optimizer_dis) if args.out_image_dir != None: if not os.path.exists(args.out_image_dir): try: os.mkdir(args.out_image_dir) except:
def main(): parser = argparse.ArgumentParser() parser.add_argument('--batchsize', '-b', type=int, default=32, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=100, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--source', default='mnist') parser.add_argument('--target', default='mnist_m') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed of z at visualization stage') parser.add_argument('--n_processes', type=int, default=16, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train gen = Generator( n_hidden=params['gen']['n_hidden'], n_resblock=params['gen']['n_resblock'], n_ch=params['gen']['n_ch'], wscale=params['gaussian_wscale'], res=params['res_image'], bn_eps=params['bn_eps']) dis = Discriminator( n_ch=params['dis']['n_ch'], wscale=params['gaussian_wscale'], bn_eps=params['bn_eps'], dr_prob=params['dis']['dropout_prob'], noise_sigma=params['dis']['noise_sigma'] ) cls = L.Classifier(DigitClassifier(n_class=params['n_class'])) if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() dis.to_gpu() cls.to_gpu() p_opt = params['optimize'] # Setup an optimizer def make_optimizer(model): optimizer = chainer.optimizers.Adam(alpha=p_opt['base_lr'], beta1=p_opt['beta1']) optimizer.setup(model) optimizer.add_hook( chainer.optimizer.WeightDecay(p_opt['weight_decay'])) return optimizer opt_gen = make_optimizer(gen) opt_dis = make_optimizer(dis) opt_cls = make_optimizer(cls) def load_dataset(name, dtype='train'): if name == 'mnist': train, _ = chainer.datasets.get_mnist(withlabel=True, ndim=3) dataset = TransformDataset(train, transform=gray2rgb) return TransformDataset(dataset, transform=scale) elif name == 'mnist_m': dataset = get_mnist_m(dtype, withlabel=True) return TransformDataset(dataset, transform=scale) else: raise NotImplementedError source = load_dataset(args.source) # from chainer.datasets import split_dataset # source, _ = split_dataset(source, split_at=1000) target_train = load_dataset(args.target, dtype='train') source_iter = MultiprocessIterator( source, args.batchsize, n_processes=args.n_processes) target_train_iter = MultiprocessIterator( target_train, args.batchsize, n_processes=args.n_processes) # Set up a trainer updater = PixelDAUpdater( models=(gen, dis, cls), iterator={'main': source_iter, 'target': target_train_iter}, optimizer={ 'gen': opt_gen, 'dis': opt_dis, 'cls': opt_cls}, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (1, 'epoch') display_interval = (10, 'iteration') for opt in [opt_gen, opt_cls, opt_dis]: trainer.extend( extensions.ExponentialShift('alpha', p_opt['alpha_decay_rate'], optimizer=opt), trigger=(p_opt['alpha_decay_steps'], 'iteration')) trainer.extend( extensions.snapshot(filename='snapshot_epoch_{.updater.epoch}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', 'cls/loss', 'validation/main/accuracy' ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) # Save two plot images to the result dir if extensions.PlotReport.available(): trainer.extend( extensions.PlotReport(['gen/loss', 'dis/loss', 'cls/loss'], 'iteration', trigger=(100, 'iteration'), file_name='loss.png')) trainer.extend( extensions.PlotReport(['validation/main/accuracy'], 'epoch', file_name='accuracy.png')) # Dump examples of generated images for every epoch trainer.extend(out_generated_image(source_iter, gen, args.gpu, args.out)) # Evaluate the model with the test dataset for each epoch target_test = load_dataset(args.target, dtype='test') target_test_iter = MultiprocessIterator( target_test, args.batchsize, n_processes=args.n_processes, repeat=False, shuffle=False) trainer.extend( extensions.Evaluator(target_test_iter, cls, device=args.gpu)) # Visualize computational graph for debug # trainer.extend(extensions.dump_graph('gen/loss', out_name='gen.dot')) # trainer.extend(extensions.dump_graph('dis/loss', out_name='dis.dot')) # trainer.extend(extensions.dump_graph('cls/loss', out_name='cls.dot')) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--device', '-d', type=str, default='-1', help='Device specifier. Either ChainerX device ' 'specifier or an integer. If non-negative integer, ' 'CuPy arrays with specified device id are used. If ' 'negative integer, NumPy arrays are used') parser.add_argument('--dataset', '-i', default='./facade/base', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', type=str, help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') group = parser.add_argument_group('deprecated arguments') group.add_argument('--gpu', '-g', dest='device', type=int, nargs='?', const=0, help='GPU ID (negative value indicates CPU)') args = parser.parse_args() if chainer.get_dtype() == numpy.float16: warnings.warn( 'This example may cause NaN in FP16 mode.', RuntimeWarning) device = chainer.get_device(args.device) if device.xp is chainerx: sys.stderr.write('This example does not support ChainerX devices.\n') sys.exit(1) print('Device: {}'.format(device)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') device.use() # Set up a neural network to train enc = Encoder(in_ch=12) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=12, out_ch=3) enc.to_device(device) dec.to_device(device) dis.to_device(device) # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = FacadeDataset(args.dataset, data_range=(1, 300)) test_d = FacadeDataset(args.dataset, data_range=(300, 379)) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) # Set up a trainer updater = FacadeUpdater( models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter}, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis}, device=device) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend(extensions.snapshot( filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend( out_image( updater, enc, dec, 5, 5, args.seed, args.out), trigger=snapshot_interval) if args.resume is not None: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
ndf = 64 lr = 0.0002 beta1 = 0.5 def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv') != -1: torch.nn.init.normal_(m.weight, 0.0, 0.02) elif classname.find('BatchNorm') != -1: torch.nn.init.normal_(m.weight, 1.0, 0.02) torch.nn.init.zeros_(m.bias) netG = Generator(ngpu, nz, ngf).to(device) netG.apply(weights_init) netD = Discriminator(ngpu, ndf).to(device) netD.apply(weights_init) criterion = nn.BCELoss() fixed_noise = torch.randn(opt.batchSize, nz, 1, 1, device=device) real_label = 1 fake_label = 0 optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999)) optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999)) schedulerD = ReduceLROnPlateau(optimizerD, mode='min', factor=0.5, patience=2, verbose=True)
def main(): FLAGS(sys.argv) running_helper = RunningHelper(FLAGS.use_mpi) global mpi_is_master mpi_is_master = running_helper.is_master # Check stage / image size / dynamic batch size / data consistency. running_helper.check_hps_consistency() # Setup Models mapping = MappingNetwork(FLAGS.ch) generator = StyleGenerator(FLAGS.ch, enable_blur=FLAGS.enable_blur) discriminator = Discriminator(ch=FLAGS.ch, enable_blur=FLAGS.enable_blur) if running_helper.keep_smoothed_gen: smoothed_generator = StyleGenerator(FLAGS.ch, enable_blur=FLAGS.enable_blur) smoothed_mapping = MappingNetwork(FLAGS.ch) models = [mapping, generator, discriminator] model_names = ['Mapping', 'Generator', 'Discriminator'] if running_helper.keep_smoothed_gen: models.append(smoothed_generator) models.append(smoothed_mapping) model_names.append('SmoothedGenerator') model_names.append('SmoothedMapping') if running_helper.device > -1: chainer.cuda.get_device_from_id(running_helper.device).use() for model in models: model.to_gpu() stage_manager = StageManager( stage_interval=running_helper.stage_interval, dynamic_batch_size=running_helper.dynamic_batch_size, make_dataset_func=running_helper.make_dataset, make_iterator_func=make_iterator_func, debug_start_instance=FLAGS.debug_start_instance) #if running_helper.is_master: # chainer.global_config.debug = True updater_args = { "models": models, "optimizer": { "map": running_helper.make_optimizer(mapping, FLAGS.adam_alpha_g / 100, FLAGS.adam_beta1, FLAGS.adam_beta2), "gen": running_helper.make_optimizer(generator, FLAGS.adam_alpha_g, FLAGS.adam_beta1, FLAGS.adam_beta2), "dis": running_helper.make_optimizer(discriminator, FLAGS.adam_alpha_d, FLAGS.adam_beta1, FLAGS.adam_beta2) }, 'stage_manager': stage_manager, 'lambda_gp': FLAGS.lambda_gp, 'smoothing': FLAGS.smoothing, 'style_mixing_rate': FLAGS.style_mixing_rate, 'use_cleargrads': running_helper.use_cleargrads, 'total_gpu': running_helper.fleet_size } updater = Updater(**updater_args) trainer = training.Trainer(updater, (lambda _trainer: _trainer.updater.stage_manager .stage_int >= FLAGS.max_stage), out=FLAGS.out) # Set up extensions if running_helper.is_master: for model, model_name in zip(models, model_names): trainer.extend(extensions.snapshot_object( model, model_name + '_{.updater.iteration}.npz'), trigger=(FLAGS.snapshot_interval, 'iteration')) trainer.extend(extensions.snapshot( filename='snapshot_iter_{.updater.iteration}.npz'), trigger=(FLAGS.snapshot_interval, 'iteration')) trainer.extend( extensions.ProgressBar(training_length=(updater.total_iteration, 'iteration'), update_interval=1)) trainer.extend(sample_generate_light(generator, mapping, FLAGS.out, rows=8, cols=8), trigger=(FLAGS.evaluation_sample_interval, 'iteration'), priority=extension.PRIORITY_WRITER) if running_helper.keep_smoothed_gen: trainer.extend(sample_generate_light(smoothed_generator, smoothed_mapping, FLAGS.out, rows=8, cols=8, subdir='preview_smoothed'), trigger=(FLAGS.evaluation_sample_interval, 'iteration'), priority=extension.PRIORITY_WRITER) report_keys = [ 'iteration', 'elapsed_time', 'stage', 'batch_size', 'image_size', 'gen/loss_adv', 'dis/loss_adv', 'dis/loss_gp' ] if FLAGS.fid_interval > 0: report_keys += 'FID' fidapi = FIDAPI(FLAGS.fid_clfs_type, FLAGS.fid_clfs_path, gpu=running_helper.device, load_real_stat=FLAGS.fid_real_stat) trainer.extend(fid_extension(fidapi, batch_generate_func( generator, mapping, trainer), seed=FLAGS.seed, report_key='FID'), trigger=(FLAGS.fid_interval, 'iteration')) if running_helper.keep_smoothed_gen: report_keys += 'S_FID' trainer.extend(fid_extension(fidapi, batch_generate_func( smoothed_generator, smoothed_mapping, trainer), seed=FLAGS.seed, report_key='S_FID'), trigger=(FLAGS.fid_interval, 'iteration')) trainer.extend( extensions.LogReport(keys=report_keys, trigger=(FLAGS.display_interval, 'iteration'))) trainer.extend(extensions.PrintReport(report_keys), trigger=(FLAGS.display_interval, 'iteration')) # Recover if possible if FLAGS.get_model_from_interation != '': resume_iteration_str = FLAGS.get_model_from_interation print('Resume from {}'.format(resume_iteration_str)) for model, model_name in zip(models, model_names): chainer.serializers.load_npz( FLAGS.out + '/' + model_name + '_%s.npz' % resume_iteration_str, model, ) chainer.serializers.load_npz( FLAGS.out + '/' + 'snapshot_iter_%s.npz' % resume_iteration_str, trainer, ) elif FLAGS.auto_resume: print("Auto Resume") candidates = [] auto_resume_dir = FLAGS.auto_resume_dir if FLAGS.auto_resume_dir != '' else FLAGS.out for fname in [ f for f in os.listdir(auto_resume_dir) if f.startswith('Generator_') and f.endswith('.npz') ]: fname = re.sub(r'^Generator_', '', fname) fname = re.sub('\.npz$', '', fname) fname_int = None try: fname_int = int(fname) except ValueError: pass if fname_int is not None: all_model_exist = True for m in model_names: if not os.path.exists(auto_resume_dir + '/' + m + '_' + fname + '.npz'): all_model_exist = False if not os.path.exists(auto_resume_dir + '/' + ('snapshot_iter_%s.npz' % fname)): all_model_exist = False if all_model_exist: candidates.append(fname) #print(candidates) candidates.sort(key=lambda _: int(_), reverse=True) if len(candidates) > 0: resume_iteration_str = candidates[0] else: resume_iteration_str = None if resume_iteration_str is not None: print('Automatic resuming: use iteration %s' % resume_iteration_str) for model, model_name in zip(models, model_names): chainer.serializers.load_npz( auto_resume_dir + '/' + model_name + '_%s.npz' % resume_iteration_str, model, ) chainer.serializers.load_npz( auto_resume_dir + '/' + 'snapshot_iter_%s.npz' % resume_iteration_str, trainer, ) # Run the training if FLAGS.enable_cuda_profiling: with cupy.cuda.profile(): trainer.run() else: #with chainer.using_config('debug', True): trainer.run() for model, model_name in zip(models, model_names): chainer.serializers.save_npz( FLAGS.out + '/' + model_name + '_latest.npz', model)
def __init__(self): warnings.filterwarnings('ignore') self.start_time = time() self.args = get_args() if self.args.checkpoint_dir_name: dir_name = self.args.checkpoint_dir_name else: dir_name = datetime.datetime.now().strftime('%y%m%d%H%M%S') self.path_to_dir = Path(__file__).resolve().parents[1] self.path_to_dir = os.path.join(self.path_to_dir, *['log', dir_name]) os.makedirs(self.path_to_dir, exist_ok=True) # tensorboard path_to_tensorboard = os.path.join(self.path_to_dir, 'tensorboard') os.makedirs(path_to_tensorboard, exist_ok=True) self.writer = SummaryWriter(path_to_tensorboard) # model saving os.makedirs(os.path.join(self.path_to_dir, 'model'), exist_ok=True) path_to_model = os.path.join(self.path_to_dir, *['model', 'model.tar']) # csv os.makedirs(os.path.join(self.path_to_dir, 'csv'), exist_ok=True) self.path_to_results_csv = os.path.join(self.path_to_dir, *['csv', 'results.csv']) path_to_args_csv = os.path.join(self.path_to_dir, *['csv', 'args.csv']) if not self.args.checkpoint_dir_name: with open(path_to_args_csv, 'a') as f: args_dict = vars(self.args) param_writer = csv.DictWriter(f, list(args_dict.keys())) param_writer.writeheader() param_writer.writerow(args_dict) # logging by hyperdash if not self.args.no_hyperdash: from hyperdash import Experiment self.exp = Experiment('Generation task on ' + self.args.dataset + ' dataset with GAN') for key in vars(self.args).keys(): exec("self.args.%s = self.exp.param('%s', self.args.%s)" % (key, key, key)) else: self.exp = None self.dataloader = get_dataloader(self.args.dataset, self.args.image_size, self.args.batch_size) sample_data = self.dataloader.__iter__().__next__()[0] image_channels = sample_data.shape[1] z = torch.randn(self.args.batch_size, self.args.z_dim) self.sample_z = z.view(z.size(0), z.size(1), 1, 1) self.Generator = Generator(self.args.z_dim, image_channels, self.args.image_size) self.Generator_optimizer = optim.Adam(self.Generator.parameters(), lr=self.args.lr_Generator, betas=(self.args.beta1, self.args.beta2)) self.writer.add_graph(self.Generator, self.sample_z) self.Generator.to(self.args.device) self.Discriminator = Discriminator(image_channels, self.args.image_size) self.Discriminator_optimizer = optim.Adam( self.Discriminator.parameters(), lr=self.args.lr_Discriminator, betas=(self.args.beta1, self.args.beta2)) self.writer.add_graph(self.Discriminator, sample_data) self.Discriminator.to(self.args.device) self.BCELoss = nn.BCELoss() self.sample_z = self.sample_z.to(self.args.device)
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='./facade/base', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=12) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=12, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = FacadeDataset(args.dataset, data_range=(1, 300)) test_d = FacadeDataset(args.dataset, data_range=(300, 379)) #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4) #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) # Set up a trainer updater = FacadeUpdater(models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter }, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_image(updater, enc, dec, 5, 5, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
class Experiment(): def __init__(self, args): self.args = args self.writer = SummaryWriter(args.output_dir) self.iter_i = 1 # data transform_list = [ transforms.Resize(args.imsize), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ] train_dataset = datasets.CIFAR10( './data', train=True, transform=transforms.Compose(transform_list), download=True) self.train_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.n_workers, shuffle=True) # network self.G = Generator(args.nz, args.ngf, args.nc).to(args.device) self.D = Discriminator(args.nc, args.ndf).to(args.device) self.criterion = nn.BCELoss() self.optimizer_G = optim.Adam(self.G.parameters(), lr=args.lr, betas=(args.beta1, args.beta2)) self.optimizer_D = optim.Adam(self.D.parameters(), lr=args.lr, betas=(args.beta1, args.beta2)) self.real_label = 1 self.fake_label = 0 self.fixed_z = torch.randn((args.batch_size, args.nz, 1, 1), device=args.device) def train(self, epoch): self.G.train() self.D.train() train_loss_G, train_loss_D = 0., 0. n_samples = 0 for data, _ in self.train_loader: batch_size = len(data) n_samples += batch_size real = data.to(self.args.device) # train D self.optimizer_D.zero_grad() # with real label = torch.full((batch_size, ), self.real_label, device=self.args.device) output_real = self.D(real) loss_D_real = self.criterion(output_real, label) loss_D_real.backward() # with fake z = torch.randn((batch_size, self.args.nz, 1, 1), device=self.args.device) fake = self.G(z) label = label.fill_(self.fake_label) output_fake = self.D(fake) loss_D_fake = self.criterion(output_fake, label) loss_D_fake.backward(retain_graph=True) loss_D = loss_D_real + loss_D_fake self.optimizer_D.step() # train G self.optimizer_G.zero_grad() label = label.fill_(self.real_label) output_fake = self.D(fake) loss_G = self.criterion(output_fake, label) loss_G.backward() self.optimizer_G.step() loss_D = loss_D.item() loss_G = loss_G.item() train_loss_D += loss_D train_loss_G += loss_G if self.iter_i % self.args.log_freq == 0: self.writer.add_scalar('Loss/D', loss_D, self.iter_i) self.writer.add_scalar('Loss/G', loss_G, self.iter_i) print('Epoch {} Train [{}/{}]: Loss/D {:.4f} Loss/G {:.4f}'. format(epoch, n_samples, len(self.train_loader.dataset), loss_D / batch_size, loss_G / batch_size)) self.iter_i += 1 dataset_size = len(self.train_loader.dataset) print('Epoch {} Train: Loss/D {:.4f} Loss/G {:.4f}'.format( epoch, train_loss_D / dataset_size, train_loss_G / dataset_size)) def test(self, epoch): self.G.eval() with torch.no_grad(): fake = self.G(self.fixed_z) grid = make_grid(fake, normalize=True).cpu() self.writer.add_image('Fake', grid, self.iter_i) # show(grid) fname = osp.join(self.args.output_dir, 'fake_epoch_{}.png'.format(epoch)) save_image(fake, fname, nrow=8) def save(self, epoch): # TODO torch.save(self.model.state_dict(), './results/checkpoint_{}.pt'.format(epoch)) def run(self): for epoch_i in range(1, 1 + self.args.epochs): self.train(epoch_i) self.test(epoch_i)
Args: labels: (LongTensor) class labels, sized [N,]. num_classes: (int) number of classes. Returns: (tensor) encoded labels, sized [N, #classes]. """ y = torch.eye(num_classes) return y[labels] # create the objects for loss function, two networks and for the two optimizers adversarial_loss = torch.nn.BCELoss() generator = Generator(latent=opt.latent, n_classes=opt.n_classes, num_filters=opt.num_filters, channels=opt.channels) discriminator = Discriminator(channels=opt.channels, num_filters=opt.num_filters, n_classes=opt.n_classes) optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.learning_rate, betas=(opt.beta_1, opt.beta_2)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.learning_rate, betas=(opt.beta_1, opt.beta_2)) # put the nets on gpu device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") generator, discriminator = generator.to(device), discriminator.to(device) print(generator) # generator.apply(weights_init) # discriminator.apply(weights_init) # start training current_epoch = 0 for epoch in range(opt.n_epochs): for i, (inputs, labels) in enumerate(dataloader): inputs = inputs.to(device)
def main(): parser = argparse.ArgumentParser(description='Train GAN') parser.add_argument('--batchsize', '-b', type=int, default=64, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=100, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='data/celebA/', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed of z at visualization stage') parser.add_argument('--snapshot_interval', type=int, default=10000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=1000, help='Interval of displaying log to console') parser.add_argument('--unrolling_steps', type=int, default=0) args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# batchsize: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train gen = Generator() dis = Discriminator(unrolling_steps=args.unrolling_steps) if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec') return optimizer opt_gen = make_optimizer(gen) opt_dis = make_optimizer(dis) # Setup a dataset all_files = os.listdir(args.dataset) image_files = [f for f in all_files if ('png' in f or 'jpg' in f)] print('{} contains {} image files'.format(args.dataset, len(image_files))) train = CelebADataset(paths=image_files, root=args.dataset) train_iter = chainer.iterators.SerialIterator(train, args.batchsize) # Set up a trainer updater = DCGANUpdater(models=(gen, dis), iterator=train_iter, optimizer={ 'gen': opt_gen, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend(extensions.snapshot( filename='snapshot_gan_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend( extensions.LogReport(trigger=display_interval, log_name='train_gan.log')) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend( extensions.PlotReport(['gen/loss', 'dis/loss'], trigger=display_interval, file_name='gan-loss.png')) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_generated_image(gen, dis, 10, 10, args.seed, args.out), trigger=snapshot_interval) if args.resume: chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='DCGAN') parser.add_argument('--batchsize', '-b', type=int, default=64, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=500, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result image') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--gennum', '-v', default=10, help='visualize image rows and columns number') parser.add_argument('--n_hidden', '-n', type=int, default=100, help='Number of hidden units (z)') parser.add_argument('--seed', type=int, default=0, help='Random seed of z at visualization stage') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# n_hidden: {}'.format(args.n_hidden)) print('# epoch: {}'.format(args.epoch)) print('') #学習モデルの作成 gen = Generator(n_hidden=args.n_hidden) dis = Discriminator() if args.gpu >= 0: #modelをGPU用に変換 chainer.cuda.get_device(args.gpu).use() gen.to_gpu() dis.to_gpu() #oputimizerのsetup def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec') return optimizer opt_gen = make_optimizer(gen) opt_dis = make_optimizer(dis) if args.dataset == '': #データセットの読み込み defaultはcifar10 train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.) else: all_files = os.listdir(args.dataset) image_files = [f for f in all_files if ('png' in f or 'jpg' in f)] print('{} contains {} image files'.format(args.dataset, len(image_files))) train = chainer.datasets\ .ImageDataset(paths=image_files, root=args.dataset) train_iter = chainer.iterators.SerialIterator(train, args.batchsize) #trainerのセットアップ updater = DCGANUpdater(models=(gen, dis), iterator=train_iter, optimizer={ 'gen': opt_gen, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.epoch}.npz'), trigger=(100, 'epoch')) trainer.extend(extensions.snapshot_object(gen, 'gen_iter_{.updater.epoch}.npz'), trigger=(100, 'epoch')) trainer.extend(extensions.snapshot_object(dis, 'dis_iter_{.updater.epoch}.npz'), trigger=(100, 'epoch')) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_generated_image(gen, dis, int(args.gennum), int(args.gennum), args.seed, args.out), trigger=snapshot_interval) if args.resume: #学習済みmodelの読み込み chainer.serializers.load_npz(args.resume, trainer) #学習の実行 trainer.run() chainer.serializers.save_npz("genmodel{0}".format(args.datasets), trainer)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--batchsize', '-b', type=int, default=50) parser.add_argument('--epoch', '-e', type=int, default=1000) parser.add_argument('--gpu', '-g', type=int, default=-1) parser.add_argument('--out', '-o', default='') parser.add_argument('--resume', '-r', default='') parser.add_argument('--n_noise', '-n', type=int, default=100) parser.add_argument('--seed', type=int, default=0) parser.add_argument('--snapshot_interval', type=int, default=1000) parser.add_argument('--display_interval', type=int, default=100) args = parser.parse_args() out_dir = 'result' if args.out != '': out_dir = '{}/{}'.format(out_dir, args.out) print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# n_hidden: {}'.format(args.n_noise)) print('# epoch: {}'.format(args.epoch)) print('# out: {}'.format(out_dir)) print('') gen = Generator(n_noise=args.n_noise, n_class=10) dis = Discriminator(n_class=10) if args.gpu >= 0: cuda.get_device_from_id(args.gpu) gen.to_gpu() dis.to_gpu() def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec') return optimizer gen_optimizer = make_optimizer(gen) dis_optimizer = make_optimizer(dis) train, _ = chainer.datasets.get_cifar10(withlabel=True) transformer = lambda data: (gen.make_noise(), ) + data train = chainer.datasets.TransformDataset(train, transformer) train_iter = chainer.iterators.SerialIterator(train, args.batchsize) updater = CGANUpdater(models=(gen, dis), iterator=train_iter, optimizer={ 'gen': gen_optimizer, 'dis': dis_optimizer }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=out_dir) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend( extensions.PlotReport(['gen/loss', 'dis/loss'], x_key='iteration', trigger=display_interval)) trainer.extend(extensions.ProgressBar(update_interval=10)) gen_func = lambda data: gen(data[0], data[1]) def data_func(gen): def _data_func(index): return (gen.make_noise(), index // 10) return _data_func trainer.extend( GenerateImage(gen_func, data_func(gen), file_name='{}/{}'.format( out_dir, 'preview/{.updater.iteration:0>8}.png'), rows=10, cols=10, seed=800, device=args.gpu, trigger=snapshot_interval)) trainer.run()
def main(): parser = argparse.ArgumentParser() parser.add_argument('--workers', type=int, help='number of data loading workers', default=2) parser.add_argument('--batch', type=int, default=50, help='input batch size') parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector') parser.add_argument('--ng_ch', type=int, default=64) parser.add_argument('--nd_ch', type=int, default=64) parser.add_argument('--epoch', type=int, default=50, help='number of epochs to train for') parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002') parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5') parser.add_argument('--outf', default='./result', help='folder to output images and model checkpoints') opt = parser.parse_args() print(opt) batch_size = opt.batch epoch_size = opt.epoch try: os.makedirs(opt.outf) except OSError: pass random.seed(0) torch.manual_seed(0) dataset = dset.SVHN(root='../svhn_root', download=True, transform=transforms.Compose([ transforms.Resize(64), transforms.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0.5), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataloader = torch.utils.data.DataLoader(dataset[:50000], batch_size=batch_size, shuffle=True, num_workers=int(opt.workers)) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print('device:', device) nz = int(opt.nz) netG = Generator().to(device) netG.apply(weights_init) print(netG) netD = Discriminator().to(device) netD.apply(weights_init) print(netD) criterion = nn.MSELoss() # criterion = nn.BCELoss() fixed_noise = torch.randn(batch_size, nz, 1, 1, device=device) real_label = 1 fake_label = 0 # setup optimizer optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=1e-5) optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=1e-5) for epoch in range(epoch_size): for itr, data in enumerate(dataloader): ############################ # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ########################### # train with real netD.zero_grad() real_image = data[0].to(device) sample_size = real_image.size(0) label = torch.full((sample_size, ), real_label, device=device) output = netD(real_image) errD_real = criterion(output, label) errD_real.backward() D_x = output.mean().item() # train with fake noise = torch.randn(sample_size, nz, 1, 1, device=device) fake_image = netG(noise) label.fill_(fake_label) output = netD(fake_image.detach()) errD_fake = criterion(output, label) errD_fake.backward() D_G_z1 = output.mean().item() errD = errD_real + errD_fake optimizerD.step() ############################ # (2) Update G network: maximize log(D(G(z))) ########################### netG.zero_grad() label.fill_(real_label) # fake labels are real for generator cost output = netD(fake_image) errG = criterion(output, label) errG.backward() D_G_z2 = output.mean().item() optimizerG.step() print( '[{}/{}][{}/{}] Loss_D: {:.3f} Loss_G: {:.3f} D(x): {:.3f} D(G(z)): {:.3f}/{:.3f}' .format(epoch + 1, opt.n_epoch, itr + 1, len(dataloader), errD.item(), errG.item(), D_x, D_G_z1, D_G_z2)) if epoch == 0 and itr == 0: vutils.save_image(real_image, '{}/real_samples.png'.format(opt.outf), normalize=True, nrow=10) fake_image = netG(fixed_noise) vutils.save_image(fake_image.detach(), '{}/fake_samples_epoch_{:03d}.png'.format( opt.outf, epoch + 1), normalize=True, nrow=10) # do checkpointing if (epoch + 1) % 100 == 0: torch.save(netG.state_dict(), '{}/netG_epoch_{}.pth'.format(opt.outf, epoch + 1)) torch.save(netD.state_dict(), '{}/netD_epoch_{}.pth'.format(opt.outf, epoch + 1))
class GAN(object): def __init__(self): warnings.filterwarnings('ignore') self.start_time = time() self.args = get_args() if self.args.checkpoint_dir_name: dir_name = self.args.checkpoint_dir_name else: dir_name = datetime.datetime.now().strftime('%y%m%d%H%M%S') self.path_to_dir = Path(__file__).resolve().parents[1] self.path_to_dir = os.path.join(self.path_to_dir, *['log', dir_name]) os.makedirs(self.path_to_dir, exist_ok=True) # tensorboard path_to_tensorboard = os.path.join(self.path_to_dir, 'tensorboard') os.makedirs(path_to_tensorboard, exist_ok=True) self.writer = SummaryWriter(path_to_tensorboard) # model saving os.makedirs(os.path.join(self.path_to_dir, 'model'), exist_ok=True) path_to_model = os.path.join(self.path_to_dir, *['model', 'model.tar']) # csv os.makedirs(os.path.join(self.path_to_dir, 'csv'), exist_ok=True) self.path_to_results_csv = os.path.join(self.path_to_dir, *['csv', 'results.csv']) path_to_args_csv = os.path.join(self.path_to_dir, *['csv', 'args.csv']) if not self.args.checkpoint_dir_name: with open(path_to_args_csv, 'a') as f: args_dict = vars(self.args) param_writer = csv.DictWriter(f, list(args_dict.keys())) param_writer.writeheader() param_writer.writerow(args_dict) # logging by hyperdash if not self.args.no_hyperdash: from hyperdash import Experiment self.exp = Experiment('Generation task on ' + self.args.dataset + ' dataset with GAN') for key in vars(self.args).keys(): exec("self.args.%s = self.exp.param('%s', self.args.%s)" % (key, key, key)) else: self.exp = None self.dataloader = get_dataloader(self.args.dataset, self.args.image_size, self.args.batch_size) sample_data = self.dataloader.__iter__().__next__()[0] image_channels = sample_data.shape[1] z = torch.randn(self.args.batch_size, self.args.z_dim) self.sample_z = z.view(z.size(0), z.size(1), 1, 1) self.Generator = Generator(self.args.z_dim, image_channels, self.args.image_size) self.Generator_optimizer = optim.Adam(self.Generator.parameters(), lr=self.args.lr_Generator, betas=(self.args.beta1, self.args.beta2)) self.writer.add_graph(self.Generator, self.sample_z) self.Generator.to(self.args.device) self.Discriminator = Discriminator(image_channels, self.args.image_size) self.Discriminator_optimizer = optim.Adam( self.Discriminator.parameters(), lr=self.args.lr_Discriminator, betas=(self.args.beta1, self.args.beta2)) self.writer.add_graph(self.Discriminator, sample_data) self.Discriminator.to(self.args.device) self.BCELoss = nn.BCELoss() self.sample_z = self.sample_z.to(self.args.device) def train(self): self.train_hist = {} self.train_hist['Generator_loss'] = 0.0 self.train_hist['Discriminator_loss'] = 0.0 # real ---> y = 1 # fake ---> y = 0 self.y_real = torch.ones(self.args.batch_size, 1).to(self.args.device) self.y_fake = torch.zeros(self.args.batch_size, 1).to(self.args.device) self.Discriminator.train() global_step = 0 # -----training ----- for epoch in range(1, self.args.n_epoch + 1): self.Generator.train() for idx, (x, _) in enumerate(self.dataloader): if idx == self.dataloader.dataset.__len__( ) // self.args.batch_size: break z = torch.randn(self.args.batch_size, self.args.z_dim) z = z.view(z.size(0), z.size(1), 1, 1) z = z.to(self.args.device) x = x.to(self.args.device) # ----- update Discriminator ----- # minimize: -{ log[D(x)] + log[1-D(G(z))] } self.Discriminator_optimizer.zero_grad() # real # ---> log[D(x)] Discriminator_real, _ = self.Discriminator(x) Discriminator_real_loss = self.BCELoss(Discriminator_real, self.y_real) # fake # ---> log[1-D(G(z))] Discriminator_fake, _ = self.Discriminator(self.Generator(z)) Discriminator_fake_loss = self.BCELoss(Discriminator_fake, self.y_fake) Discriminator_loss = Discriminator_real_loss + Discriminator_fake_loss self.train_hist[ 'Discriminator_loss'] = Discriminator_loss.item() Discriminator_loss.backward() self.Discriminator_optimizer.step() # ----- update Generator ----- # As stated in the original paper, # we want to train the Generator # by minimizing log(1−D(G(z))) # in an effort to generate better fakes. # As mentioned, this was shown by Goodfellow # to not provide sufficient gradients, # especially early in the learning process. # As a fix, we instead wish to maximize log(D(G(z))). # ---> minimize: -log[D(G(z))] self.Generator_optimizer.zero_grad() Discriminator_fake, _ = self.Discriminator(self.Generator(z)) Generator_loss = self.BCELoss(Discriminator_fake, self.y_real) self.train_hist['Generator_loss'] = Generator_loss.item() Generator_loss.backward() self.Generator_optimizer.step() # ----- logging by tensorboard, csv and hyperdash # tensorboard self.writer.add_scalar('loss/Generator_loss', Generator_loss.item(), global_step) self.writer.add_scalar('loss/Discriminator_loss', Discriminator_loss.item(), global_step) # csv with open(self.path_to_results_csv, 'a') as f: result_writer = csv.DictWriter( f, list(self.train_hist.keys())) if epoch == 1 and idx == 0: result_writer.writeheader() result_writer.writerow(self.train_hist) # hyperdash if self.exp: self.exp.metric('Generator loss', Generator_loss.item()) self.exp.metric('Discriminator loss', Discriminator_loss.item()) if (idx % 10) == 0: self._plot_sample(global_step) global_step += 1 elapsed_time = time() - self.start_time print('\nTraining Finish, elapsed time ---> %f' % (elapsed_time)) def _plot_sample(self, global_step): with torch.no_grad(): total_n_sample = min(self.args.n_sample, self.args.batch_size) image_frame_dim = int(np.floor(np.sqrt(total_n_sample))) samples = self.Generator(self.sample_z) samples = samples.cpu().data.numpy().transpose(0, 2, 3, 1) samples = (samples + 1) / 2 fig = plt.figure(figsize=(24, 15)) for i in range(image_frame_dim * image_frame_dim): ax = fig.add_subplot( image_frame_dim, image_frame_dim * 2, (int(i / image_frame_dim) + 1) * image_frame_dim + i + 1, xticks=[], yticks=[]) if samples[i].shape[2] == 3: ax.imshow(samples[i]) else: ax.imshow(samples[i][:, :, 0], cmap='gray') self.writer.add_figure('sample images generated by GAN', fig, global_step)
help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=5) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=5, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec)
def main(): parser = argparse.ArgumentParser( description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=500, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--out', '-o', default='result_dehighlight', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=10000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=10000, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=3) dec = Decoder(out_ch=1) dis = Discriminator(in_ch=3, out_ch=1) gen = Generator(in_ch=3, out_ch=1) serializers.load_npz("depro.npz", depro) gen.encoder = enc gen.decoder = dec if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() gen.to_gpu() depro.disable_update() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) #will fix train_d = NyuDataset("E:/nyu_depth_v2_labeled.mat", startnum=0, endnum=1000) test_d = NyuDataset("E:/nyu_depth_v2_labeled.mat", startnum=1000, endnum=1449) #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4) #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) test_iter2 = chainer.iterators.SerialIterator(test_d, args.batchsize, repeat=False, shuffle=False) # Set up a trainer updater = PicUpdater(models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter }, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss_enc', 'dec/loss_dec', 'dis/loss_dis', "validation/main/loss" ]), trigger=display_interval) trainer.extend(extensions.Evaluator(test_iter2, gen, device=args.gpu)) trainer.extend( extensions.PlotReport(['dec/loss_dec'], x_key='iteration', file_name='dec_loss.png', trigger=display_interval)) trainer.extend( extensions.PlotReport(['dis/loss_dis'], x_key='iteration', file_name='dis_loss.png', trigger=display_interval)) trainer.extend( extensions.PlotReport(["validation/main/loss"], x_key='iteration', file_name='gen_loss.png', trigger=display_interval)) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_image(updater, depro, enc, dec, 3, 3, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='Train script') parser.add_argument('--data_path', type=str, default="data/datasets/cifar10/train", help='dataset directory path') parser.add_argument('--class_name', '-class', type=str, default='all_class', help='class name (default: all_class(str))') parser.add_argument('--num_workers', type=int, default=4, help='number of parallel data loading processes') parser.add_argument('--batchsize', '-b', type=int, default=16) parser.add_argument('--max_iter', '-m', type=int, default=40000) parser.add_argument('--gpu', '-g', type=int, default=0, help='GPU ID (negative value indicates CPU)') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--snapshot_interval', type=int, default=2500, help='Interval of snapshot') parser.add_argument('--evaluation_interval', type=int, default=5000, help='Interval of evaluation') parser.add_argument('--out_image_interval', type=int, default=1250, help='Interval of evaluation') parser.add_argument('--stage_interval', type=int, default=40000, help='Interval of stage progress') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') parser.add_argument( '--n_dis', type=int, default=1, help='number of discriminator update per generator update') parser.add_argument('--lam', type=float, default=10, help='gradient penalty') parser.add_argument('--gamma', type=float, default=750, help='gradient penalty') parser.add_argument('--pooling_comp', type=float, default=1.0, help='compensation') parser.add_argument('--pretrained_generator', type=str, default="") parser.add_argument('--pretrained_discriminator', type=str, default="") parser.add_argument('--initial_stage', type=float, default=0.0) parser.add_argument('--generator_smoothing', type=float, default=0.999) args = parser.parse_args() record_setting(args.out) report_keys = [ "stage", "loss_dis", "loss_gp", "loss_gen", "g", "inception_mean", "inception_std", "FID" ] max_iter = args.max_iter if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu).use() generator = Generator() generator_smooth = Generator() discriminator = Discriminator(pooling_comp=args.pooling_comp) # select GPU if args.gpu >= 0: generator.to_gpu() generator_smooth.to_gpu() discriminator.to_gpu() print("use gpu {}".format(args.gpu)) if args.pretrained_generator != "": chainer.serializers.load_npz(args.pretrained_generator, generator) if args.pretrained_discriminator != "": chainer.serializers.load_npz(args.pretrained_discriminator, discriminator) copy_param(generator_smooth, generator) # Setup an optimizer def make_optimizer(model, alpha=0.001, beta1=0.0, beta2=0.99): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1, beta2=beta2) optimizer.setup(model) # optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_gen = make_optimizer(generator) opt_dis = make_optimizer(discriminator) if args.class_name == 'all_class': data_path = args.data_path one_class_flag = False else: data_path = os.path.join(args.data_path, args.class_name) one_class_flag = True train_dataset = ImageDataset(data_path, one_class_flag=one_class_flag) train_iter = chainer.iterators.MultiprocessIterator( train_dataset, args.batchsize, n_processes=args.num_workers) # Set up a trainer updater = Updater(models=(generator, discriminator, generator_smooth), iterator={'main': train_iter}, optimizer={ 'opt_gen': opt_gen, 'opt_dis': opt_dis }, device=args.gpu, n_dis=args.n_dis, lam=args.lam, gamma=args.gamma, smoothing=args.generator_smoothing, initial_stage=args.initial_stage, stage_interval=args.stage_interval) trainer = training.Trainer(updater, (max_iter, 'iteration'), out=args.out) trainer.extend(extensions.snapshot_object( generator, 'generator_{.updater.iteration}.npz'), trigger=(args.snapshot_interval, 'iteration')) trainer.extend(extensions.snapshot_object( generator_smooth, 'generator_smooth_{.updater.iteration}.npz'), trigger=(args.snapshot_interval, 'iteration')) trainer.extend(extensions.snapshot_object( discriminator, 'discriminator_{.updater.iteration}.npz'), trigger=(args.snapshot_interval, 'iteration')) trainer.extend( extensions.LogReport(keys=report_keys, trigger=(args.display_interval, 'iteration'))) trainer.extend(extensions.PrintReport(report_keys), trigger=(args.display_interval, 'iteration')) trainer.extend(sample_generate(generator_smooth, args.out), trigger=(args.out_image_interval, 'iteration'), priority=extension.PRIORITY_WRITER) trainer.extend(sample_generate_light(generator_smooth, args.out), trigger=(args.evaluation_interval // 10, 'iteration'), priority=extension.PRIORITY_WRITER) trainer.extend(calc_inception(generator_smooth), trigger=(args.evaluation_interval, 'iteration'), priority=extension.PRIORITY_WRITER) trainer.extend(calc_FID(generator_smooth), trigger=(args.evaluation_interval, 'iteration'), priority=extension.PRIORITY_WRITER) trainer.extend(extensions.ProgressBar(update_interval=10)) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='chainer implementation of pix2pix') parser.add_argument('--batchsize', '-b', type=int, default=1, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=200, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='./facade/base', help='Directory of image files.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--seed', type=int, default=0, help='Random seed') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train enc = Encoder(in_ch=12) dec = Decoder(out_ch=3) dis = Discriminator(in_ch=12, out_ch=3) if args.gpu >= 0: chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current enc.to_gpu() # Copy the model to the GPU dec.to_gpu() dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.00001), 'hook_dec') return optimizer opt_enc = make_optimizer(enc) opt_dec = make_optimizer(dec) opt_dis = make_optimizer(dis) train_d = FacadeDataset(args.dataset, data_range=(1,300)) test_d = FacadeDataset(args.dataset, data_range=(300,379)) #train_iter = chainer.iterators.MultiprocessIterator(train_d, args.batchsize, n_processes=4) #test_iter = chainer.iterators.MultiprocessIterator(test_d, args.batchsize, n_processes=4) train_iter = chainer.iterators.SerialIterator(train_d, args.batchsize) test_iter = chainer.iterators.SerialIterator(test_d, args.batchsize) # Set up a trainer updater = FacadeUpdater( models=(enc, dec, dis), iterator={ 'main': train_iter, 'test': test_iter}, optimizer={ 'enc': opt_enc, 'dec': opt_dec, 'dis': opt_dis}, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend(extensions.snapshot( filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( enc, 'enc_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dec, 'dec_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'enc/loss', 'dec/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend( out_image( updater, enc, dec, 5, 5, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser() parser.add_argument('--workers', type=int, help='number of data loading workers', default=2) parser.add_argument('--batch_size', type=int, default=50, help='input batch size') parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector') parser.add_argument('--nch_g', type=int, default=64) parser.add_argument('--nch_d', type=int, default=64) parser.add_argument('--n_epoch', type=int, default=200, help='number of epochs to train for') parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002') parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5') parser.add_argument('--outf', default='./result_lsgan', help='folder to output images and model checkpoints') opt = parser.parse_args() print(opt) try: os.makedirs(opt.outf) except OSError: pass # 乱数のシード(種)を固定 random.seed(0) np.random.seed(0) torch.manual_seed(0) # STL-10のトレーニングデータセットとテストデータセットを読み込む trainset = dset.STL10(root='./dataset/stl10_root', download=True, split='train+unlabeled', transform=transforms.Compose([ transforms.RandomResizedCrop(64, scale=(88 / 96, 1.0), ratio=(1., 1.)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.05, contrast=0.05, saturation=0.05, hue=0.05), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) # ラベルを使用しないのでラベルなしを混在した'train+unlabeled'を読み込む testset = dset.STL10(root='./dataset/stl10_root', download=True, split='test', transform=transforms.Compose([ transforms.RandomResizedCrop(64, scale=(88 / 96, 1.0), ratio=(1., 1.)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.05, contrast=0.05, saturation=0.05, hue=0.05), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataset = trainset + testset # STL-10のトレーニングデータセットとテストデータセットを合わせて訓練データとする # 訓練データをセットしたデータローダーを作成する dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=int(opt.workers)) # 学習に使用するデバイスを得る。可能ならGPUを使用する device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print('device:', device) # 生成器G。ランダムベクトルから贋作画像を生成する netG = Generator(nz=opt.nz, nch_g=opt.nch_g).to(device) netG.apply(weights_init) # weights_init関数で初期化 print(netG) # 識別器D。画像が、元画像か贋作画像かを識別する netD = Discriminator(nch_d=opt.nch_d).to(device) netD.apply(weights_init) print(netD) criterion = nn.MSELoss() # 損失関数は平均二乗誤差損失 # オプティマイザ−のセットアップ optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=1e-5) # 識別器D用 optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=1e-5) # 生成器G用 fixed_noise = torch.randn(opt.batch_size, opt.nz, 1, 1, device=device) # 確認用の固定したノイズ # 学習のループ for epoch in range(opt.n_epoch): for itr, data in enumerate(dataloader): real_image = data[0].to(device) # 元画像 sample_size = real_image.size(0) # 画像枚数 noise = torch.randn(sample_size, opt.nz, 1, 1, device=device) # 正規分布からノイズを生成 real_target = torch.full((sample_size, ), 1., device=device) # 元画像に対する識別信号の目標値「1」 fake_target = torch.full((sample_size, ), 0., device=device) # 贋作画像に対する識別信号の目標値「0」 ############################ # 識別器Dの更新 ########################### netD.zero_grad() # 勾配の初期化 output = netD(real_image) # 識別器Dで元画像に対する識別信号を出力 errD_real = criterion(output, real_target) # 元画像に対する識別信号の損失値 D_x = output.mean().item() fake_image = netG(noise) # 生成器Gでノイズから贋作画像を生成 output = netD(fake_image.detach()) # 識別器Dで元画像に対する識別信号を出力 errD_fake = criterion(output, fake_target) # 贋作画像に対する識別信号の損失値 D_G_z1 = output.mean().item() errD = errD_real + errD_fake # 識別器Dの全体の損失 errD.backward() # 誤差逆伝播 optimizerD.step() # Dのパラメーターを更新 ############################ # 生成器Gの更新 ########################### netG.zero_grad() # 勾配の初期化 output = netD(fake_image) # 更新した識別器Dで改めて贋作画像に対する識別信号を出力 errG = criterion( output, real_target) # 生成器Gの損失値。Dに贋作画像を元画像と誤認させたいため目標値は「1」 errG.backward() # 誤差逆伝播 D_G_z2 = output.mean().item() optimizerG.step() # Gのパラメータを更新 print( '[{}/{}][{}/{}] Loss_D: {:.3f} Loss_G: {:.3f} D(x): {:.3f} D(G(z)): {:.3f}/{:.3f}' .format(epoch + 1, opt.n_epoch, itr + 1, len(dataloader), errD.item(), errG.item(), D_x, D_G_z1, D_G_z2)) if epoch == 0 and itr == 0: # 初回に元画像を保存する vutils.save_image(real_image, '{}/real_samples.png'.format(opt.outf), normalize=True, nrow=10) ############################ # 確認用画像の生成 ############################ fake_image = netG(fixed_noise) # 1エポック終了ごとに確認用の贋作画像を生成する vutils.save_image(fake_image.detach(), '{}/fake_samples_epoch_{:03d}.png'.format( opt.outf, epoch + 1), normalize=True, nrow=10) ############################ # モデルの保存 ############################ if (epoch + 1) % 50 == 0: # 50エポックごとにモデルを保存する torch.save(netG.state_dict(), '{}/netG_epoch_{}.pth'.format(opt.outf, epoch + 1)) torch.save(netD.state_dict(), '{}/netD_epoch_{}.pth'.format(opt.outf, epoch + 1))
def main(): parser = argparse.ArgumentParser(description='Chainer example: DCGAN') parser.add_argument('--batchsize', '-b', type=int, default=50, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=1000, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='', help='Directory of image files. Default is cifar-10.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--n_hidden', '-n', type=int, default=100, help='Number of hidden units (z)') parser.add_argument('--seed', type=int, default=0, help='Random seed of z at visualization stage') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# n_hidden: {}'.format(args.n_hidden)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train gen = Generator(n_hidden=args.n_hidden) dis = Discriminator() if args.gpu >= 0: # Make a specified GPU current chainer.backends.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() # Copy the model to the GPU dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer_hooks.WeightDecay(0.0001), 'hook_dec') return optimizer opt_gen = make_optimizer(gen) opt_dis = make_optimizer(dis) if args.dataset == '': # Load the CIFAR10 dataset if args.dataset is not specified train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.) else: all_files = os.listdir(args.dataset) image_files = [f for f in all_files if ('png' in f or 'jpg' in f)] print('{} contains {} image files'.format(args.dataset, len(image_files))) train = chainer.datasets\ .ImageDataset(paths=image_files, root=args.dataset) train_iter = chainer.iterators.SerialIterator(train, args.batchsize) # Set up a trainer updater = DCGANUpdater(models=(gen, dis), iterator=train_iter, optimizer={ 'gen': opt_gen, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_generated_image(gen, dis, 10, 10, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='ChainerMN example: DCGAN') parser.add_argument('--batchsize', '-b', type=int, default=50, help='Number of images in each mini-batch') parser.add_argument('--communicator', type=str, default='hierarchical', help='Type of communicator') parser.add_argument('--epoch', '-e', type=int, default=1000, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', action='store_true', help='Use GPU') parser.add_argument('--dataset', '-i', default='', help='Directory of image files. Default is cifar-10.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--gen_model', '-r', default='', help='Use pre-trained generator for training') parser.add_argument('--dis_model', '-d', default='', help='Use pre-trained discriminator for training') parser.add_argument('--n_hidden', '-n', type=int, default=100, help='Number of hidden units (z)') parser.add_argument('--seed', type=int, default=0, help='Random seed of z at visualization stage') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() # Prepare ChainerMN communicator. if args.gpu: if args.communicator == 'naive': print("Error: 'naive' communicator does not support GPU.\n") exit(-1) comm = chainermn.create_communicator(args.communicator) device = comm.intra_rank else: if args.communicator != 'naive': print('Warning: using naive communicator ' 'because only naive supports CPU-only execution') comm = chainermn.create_communicator('naive') device = -1 if comm.rank == 0: print('==========================================') print('Num process (COMM_WORLD): {}'.format(comm.size)) if args.gpu: print('Using GPUs') print('Using {} communicator'.format(args.communicator)) print('Num hidden unit: {}'.format(args.n_hidden)) print('Num Minibatch-size: {}'.format(args.batchsize)) print('Num epoch: {}'.format(args.epoch)) print('==========================================') # Set up a neural network to train gen = Generator(n_hidden=args.n_hidden) dis = Discriminator() if device >= 0: # Make a specified GPU current chainer.cuda.get_device_from_id(device).use() gen.to_gpu() # Copy the model to the GPU dis.to_gpu() # Setup an optimizer def make_optimizer(model, comm, alpha=0.0002, beta1=0.5): # Create a multi node optimizer from a standard Chainer optimizer. optimizer = chainermn.create_multi_node_optimizer( chainer.optimizers.Adam(alpha=alpha, beta1=beta1), comm) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec') return optimizer opt_gen = make_optimizer(gen, comm) opt_dis = make_optimizer(dis, comm) # Split and distribute the dataset. Only worker 0 loads the whole dataset. # Datasets of worker 0 are evenly split and distributed to all workers. if comm.rank == 0: if args.dataset == '': # Load the CIFAR10 dataset if args.dataset is not specified train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.) else: all_files = os.listdir(args.dataset) image_files = [f for f in all_files if ('png' in f or 'jpg' in f)] print('{} contains {} image files' .format(args.dataset, len(image_files))) train = chainer.datasets\ .ImageDataset(paths=image_files, root=args.dataset) else: train = None train = chainermn.scatter_dataset(train, comm) train_iter = chainer.iterators.SerialIterator(train, args.batchsize) # Set up a trainer updater = DCGANUpdater( models=(gen, dis), iterator=train_iter, optimizer={ 'gen': opt_gen, 'dis': opt_dis}, device=device) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) # Some display and output extensions are necessary only for one worker. # (Otherwise, there would just be repeated outputs.) if comm.rank == 0: snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') # Save only model parameters. # `snapshot` extension will save all the trainer module's attribute, # including `train_iter`. # However, `train_iter` depends on scattered dataset, which means that # `train_iter` may be different in each process. # Here, instead of saving whole trainer module, only the network models # are saved. trainer.extend(extensions.snapshot_object( gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', 'elapsed_time', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend( out_generated_image( gen, dis, 10, 10, args.seed, args.out), trigger=snapshot_interval) # Start the training using pre-trained model, saved by snapshot_object if args.gen_model: chainer.serializers.load_npz(args.gen_model, gen) if args.dis_model: chainer.serializers.load_npz(args.dis_model, dis) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='Chainer example: DCGAN') parser.add_argument('--batchsize', '-b', type=int, default=50, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=1000, help='Number of sweeps over the dataset to train') parser.add_argument('--device', '-d', type=str, default='-1', help='Device specifier. Either ChainerX device ' 'specifier or an integer. If non-negative integer, ' 'CuPy arrays with specified device id are used. If ' 'negative integer, NumPy arrays are used') parser.add_argument('--dataset', '-i', default='', help='Directory of image files. Default is cifar-10.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', type=str, help='Resume the training from snapshot') parser.add_argument('--n_hidden', '-n', type=int, default=100, help='Number of hidden units (z)') parser.add_argument('--seed', type=int, default=0, help='Random seed of z at visualization stage') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') group = parser.add_argument_group('deprecated arguments') group.add_argument('--gpu', '-g', dest='device', type=int, nargs='?', const=0, help='GPU ID (negative value indicates CPU)') args = parser.parse_args() device = chainer.get_device(args.device) device.use() print('Device: {}'.format(device)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# n_hidden: {}'.format(args.n_hidden)) print('# epoch: {}'.format(args.epoch)) print('') # Set up a neural network to train gen = Generator(n_hidden=args.n_hidden) dis = Discriminator() gen.to_device(device) # Copy the model to the device dis.to_device(device) # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook( chainer.optimizer_hooks.WeightDecay(0.0001), 'hook_dec') return optimizer opt_gen = make_optimizer(gen) opt_dis = make_optimizer(dis) if args.dataset == '': # Load the CIFAR10 dataset if args.dataset is not specified train, _ = chainer.datasets.get_cifar10(withlabel=False, scale=255.) else: all_files = os.listdir(args.dataset) image_files = [f for f in all_files if ('png' in f or 'jpg' in f)] print('{} contains {} image files' .format(args.dataset, len(image_files))) train = chainer.datasets\ .ImageDataset(paths=image_files, root=args.dataset) # Setup an iterator train_iter = chainer.iterators.SerialIterator(train, args.batchsize) # Setup an updater updater = DCGANUpdater( models=(gen, dis), iterator=train_iter, optimizer={ 'gen': opt_gen, 'dis': opt_dis}, device=device) # Setup a trainer trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend( out_generated_image( gen, dis, 10, 10, args.seed, args.out), trigger=snapshot_interval) if args.resume is not None: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run()
def main(): parser = argparse.ArgumentParser(description='Chainer example: DCGAN') parser.add_argument('--batchsize', '-b', type=int, default=50, help='Number of images in each mini-batch') parser.add_argument('--epoch', '-e', type=int, default=1000, help='Number of sweeps over the dataset to train') parser.add_argument('--gpu', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') parser.add_argument('--dataset', '-i', default='filelist.txt', help='List file of image files.') parser.add_argument('--dataset_root', default='.', help='Root directory to retrieve images from.') parser.add_argument('--out', '-o', default='result', help='Directory to output the result') parser.add_argument('--resume', '-r', default='', help='Resume the training from snapshot') parser.add_argument('--n_hidden', '-n', type=int, default=100, help='Number of hidden units (z)') parser.add_argument('--seed', type=int, default=0, help='Random seed of z at visualization stage') parser.add_argument('--snapshot_interval', type=int, default=1000, help='Interval of snapshot') parser.add_argument('--display_interval', type=int, default=100, help='Interval of displaying log to console') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(args.batchsize)) print('# n_hidden: {}'.format(args.n_hidden)) print('# epoch: {}'.format(args.epoch)) # Set up a neural network to train gen = Generator(n_hidden=args.n_hidden) dis = Discriminator() if args.gpu >= 0: # Make a specified GPU current chainer.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() # Copy the model to the GPU dis.to_gpu() # Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec') return optimizer opt_gen = make_optimizer(gen) opt_dis = make_optimizer(dis) # Set up dataset train = chainer.datasets.LabeledImageDataset(args.dataset, root=args.dataset_root) print('# data-size: {}'.format(len(train))) print('# data-shape: {}'.format(train[0][0].shape)) print('') train_iter = chainer.iterators.SerialIterator(train, args.batchsize) # Set up a trainer updater = GANUpdater(models=(gen, dis), iterator=train_iter, optimizer={ 'gen': opt_gen, 'dis': opt_dis }, device=args.gpu) trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out) snapshot_interval = (args.snapshot_interval, 'iteration') display_interval = (args.display_interval, 'iteration') trainer.extend( extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'gen/loss', 'dis/loss', ]), trigger=display_interval) trainer.extend( extensions.PlotReport(['gen/loss', 'dis/loss'], x_key='iteration', trigger=display_interval)) trainer.extend(extensions.ProgressBar(update_interval=10)) trainer.extend(out_generated_image(gen, dis, 10, 10, args.seed, args.out), trigger=snapshot_interval) if args.resume: # Resume from a snapshot chainer.serializers.load_npz(args.resume, trainer) # Run the training trainer.run() # Save generator model if args.gpu >= 0: gen.to_cpu() chainer.serializers.save_npz(os.path.join(args.out, 'gen.npz'), gen)