def test_pix2pix(pathdataset: str, pathoutput: str, pathmodel: str) -> str: """ Use the Pix2Pix trained model to generate ("translate") images from the Test dataset Args: pathdataset (string) : Full path name to the test dataset (TFRecords file) pathoutput (string) : Full path name to the Output ("translated") images pathmodel (string) : Full path name to the Pix2Pix model checkpoint Returns: output_path (string) : Path to generated/translated images """ # ------------------------------ # In order to be able to convert # a Python Function directly # into a Kubeflow component, # we need to move the python # includes inside that python # function. # ------------------------------ import numpy as np import tensorflow as tf # The below tag comment is used by a tool script from this project to automatically nest # the python code of the imports function tagged KFP-NEST-IMPORT, just right after this tag. # This is only usefull when using the SDK's kfp.components.func_to_container_op method, # which allows to convert a Python function to a pipeline component and returns a factory function # #KFP-NEST-HERE # ------------------------------ # Define some hyperparameters # (Not managed as Kubeflow # pipeline parameters) # ------------------------------ BATCH_SIZE = 1 # ------------------------------ # Restore the Pix2Pix GAN # # Rebuild the Generator # (Skip the Discriminator # and Optimizers) # ------------------------------ generator = Generator() # Layer objects in TensorFlow may delay the creation of variables # to their first call, when input shapes are available. _ = generator( np.random.uniform(-127.5, 127.5, (1, 256, 256, 3)).astype(np.float32)) # ------------------------------ # Restore the last Checkpoint # (Skip the Discriminator # and Optimizers) # ----------------------------- # Configure the model checkpoints saving checkpoint = tf.train.Checkpoint(generator=generator) # Restoring the latest checkpoint manager = tf.train.CheckpointManager(checkpoint, directory=pathmodel, max_to_keep=1, checkpoint_name='ckpt') _ = checkpoint.restore(manager.latest_checkpoint).expect_partial() if manager.latest_checkpoint: print("Model restored from {}".format(manager.latest_checkpoint)) else: print("[WARNING] Initializing model from scratch.") # ------------------------------ # One step Inference # nested Helper function # ------------------------------ #@tf.function # Compile function into a graph for faster execution # TODO : Using Autograd does not seem to work with Checkpoint Restore ?? def generate_image(img): """ Generate an image using the Pix2pix Generator """ return generator(img, training=False) # ------------------------------ # Extract the Test Dataset # from the TFRecords file # ------------------------------ test_dataset = get_dataset(pathdataset, BATCH_SIZE, shuffle=False) # ------------------------------ # Loop over the training batches # ------------------------------ for image_features in test_dataset: # Loop over all images in the batch for record_idx in range(BATCH_SIZE): # ------------------------------ # Extract the individual # features and reconstruct # the input images to feed # the Neural Networks # ------------------------------ a_image, b_image, file_name = decode_tfrecords_example( image_features, record_idx) print("Processing image ", file_name) # ------------------------------ # Apply same Data preparation, # as for the training, on # b_image - the input image - # (but without Data augmentation) # in order to feed into the # Generator Neural Network Input # ------------------------------ img_b = normalize(b_image) #[ height, width, OUTPUT_CHANNELS] img_b = tf.image.resize(img_b, [256, 256]) # Add a batch dimension to have 4D Tensor images to feed the Neural Network input img_b = tf.expand_dims(img_b, axis=0) # [ 1, height, width, channel] # ------------------------------ # Generate output image # - One step inference - # ------------------------------ fake_b = generate_image(img_b) # ------------------------------ # Save the source, target # and generated images to disk # (incl. decoding if necessary) # ------------------------------ img_name = pathoutput + "/img_a-" + file_name + ".jpg" save_image(a_image, img_name) img_name = pathoutput + "/img_b-" + file_name + ".jpg" save_image(b_image, img_name) img_name = pathoutput + "/fake_b-" + file_name + ".jpg" save_image(fake_b, img_name) print("End of inference on Test Dataset") # ------------------------------ # Write the Output of the # Kubeflow Pipeline Component # ------------------------------ try: # This works only inside Docker containers with open('/output.txt', 'w') as f: f.write(pathoutput) except PermissionError: pass return pathoutput
def project(path_ckpt, path_files, step=1000): device = "cuda" parser = argparse.ArgumentParser() parser.add_argument('-f', type=str, help='jup kernel') # parser.add_argument("--ckpt", type=str, required=True) parser.add_argument("--size", type=int, default=512) parser.add_argument("--lr_rampup", type=float, default=0.05) parser.add_argument("--lr_rampdown", type=float, default=0.25) parser.add_argument("--lr", type=float, default=0.1) parser.add_argument("--noise", type=float, default=0.05) parser.add_argument("--noise_ramp", type=float, default=0.75) parser.add_argument("--step", type=int, default=1000) parser.add_argument("--noise_regularize", type=float, default=1e5) parser.add_argument("--mse", type=float, default=0) # parser.add_argument("--w_plus", action="store_true") # parser.add_argument("files", metavar="FILES", nargs="+") args = parser.parse_args() args.ckpt = path_ckpt args.files = path_files args.w_plus = False args.step = step n_mean_latent = 10000 resize = min(args.size, 256) transform = transforms.Compose([ transforms.Resize(resize), transforms.CenterCrop(resize), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) imgs = [] for imgfile in args.files: img = transform(Image.open(imgfile).convert("RGB")) imgs.append(img) imgs = torch.stack(imgs, 0).to(device) g_ema = Generator(args.size, 512, 8) g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False) g_ema.eval() g_ema = g_ema.to(device) with torch.no_grad(): noise_sample = torch.randn(n_mean_latent, 512, device=device) latent_out = g_ema.style(noise_sample) latent_mean = latent_out.mean(0) latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5 percept = lpips.PerceptualLoss( model="net-lin", net="vgg", use_gpu=device.startswith("cuda") ) noises_single = g_ema.make_noise() noises = [] for noise in noises_single: noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_()) latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(imgs.shape[0], 1) if args.w_plus: latent_in = latent_in.unsqueeze(1).repeat(1, g_ema.n_latent, 1) latent_in.requires_grad = True for noise in noises: noise.requires_grad = True optimizer = optim.Adam([latent_in] + noises, lr=args.lr) pbar = tqdm(range(args.step)) latent_path = [] for i in pbar: t = i / args.step lr = get_lr(t, args.lr) optimizer.param_groups[0]["lr"] = lr noise_strength = latent_std * args.noise * max(0, 1 - t / args.noise_ramp) ** 2 latent_n = latent_noise(latent_in, noise_strength.item()) img_gen, _ = g_ema([latent_n], input_is='latent', noise=noises) batch, channel, height, width = img_gen.shape if height > resize: factor = height // resize img_gen = img_gen.reshape( batch, channel, height // factor, factor, width // factor, factor ) img_gen = img_gen.mean([3, 5]) p_loss = percept(img_gen, imgs).sum() n_loss = noise_regularize(noises) mse_loss = F.mse_loss(img_gen, imgs) loss = p_loss + args.noise_regularize * n_loss + args.mse * mse_loss optimizer.zero_grad() loss.backward() optimizer.step() noise_normalize_(noises) if (i + 1) % 100 == 0: latent_path.append(latent_in.detach().clone()) pbar.set_description(( f"perceptual: {p_loss.item():.8f}; noise regularize: {n_loss.item():.8f}; mse: {mse_loss.item():.8f}; lr: {lr:.4f}" )) img_gen, _ = g_ema([latent_path[-1]], input_is='latent', noise=noises) filename = os.path.splitext(os.path.basename(args.files[0]))[0] + ".pt" img_ar = make_image(img_gen) result_file = {} for i, input_name in enumerate(args.files): noise_single = [] for noise in noises: noise_single.append(noise[i: i + 1]) result_file[input_name] = { "img": img_gen[i], "latent": latent_in[i], "noise": noise_single, } img_name = os.path.splitext(os.path.basename(input_name))[0] + "-project.png" pil_img = Image.fromarray(img_ar[i]) pil_img.save(img_name) torch.save(result_file, filename) print(filename) return img_gen, latent_path, latent_in
def train(): data, info = tfds.load("mnist", with_info=True, data_dir='/data/tensorflow_datasets') train_data = data['train'] if not os.path.exists('./images'): os.makedirs('./images') # settting hyperparameter latent_dim = 100 epochs = 800 batch_size = 200 buffer_size = 6000 save_interval = 50 generator = Generator() discriminator = Discriminator() gen_optimizer = tf.keras.optimizers.Adam(0.0002) disc_optimizer = tf.keras.optimizers.Adam(0.0002) train_dataset = train_data.map(normalize).shuffle(buffer_size).batch( batch_size) cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True) @tf.function def train_step(images): noise = tf.random.normal([batch_size, latent_dim]) with tf.GradientTape(persistent=True) as tape: generated_images = generator(noise) real_output = discriminator(images) generated_output = discriminator(generated_images) gen_loss = generator_loss(cross_entropy, generated_output) disc_loss = discriminator_loss(cross_entropy, real_output, generated_output) grad_disc = tape.gradient(disc_loss, discriminator.trainable_variables) grad_gen = tape.gradient(gen_loss, generator.trainable_variables) disc_optimizer.apply_gradients( zip(grad_disc, discriminator.trainable_variables)) gen_optimizer.apply_gradients( zip(grad_gen, generator.trainable_variables)) return gen_loss, disc_loss seed = tf.random.normal([16, latent_dim]) for epoch in range(epochs): start = time.time() total_gen_loss = 0 total_disc_loss = 0 for images in train_dataset: gen_loss, disc_loss = train_step(images) total_gen_loss += gen_loss total_disc_loss += disc_loss print('Time for epoch {} is {} sec - gen_loss = {}, disc_loss = {}'. format(epoch + 1, time.time() - start, total_gen_loss / batch_size, total_disc_loss / batch_size)) if epoch % save_interval == 0: save_imgs(epoch, generator, seed)
import torch from torch import optim from torch.autograd import Variable import torchvision import os import matplotlib.pyplot as plt from model import Generator, Discriminator from utils import same_seeds from data import get_dataset from torch.utils.data import Dataset, DataLoader import torch.nn as nn import sys if __name__ == "__main__": # hyperparameters batch_size = 64 z_dim = 100 lr = 1e-4 n_epoch = 30 model_name = sys.argv[1] file_name = sys.argv[2] same_seeds(0) # model G = Generator(in_dim=z_dim).cuda() G.load_state_dict(torch.load(model_name)) G.eval() z_sample = Variable(torch.randn(100, z_dim)).cuda() f_imgs_sample = (G(z_sample).data + 1) / 2.0 torchvision.utils.save_image(f_imgs_sample, file_name, nrow=10)
from model import Generator, Discriminator MAX_EPOCHS = 500 LAMBDA = 1 BATCH_SIZE = 4 transform = transforms.Compose([ToTensor()]) train_set = FacadeDataset("datasets/facades/train", transform=transform) val_set = FacadeDataset("datasets/facades/val", transform=transform) test_set = FacadeDataset("datasets/facades/test", transform=transform) train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, \ shuffle=True, num_workers=0) generator = Generator().cuda() discriminator = Discriminator().cuda() gen_optim = torch.optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999)) dis_optim = torch.optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999)) gen_loss = torch.tensor(0) dis_loss = torch.tensor(0) for epoch in range(MAX_EPOCHS): pbar = tqdm(train_loader, dynamic_ncols=True) pbar.set_description(f'e{epoch}') for idx, batch in enumerate(pbar):
# test_path = '../dataset/crop_people/test.txt' show_result = './prediction' if not os.path.exists(show_result): os.mkdir(show_result) # dirname = os.path.dirname(test_path) # with open(test_path, 'r') as fid: # test_list = [l.strip() for l in fid.readlines()] # test_img_files = [os.path.join(dirname, 'image', f) for f in test_list] # test_label_files = [os.path.join(dirname, 'label', f) for f in test_list] # test_img_list = tf.data.Dataset.list_files(test_img_files,shuffle= False) # test_label_list = tf.data.Dataset.list_files(test_label_files,shuffle= False) # test_dataset = tf.data.Dataset.zip((test_img_list,test_label_list)) # test_dataset = test_dataset.map(load_image_test) # test_dataset = test_dataset.batch(1) generator = Generator() discriminator = Discriminator() generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5) checkpoint_best_dir = './Unet_attention_data_checkpoints/best' checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer, discriminator_optimizer=discriminator_optimizer, generator=generator, discriminator=discriminator) checkpoint.restore(tf.train.latest_checkpoint(checkpoint_best_dir)) test_psnr = 0 test_ssim = 0 for n, (input_image, target) in test_dataset.enumerate():
batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=train_sampler) else: dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) dataloader = inf_train_gen(dataloader) #models generator = Generator().to(device) discriminator = Discriminator().to(device) vgg = Vgg16(requires_grad=False).to(device) if args.pre_train: if args.distributed: g_checkpoint = torch.load( args.checkpoint_path + 'generator_checkpoint_{}.ckpt'.format(args.last_iter), map_location=lambda storage, loc: storage.cuda(args.local_rank )) d_checkpoint = torch.load( args.checkpoint_path + 'discriminator_checkpoint_{}.ckpt'.format(args.last_iter), map_location=lambda storage, loc: storage.cuda(args.local_rank ))
def main(): # Load the data data = GANstronomyDataset(opts.DATA_PATH, split=opts.TVT_SPLIT) data.set_split_index(0) data_loader = torch.utils.data.DataLoader(data, batch_size=opts.BATCH_SIZE, shuffle=True) num_classes = data.num_classes() # Make the output directory util.create_dir(opts.RUN_PATH) util.create_dir(opts.IMG_OUT_PATH) util.create_dir(opts.MODEL_OUT_PATH) # Copy opts.py and model.py to opts.RUN_PATH as a record shutil.copy2('opts.py', opts.RUN_PATH) shutil.copy2('model.py', opts.RUN_PATH) shutil.copy2('train.py', opts.RUN_PATH) # Instantiate the models G = Generator(opts.EMBED_SIZE, num_classes).to(opts.DEVICE) G_optimizer = torch.optim.Adam(G.parameters(), lr=opts.ADAM_LR, betas=opts.ADAM_B) D = Discriminator(num_classes).to(opts.DEVICE) D_optimizer = torch.optim.Adam(D.parameters(), lr=opts.ADAM_LR, betas=opts.ADAM_B) if opts.MODEL_PATH is None: start_iepoch, start_ibatch = 0, 0 else: print('Attempting to resume training using model in %s...' % opts.MODEL_PATH) start_iepoch, start_ibatch = load_state_dicts(opts.MODEL_PATH, G, G_optimizer, D, D_optimizer) for iepoch in range(opts.NUM_EPOCHS): for ibatch, data_batch in enumerate(data_loader): # To try to resume training, just continue if iepoch and ibatch are less than their starts if iepoch < start_iepoch or (iepoch == start_iepoch and ibatch < start_ibatch): if iepoch % opts.INTV_PRINT_LOSS == 0 and not ibatch: print('Skipping epoch %d...' % iepoch) continue recipe_ids, recipe_embs, img_ids, imgs, classes, noisy_real, noisy_fake = data_batch # Make sure we're not training on validation or test data! if opts.SAFETY_MODE: for recipe_id in recipe_ids: assert data.get_recipe_split_index(recipe_id) == 0 batch_size, recipe_embs, imgs, classes, classes_one_hot = util.get_variables( recipe_ids, recipe_embs, img_ids, imgs, classes, num_classes) noisy_real, noisy_fake = util.get_variables2( noisy_real, noisy_fake) # Adversarial ground truths all_real = Variable(FloatTensor(batch_size, 1).fill_(1.0), requires_grad=False).to(opts.DEVICE) all_fake = Variable(FloatTensor(batch_size, 1).fill_(0.0), requires_grad=False).to(opts.DEVICE) # Train Discriminator imgs_gen = G(recipe_embs, classes_one_hot).detach() for _ in range(opts.NUM_UPDATE_D): D_optimizer.zero_grad() fake_probs = D(imgs_gen, classes_one_hot) real_probs = D(imgs, classes_one_hot) D_loss = ( BCELoss(fake_probs, noisy_fake if opts.NOISY_LABELS else all_fake) + BCELoss(real_probs, noisy_real if opts.NOISY_LABELS else all_real)) D_loss.backward() D_optimizer.step() # Train Generator for _ in range(opts.NUM_UPDATE_G): G_optimizer.zero_grad() imgs_gen = G(recipe_embs, classes_one_hot) fake_probs = D(imgs_gen, classes_one_hot) G_BCE_loss = BCELoss(fake_probs, all_real) G_MSE_loss = MSELoss(imgs_gen, imgs) G_loss = opts.A_BCE * G_BCE_loss + opts.A_MSE * G_MSE_loss G_loss.backward() G_optimizer.step() if iepoch % opts.INTV_PRINT_LOSS == 0 and not ibatch: print_loss(G_BCE_loss, G_MSE_loss, D_loss, iepoch) if iepoch % opts.INTV_SAVE_IMG == 0 and not ibatch: # Save a training image get_img_gen(data, 0, G, iepoch, opts.IMG_OUT_PATH) # Save a validation image get_img_gen(data, 1, G, iepoch, opts.IMG_OUT_PATH) if iepoch % opts.INTV_SAVE_MODEL == 0 and not ibatch: print('Saving model...') save_model(G, G_optimizer, D, D_optimizer, iepoch, opts.MODEL_OUT_PATH) save_model(G, G_optimizer, D, D_optimizer, 'FINAL', opts.MODEL_OUT_PATH) print('\a') # Ring the bell to alert the human
def train(epochs, batchsize, iterations, data_path, modeldir, nc_size): # Dataset definition dataset = DatasetLoader(data_path, nc_size) # Model & Optimizer definition generator = Generator(nc_size) generator.to_gpu() gen_opt = set_optimizer(generator, 0.00005, 0.9) discriminator = Discriminator(nc_size) discriminator.to_gpu() dis_opt = set_optimizer(discriminator, 0.00005, 0.9) for epoch in range(epochs): sum_dis_loss = 0 sum_gen_loss = 0 for batch in range(0, iterations, batchsize): x, x_label, y, y_label, z, z_label = dataset.train(batchsize) # Discriminator update # Adversairal loss a = y_label - x_label fake = generator(x, a) fake.unchain_backward() loss = adversarial_loss_dis(discriminator, fake, y) # Interpolation loss rnd = np.random.randint(2) if rnd == 0: alpha = xp.random.uniform(0, 0.5, size=batchsize) else: alpha = xp.random.uniform(0.5, 1.0, size=batchsize) alpha = chainer.as_variable(alpha.astype(xp.float32)) alpha = F.tile(F.expand_dims(alpha, axis=1), (1, nc_size)) fake_0 = generator(x, y_label - y_label) fake_1 = generator(x, alpha * a) fake_0.unchain_backward() fake_1.unchain_backward() loss += 10 * interpolation_loss_dis(discriminator, fake_0, fake, fake_1, alpha, rnd) # Matching loss v2 = y_label - z_label v3 = z_label - x_label loss += matching_loss_dis(discriminator, x, fake, y, z, a, v2, v3) discriminator.cleargrads() loss.backward() dis_opt.update() loss.unchain_backward() sum_dis_loss += loss.data # Generator update # Adversarial loss fake = generator(x, a) loss = adversarial_loss_gen(discriminator, fake) # Interpolation loss rnd = np.random.randint(2) if rnd == 0: alpha = xp.random.uniform(0, 0.5, size=batchsize) else: alpha = xp.random.uniform(0.5, 1.0, size=batchsize) alpha = chainer.as_variable(alpha.astype(xp.float32)) alpha = F.tile(F.expand_dims(alpha, axis=1), (1, nc_size)) fake_alpha = generator(x, alpha * a) loss += 10 * interpolation_loss_gen(discriminator, fake_alpha) # Matching loss loss += matching_loss_gen(discriminator, x, fake, a) # Cycle-consistency loss cyc = generator(fake, -a) loss += 10 * F.mean_absolute_error(cyc, x) # Self-reconstruction loss fake_0 = generator(x, y_label - y_label) loss += 10 * F.mean_absolute_error(fake_0, x) generator.cleargrads() loss.backward() gen_opt.update() loss.unchain_backward() sum_gen_loss += loss.data if batch == 0: serializers.save_npz(f"{modeldir}/generator_{epoch}.model", generator) print( f"epoch: {epoch} disloss: {sum_dis_loss/iterations} genloss: {sum_gen_loss/iterations}" )
dataset = Cifar100() stm = PreResNet(depth=32, num_classes=100) ltm = PreResNet(depth=32, num_classes=100) seg = [] for module in ltm.modules(): if isinstance(module, nn.Conv2d) and module.weight.data.shape[ 1] % args.units_x == 0 and module.weight.data.shape[2] == 3: seg.append(module.weight.data.shape[0] // args.units_x) print("seg: ", seg) dim_output = 64 # output dim of PreResNet gen = Generator(in_features=dim_output + args.hidden_dim, out_features=sum(seg) * args.num_units, total_class=args.total_class, hidden_dim=args.hidden_dim, num_units=args.num_units, units_x=args.units_x, units_y=args.units_y, seg=seg) units = torch.rand(args.num_units, args.units_x, args.units_y).unsqueeze(0) units = units.repeat(args.batch_size, 1, 1, 1) # create trainer trainer = ModelTrainer(stm=stm, ltm=ltm, gen=gen, dataset=dataset, units=units) trainer.train(batch_size=args.batch_size, epoches=args.epoches, lr=args.lr)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--input_dir', help = 'Directory containing xxx_i_s and xxx_i_t with same prefix', default = cfg.example_data_dir) parser.add_argument('--save_dir', help = 'Directory to save result', default = cfg.predict_result_dir) parser.add_argument('--checkpoint', help = 'ckpt', default = cfg.ckpt_path) args = parser.parse_args() assert args.input_dir is not None assert args.save_dir is not None assert args.checkpoint is not None print_log('model compiling start.', content_color = PrintColor['yellow']) G = Generator(in_channels = 3).to(device) D1 = Discriminator(in_channels = 6).to(device) D2 = Discriminator(in_channels = 6).to(device) vgg_features = Vgg19().to(device) G_solver = torch.optim.Adam(G.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2)) D1_solver = torch.optim.Adam(D1.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2)) D2_solver = torch.optim.Adam(D2.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2)) checkpoint = torch.load(args.checkpoint) G.load_state_dict(checkpoint['generator']) D1.load_state_dict(checkpoint['discriminator1']) D2.load_state_dict(checkpoint['discriminator2']) G_solver.load_state_dict(checkpoint['g_optimizer']) D1_solver.load_state_dict(checkpoint['d1_optimizer']) D2_solver.load_state_dict(checkpoint['d2_optimizer']) trfms = To_tensor() example_data = example_dataset(data_dir= args.input_dir, transform = trfms) example_loader = DataLoader(dataset = example_data, batch_size = 1, shuffle = False) example_iter = iter(example_loader) print_log('Model compiled.', content_color = PrintColor['yellow']) print_log('Predicting', content_color = PrintColor['yellow']) with torch.no_grad(): for step in tqdm(range(len(example_data)//2)): try: inp = example_iter.next() except StopIteration: example_iter = iter(example_loader) inp = example_iter.next() i_t = inp[0].to(device) i_s = inp[1].to(device) name = str(inp[2][0]) o_sk, o_t, o_b, o_f = G(i_t, i_s) o_sk = o_sk.squeeze(0).detach().to('cpu') o_t = o_t.squeeze(0).detach().to('cpu') o_b = o_b.squeeze(0).detach().to('cpu') o_f = o_f.squeeze(0).detach().to('cpu') if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) o_sk = F.to_pil_image((o_sk)*255.0) o_t = F.to_pil_image((o_t + 1)*127.5) o_b = F.to_pil_image((o_b + 1)*127.5) o_f = F.to_pil_image((o_f + 1)*127.5) o_f.save(os.path.join(args.save_dir, name + 'o_f.png'))
num_gpus = len(gpus) parser = argparse.ArgumentParser(description='Progressive Growing of GANs') parser.add_argument('path', type=str, help='path of specified dataset', default='/raid/datasets/img_align_celeba') parser.add_argument('-d', '--data', default='celeba', type=str, choices=['celeba', 'lsun'], help=('Specify dataset. ' 'Currently CelebA and LSUN is supported')) generator = Generator(code_size, n_label).cuda() discriminator = Discriminator(n_label).cuda() g_running = Generator(code_size, n_label).cuda() g_running.train(False) generator = torch.nn.DataParallel(generator, device_ids=gpus) discriminator = torch.nn.DataParallel(discriminator, device_ids=gpus) g_running = torch.nn.DataParallel(g_running, device_ids=gpus) class_loss = nn.CrossEntropyLoss() g_optimizer = optim.Adam(generator.parameters(), lr=0.001, betas=(0.0, 0.99)) d_optimizer = optim.Adam(discriminator.parameters(), lr=0.001, betas=(0.0, 0.99))
type=str, default="factor", help="filename prefix to result samples", ) parser.add_argument( "factor", type=str, help="name of the closed form factorization result factor file", ) args = parser.parse_args() eigvec = torch.load(args.factor)["eigvec"].to(args.device) ckpt = torch.load(args.ckpt) g = Generator(args.size, 512, 8, channel_multiplier=args.channel_multiplier).to(args.device) g.load_state_dict(ckpt["g_ema"], strict=False) trunc = g.mean_latent(4096) for i in range(1, 512): for j in range(1, 50): latent = torch.randn(args.n_sample, 512, device=args.device) latent = g.get_latent(latent) direction = j * eigvec[:, i].unsqueeze(0) img, _ = g( [latent], truncation=args.truncation, truncation_latent=trunc,
MODEL_NAME = "./pretrained/checkpoints/checkpoint_13399.ckpt" GENERATE_N_SAMPLES = 100000 OUTPUT_DIR = "./samples/generate.txt" BATCH_SIZE = 64 SEQ_LEN = 10 DIM = 128 with open("pretrained/char2int.pickle", "rb") as f: charmap = pickle.load(f) f.close() with open("pretrained/chars.pickle", "rb") as f: chars = pickle.load(f) f.close() fake_inputs = Generator(BATCH_SIZE, SEQ_LEN, DIM, len(charmap)) with tf.Session() as session: def generate_samples(): samples = session.run(fake_inputs) samples = np.argmax(samples, axis=2) decoded_samples = [] for i in range(len(samples)): decoded = [] for j in range(len(samples[i])): decoded.append(chars[samples[i][j]]) decoded_samples.append(tuple(decoded)) return decoded_samples def save(samples):
parser.add_argument('--upscale_factor', default=1, type=int, choices=[1, 4, 8], help='super resolution upscale factor') parser.add_argument('--num_epochs', default=1, type=int, help='train epoch number') opt = parser.parse_args() CROP_SIZE = opt.crop_size #裁剪会带来拼尽问题嘛 UPSCALE_FACTOR = opt.upscale_factor #上采样 NUM_EPOCHS = opt.num_epochs #轮数 val_set = TestDatasetFromFolder('/data/lpw/FusionDataset/tmp_val/', upscale_factor=UPSCALE_FACTOR) #测试集导入 for pthdir in os.listdir('/data/lpw/ResnetFusion/epochs/'): for MODEL_NAME in os.listdir('/data/lpw/ResnetFusion/epochs/'+pthdir): # MODEL_NAME = 'netG_epoch_1_4000.pth' netG = Generator(UPSCALE_FACTOR).eval() netG.cuda() netG.load_state_dict(torch.load('/data/lpw/ResnetFusion/epochs/' + pthdir+'/'+MODEL_NAME)) val_loader = DataLoader(dataset=val_set, num_workers=1, batch_size=1, shuffle=False) epoch =1 out_path = '/data/lpw/ResnetFusion/quota_results/SRF_' + str(UPSCALE_FACTOR) + '/'+pthdir+'/'+MODEL_NAME+'/'#输出路径 print(out_path) if not os.path.exists(out_path): os.makedirs(out_path) val_bar = tqdm(val_loader) #验证集的进度条 val_images = [] index = 1 for val_lr , val_lr_restore, val_hr in val_bar: batch_size = val_lr.size(0)
def train(args): device = torch.device('cuda' if torch.cuda.is_available() and args.enable_cuda else 'cpu') # transforms applied transform = transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) if args.dataset == 'celeba': img_dir = 'cropped' ann_dir = 'list_attr_celeba.csv' train_dataset = CelebA(args.root_dir, img_dir, ann_dir, transform=transform, train=True) test_dataset = CelebA(args.root_dir, img_dir, ann_dir, transform=transform, train=False) attnames = list(train_dataset.df.columns)[1:] elif args.dataset == 'sunattributes': img_dir = 'cropped' ann_dir = 'SUNAttributeDB' train_dataset = SUN_Attributes(args.root_dir, img_dir, ann_dir, transform=transform, train=True) test_dataset = SUN_Attributes(args.root_dir, img_dir, ann_dir, transform=transform, train=False) attnames = train_dataset.attrnames else: raise NotImplementedError() fsize = train_dataset.feature_size trainloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn, drop_last=True) testloader = DataLoader(test_dataset, batch_size=args.show_size, shuffle=True, collate_fn=collate_fn, drop_last=True) ''' dataloader returns dictionaries. sample : {'image' : (bs, 64, 64, 3), 'attributes' : (bs, att_size)} ''' # model, optimizer, criterion if args.residual: gen = Generator_Res(in_c = args.nz + fsize) dis = Discriminator_Res(ny=fsize) else: gen = Generator(in_c = args.nz + fsize) dis = Discriminator(ny=fsize) MODELPATH = '../model/celeba/res_False/gen_epoch_{}.ckpt'.format(args.model_ep) gen.load_state_dict(torch.load(MODELPATH)) enc_y = Encoder(fsize).to(device) MODELPATH = '../model/celeba/res_False/enc_y_epoch_{}.ckpt'.format(args.enc_ep) enc_y.load_state_dict(torch.load(MODELPATH)) enc_z = Encoder(args.nz, for_y=False).to(device) MODELPATH = '../model/celeba/res_False/enc_y_epoch_{}.ckpt'.format(args.enc_ep) enc_z.load_state_dict(torch.load(MODELPATH)) gen.eval() enc_y.eval() enc_z.eval() model = AttrEncoder(outdims=fsize).to(device) # initialize weights for encoders att_optim = optim.Adam(model.parameters(), lr=args.learning_rate, betas=args.betas) criterion = nn.BCELoss().to(device) noise = torch.randn((args.batch_size, args.nz)).to(device) if args.use_tensorboard: writer.add_text("Text", "begin training, lr={}".format(args.learning_rate)) print("begin training, lr={}".format(args.learning_rate), flush=True) stepcnt = 0 for ep in range(args.num_epoch): for it, sample in enumerate(trainloader): elapsed = time.time() image = sample['image'].to(device) att = sample['attributes'].to(device) out = model(image) loss = criterion(out, att) z = enc_z(image) y = enc_y(image) out2 = gen(z, y) loss2 = criterion(out2, att) loss.backward() loss2.backward() att_optim.step() if it % args.log_every == (args.log_every - 1): if args.use_tensorboard: writer.add_scalar('loss', loss, it+1) print("{}th iter \t loss: {:.8f} \t time per iter: {:.05f}s".format(it+1, loss.detach().cpu(), (time.time() - elapsed) / args.log_every), flush=True) cnt, allcnt = eval(model, testloader, device) print("-" * 50) print("epoch {} done. accuracy: {:.03f}%. num guessed: [{:05d}/{:05d}]".format(ep+1, cnt / allcnt * 100, cnt, allcnt)) print("-" * 50, flush=True) if ep % args.save_every == (args.save_every - 1): torch.save(model.state_dict(), "../model/atteval_epoch_{}.model".format(ep+1)) cnt, allcnt = eval(model, testloader, device, early=False) print("-" * 50) print("training done. final accuracy: {:.03f}% num guessed: [{:07d}/{:07d}]".format(ep+1, cnt / allcnt * 100, cnt, allcnt)) print("-" * 50, flush=True)
def main(_): strategy = tf.distribute.MirroredStrategy() NUM_GPU = len(tf.config.experimental.list_physical_devices('GPU')) train_ds, ds_info = tfds.load(FLAGS.dataset, split='train', shuffle_files=True, with_info=True) #dataset is very big, don't want to wait long if FLAGS.dataset == 'lsun/bedroom': train_ds = train_ds.take(300000) output_channels = 3 if FLAGS.dataset == 'cifar10': output_channels = 3 if FLAGS.dataset == 'mnist': output_channels = 1 OUTPUT_DIM = FLAGS.image_size * FLAGS.image_size * output_channels def preprocess(image): """Normalize the images to [-1.0, 1.0]""" image = image['image'] image = tf.image.resize_with_pad(image, FLAGS.image_size, FLAGS.image_size) return (tf.cast(image, tf.float32) - 127.5) / 127.5 train_ds = train_ds.map(preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE) train_ds = train_ds.cache() train_ds = train_ds.shuffle(ds_info.splits['train'].num_examples) train_ds = train_ds.batch(FLAGS.batch_size) train_ds = train_ds.prefetch(tf.data.experimental.AUTOTUNE) train_ds = strategy.experimental_distribute_dataset(train_ds) with strategy.scope(): generator_optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.lr, beta_1=0.5, beta_2=0.9) discriminator_optimizer = tf.keras.optimizers.Adam( learning_rate=FLAGS.lr, beta_1=0.5, beta_2=0.9) inputs = tf.keras.Input(shape=(FLAGS.latent_vector, ), name="latent_vector") outputs = Generator(FLAGS.num_filters, FLAGS.latent_vector, output_channels, OUTPUT_DIM)(inputs) generator = tf.keras.Model(inputs=inputs, outputs=outputs) inputs = tf.keras.Input(shape=(FLAGS.image_size * FLAGS.image_size * output_channels), name="imgs") outputs = Discriminator(FLAGS.num_filters, FLAGS.image_size, output_channels)(inputs) discriminator = tf.keras.Model(inputs=inputs, outputs=outputs) @tf.function def train_gen(): noise = tf.random.normal( [FLAGS.batch_size // NUM_GPU, FLAGS.latent_vector]) with tf.GradientTape() as gen_tape: generated_images = generator(noise, training=True) fake_output = discriminator(tf.reshape(generated_images, [-1, OUTPUT_DIM]), training=False) gen_loss = -tf.reduce_mean(fake_output) gradients_of_generator = gen_tape.gradient( gen_loss, generator.trainable_variables) generator_optimizer.apply_gradients( zip(gradients_of_generator, generator.trainable_variables)) return tf.reduce_mean(gen_loss) @tf.function def train_disc(images): image = tf.reshape(images, [-1, OUTPUT_DIM]) noise = tf.random.normal([images.shape[0], FLAGS.latent_vector]) fake_images = generator(noise, training=True) with tf.GradientTape() as disc_tape: disc_real = discriminator(images, training=True) disc_fake = discriminator(fake_images, training=True) disc_loss = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real) alpha = tf.random.uniform(shape=[ images.shape[0], 1, ], minval=0., maxval=1.) differences = fake_images - image interpolates = image + (alpha * differences) gradients = tf.gradients(discriminator(interpolates), [interpolates])[0] slopes = tf.math.sqrt(tf.reduce_sum(tf.square(gradients), axis=[1])) gradient_penalty = tf.reduce_mean((slopes - 1.)**2) disc_loss += 10 * gradient_penalty gradients_of_discriminator = disc_tape.gradient( disc_loss, discriminator.trainable_variables) discriminator_optimizer.apply_gradients( zip(gradients_of_discriminator, discriminator.trainable_variables)) return tf.reduce_mean(disc_loss) @tf.function def distributed_disc_step(dist_inputs): per_replica_disc_loss = strategy.run(train_disc, args=[dist_inputs]) return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_disc_loss, axis=None) @tf.function def distributed_gen_step(): per_replica_gen_loss = strategy.run(train_gen, args=()) return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_gen_loss, axis=None) def save_images(model, ep, vector): predictions = tf.clip_by_value(model(vector, training=False), -1, 1) plt.figure(figsize=(5, 5)) for i in range(predictions.shape[0]): plt.subplot(4, 4, i + 1) pred = tf.reshape( predictions[i], [FLAGS.image_size, FLAGS.image_size, output_channels]) plt.imshow((pred.numpy() * 127.5 + 127.5).astype(np.uint8)) plt.axis('off') plt.savefig(FLAGS.save_folder + '/image_at_epoch_{:02d}.png'.format(ep)) if not os.path.exists(FLAGS.save_folder): os.makedirs(FLAGS.save_folder) noise_vector = tf.random.normal([FLAGS.num_examples, FLAGS.latent_vector]) for epoch in tqdm(range(FLAGS.epochs)): iterator = iter(train_ds) gen_loss = 0 disc_loss = 0 num_batch = 0 iterations = 0 flag = True while flag: gen_loss += distributed_gen_step() iterations += 1 for _ in range(FLAGS.disc_iters): optional = iterator.get_next_as_optional() if optional.has_value().numpy() == False: flag = False else: data = optional.get_value() d_loss = distributed_disc_step(data) disc_loss += d_loss num_batch += 1 disc_loss /= num_batch gen_loss /= iterations print("Epoch {}, gen_loss {:.5f} \n disc_loss {:.5f}\n".format( epoch, gen_loss, disc_loss)) save_images(generator, epoch, noise_vector) save_images(generator, FLAGS.epochs, noise_vector)
def convert(config): os.makedirs(join(config.convert_dir, config.resume_model), exist_ok=True) sampling_rate, num_mcep, frame_period = config.sampling_rate, 36, 5 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Restore model print(f'Loading the trained models from step {config.resume_model}...') generator = Generator().to(device) g_path = join(config.model_save_dir, f'{config.resume_model}-G.ckpt') generator.load_state_dict( torch.load(g_path, map_location=lambda storage, loc: storage)) # for all possible speaker pairs in config.speakers for i in range(0, len(config.speakers)): for j in range(0, len(config.speakers)): if i != j: target_dir = join( config.convert_dir, str(config.resume_model), f'{config.speakers[i]}_to_{config.speakers[j]}') os.makedirs(target_dir, exist_ok=True) # Load speakers data_loader = ConvertDataset(config, src_spk=config.speakers[i], trg_spk=config.speakers[j]) print('---------------------------------------') print('Source: ', config.speakers[i], ' Target: ', config.speakers[j]) print('---------------------------------------') # Read a batch of testdata src_test_wavfiles = data_loader.get_batch_test_data( batch_size=config.num_converted_wavs) src_test_wavs = [ load_wav(wavfile, sampling_rate) for wavfile in src_test_wavfiles ] with torch.no_grad(): for idx, wav in enumerate(src_test_wavs): print(f'({idx}), file length: {len(wav)}') wav_name = basename(src_test_wavfiles[idx]) # convert wav to mceps f0, _, sp, ap = world_decompose( wav=wav, fs=sampling_rate, frame_period=frame_period) f0_converted = pitch_conversion( f0=f0, mean_log_src=data_loader.logf0s_mean_src, std_log_src=data_loader.logf0s_std_src, mean_log_target=data_loader.logf0s_mean_trg, std_log_target=data_loader.logf0s_std_trg) coded_sp = world_encode_spectral_envelop( sp=sp, fs=sampling_rate, dim=num_mcep) print("Before being fed into G: ", coded_sp.shape) coded_sp_norm = (coded_sp - data_loader.mcep_mean_src ) / data_loader.mcep_std_src coded_sp_norm_tensor = torch.FloatTensor( coded_sp_norm.T).unsqueeze_(0).unsqueeze_(1).to( device) spk_conds = torch.FloatTensor( data_loader.spk_c_trg).to(device) # generate converted speech coded_sp_converted_norm = generator( coded_sp_norm_tensor, spk_conds).data.cpu().numpy() coded_sp_converted = np.squeeze( coded_sp_converted_norm ).T * data_loader.mcep_std_trg + data_loader.mcep_mean_trg coded_sp_converted = np.ascontiguousarray( coded_sp_converted) print("After being fed into G: ", coded_sp_converted.shape) # convert back to wav wav_transformed = world_speech_synthesis( f0=f0_converted, coded_sp=coded_sp_converted, ap=ap, fs=sampling_rate, frame_period=frame_period) wav_id = wav_name.split('.')[0] # SAVE TARGET SYNTHESIZED librosa.output.write_wav( join(target_dir, f'{wav_id}-vcto-{data_loader.trg_spk}.wav'), wav_transformed, sampling_rate) # SAVE COPY OF TARGET REFERENCE wav_num = wav_name.split('.')[0].split('_')[1] copy( f'{config.wav_dir}/{config.speakers[j]}/{config.speakers[j]}_{wav_num}.wav', target_dir)
args = parser.parse_args() sys.path.append(args.repo) import dnnlib from dnnlib import tflib tflib.init_tf() with open(args.path, 'rb') as f: generator, discriminator, g_ema = pickle.load(f) size = g_ema.output_shape[2] g = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier) state_dict = g.state_dict() state_dict = fill_statedict(state_dict, g_ema.vars, size) g.load_state_dict(state_dict) latent_avg = torch.from_numpy(g_ema.vars['dlatent_avg'].value().eval()) ckpt = {'g_ema': state_dict, 'latent_avg': latent_avg} if args.gen: g_train = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier) g_train_state = g_train.state_dict()
def main(): parser = argparse.ArgumentParser() parser.add_argument('--batchsize', '-b', type=int, default=12) parser.add_argument('--epoch', '-e', type=int, default=500) parser.add_argument('--gpu', '-g', type=int, default=-1) parser.add_argument('--root', '-R', default='/mnt/sakuradata10-striped/gao/background') parser.add_argument('--out', '-o', default='/mnt/sakuradata10-striped/gao/results/cartoongan') parser.add_argument('--resume', '-r', default='', help='snapshot No.') parser.add_argument('--model_num', '-m', default='', help='generater No.') parser.add_argument('--snapshot_interval', type=int, default=1000) parser.add_argument('--test_interval', type=int, default=100) parser.add_argument('--display_interval', type=int, default=5) parser.add_argument('--size', type=int, default=256) parser.add_argument('--use_gan', '-G', action='store_true') args = parser.parse_args() print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}\n'.format(args.batchsize)) # Setup models gen = Generator() vgg = VGG() if args.use_gan: dis = Discriminator() else: dis = None # Setup datasets photos = PhotoDataset(os.path.join(args.root, 'photos_resized', '*'), crop_size=args.size) photos_iter = chainer.iterators.SerialIterator(photos, args.batchsize) if args.use_gan: illusts = ImageDataset(os.path.join(args.root, 'illusts', '*', '*', '*'), crop_size=args.size) illusts_iter = chainer.iterators.SerialIterator(illusts, args.batchsize) else: illusts_iter = None # models to gpu if args.gpu >= 0: chainer.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() vgg.to_gpu() if args.use_gan: dis.to_gpu() # Setup optimizer parameters. opt = chainer.optimizers.Adam(alpha=0.0002) opt.setup(gen) if args.use_gan: opt_d = chainer.optimizers.Adam(alpha=0.0002) opt_d.setup(dis) else: opt_d = None # Set up a trainer optimizers = {'gen': opt, 'dis': opt_d} if args.use_gan else {'gen': opt} iterators = {'main': photos_iter, 'illusts': illusts_iter} if args.use_gan else {'main': photos_iter} updater = CartoonGAN( models=(gen, dis, vgg), iterator=iterators, optimizer=optimizers, device=args.gpu, w=10 ) out = os.path.join(args.out, 'gan') if args.use_gan else os.path.join(args.out, 'initial') trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=out) # Load npz if necessary if args.resume: chainer.serializers.load_npz(os.path.join(out, 'snapshot_iter_'+args.resume+'.npz'), trainer) print('snapshot {} loaded\n'.format(args.resume)) elif args.model_num: chainer.serializers.load_npz(os.path.join(args.out, 'initial', 'gen_iter_'+args.model_num+'.npz'), gen) print('model {} loaded\n'.format(args.model_num)) # trainer extensions snapshot_interval = (args.snapshot_interval, 'iteration') test_interval = (args.test_interval, 'iteration') trainer.extend(extensions.dump_graph('gen/loss')) 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.LogReport(trigger=(args.display_interval, 'iteration'), )) report = ['epoch', 'iteration', 'gen/loss', 'gen/content', 'gen/mae'] if args.use_gan: report += ['gen/adv', 'dis/illust', 'dis/edge', 'dis/photo', 'dis/loss'] trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval) trainer.extend(extensions.PrintReport(report)) trainer.extend(extensions.ProgressBar(update_interval=args.display_interval)) trainer.extend(photos.visualizer(), trigger=test_interval) trainer.run() # Save the trained model chainer.serializers.save_npz(os.path.join(args.out, 'model_final'), gen) chainer.serializers.save_npz(os.path.join(args.out, 'optimizer_final'), opt)
if (i + 1) % 10000 == 0: no = str(i + 1).zfill(7) torch.save(generator.state_dict(), f'checkpoint/generator_{no}.pt') torch.save(discriminator.state_dict(), f'checkpoint/discriminator_{no}.pt') torch.save(g_optimizer.state_dict(), f'checkpoint/gen_optimizer_{no}.pt') torch.save(d_optimizer.state_dict(), f'checkpoint/dis_optimizer_{no}.pt') pbar.set_description( (f'{i + 1}; G: {gen_loss_val:.5f};' f' D: {disc_loss_val:.5f}') ) if __name__ == '__main__': args = parser.parse_args() print(args) n_class = len(glob.glob(os.path.join(args.path, '*/'))) if args.model == 'dcgan': from model import Generator, Discriminator elif args.model == 'resnet': from model_resnet import Generator, Discriminator generator = Generator(args.code, n_class).to(device) discriminator = Discriminator(n_class).to(device) g_optimizer = optim.Adam(generator.parameters(), lr=args.lr_g, betas=(0, 0.9)) d_optimizer = optim.Adam(discriminator.parameters(), lr=args.lr_d, betas=(0, 0.9)) train(args, n_class, generator, discriminator)
def main(): os.environ['CUDA_VISIBLE_DEVICES'] = str(cfg.gpu) train_name = get_train_name() print_log('Initializing SRNET', content_color = PrintColor['yellow']) train_data = datagen_srnet(cfg) train_data = DataLoader(dataset = train_data, batch_size = cfg.batch_size, shuffle = False, collate_fn = custom_collate, pin_memory = True) trfms = To_tensor() example_data = example_dataset(transform = trfms) example_loader = DataLoader(dataset = example_data, batch_size = len(example_data), shuffle = False) print_log('training start.', content_color = PrintColor['yellow']) G = Generator(in_channels = 3).cuda() D1 = Discriminator(in_channels = 6).cuda() D2 = Discriminator(in_channels = 6).cuda() vgg_features = Vgg19().cuda() G_solver = torch.optim.Adam(G.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2)) D1_solver = torch.optim.Adam(D1.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2)) D2_solver = torch.optim.Adam(D2.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2)) g_scheduler = torch.optim.lr_scheduler.MultiStepLR(G_solver, milestones=[30, 200], gamma=0.5) d1_scheduler = torch.optim.lr_scheduler.MultiStepLR(D1_solver, milestones=[30, 200], gamma=0.5) d2_scheduler = torch.optim.lr_scheduler.MultiStepLR(D2_solver, milestones=[30, 200], gamma=0.5) requires_grad(G, False) requires_grad(D1, True) requires_grad(D2, True) disc_loss_val = 0 gen_loss_val = 0 grad_loss_val = 0 trainiter = iter(train_data) example_iter = iter(example_loader) K = torch.nn.ZeroPad2d((0, 1, 1, 0)) for step in tqdm(range(cfg.max_iter)): D1_solver.zero_grad() D2_solver.zero_grad() if ((step+1) % cfg.save_ckpt_interval == 0): torch.save( { 'generator': G.module.state_dict(), 'discriminator1': D1.module.state_dict(), 'discriminator2': D2.module.state_dict(), 'g_optimizer': G_solver.state_dict(), 'd1_optimizer': D1_solver.state_dict(), 'd2_optimizer': D2_solver.state_dict(), }, f'checkpoint/train_step-{step+1}.model', ) i_t, i_s, t_sk, t_t, t_b, t_f, mask_t = trainiter.next() i_t = i_t.cuda() i_s = i_s.cuda() t_sk = t_sk.cuda() t_t = t_t.cuda() t_b = t_b.cuda() t_f = t_f.cuda() mask_t = mask_t.cuda() #inputs = [i_t, i_s] labels = [t_sk, t_t, t_b, t_f] o_sk, o_t, o_b, o_f = G(i_t, i_s) o_sk = K(o_sk) o_t = K(o_t) o_b = K(o_b) o_f = K(o_f) #print(o_sk.shape, o_t.shape, o_b.shape, o_f.shape) #print('------') #print(i_s.shape) i_db_true = torch.cat((t_b, i_s), dim = 1) i_db_pred = torch.cat((o_b, i_s), dim = 1) i_df_true = torch.cat((t_f, i_t), dim = 1) i_df_pred = torch.cat((o_f, i_t), dim = 1) o_db_true = D1(i_db_true) o_db_pred = D1(i_db_pred) o_df_true = D2(i_df_true) o_df_pred = D2(i_df_pred) i_vgg = torch.cat((t_f, o_f), dim = 0) out_vgg = vgg_features(i_vgg) db_loss = build_discriminator_loss(o_db_true, o_db_pred) df_loss = build_discriminator_loss(o_df_true, o_df_pred) db_loss.backward() df_loss.backward() D1_solver.step() D2_solver.step() d1_scheduler.step() d2_scheduler.step() clip_grad(D1) clip_grad(D2) if ((step+1) % 5 == 0): requires_grad(G, True) requires_grad(D1, False) requires_grad(D2, False) G_solver.zero_grad() o_sk, o_t, o_b, o_f = G(i_t, i_s) o_sk = K(o_sk) o_t = K(o_t) o_b = K(o_b) o_f = K(o_f) #print(o_sk.shape, o_t.shape, o_b.shape, o_f.shape) #print('------') #print(i_s.shape) i_db_true = torch.cat((t_b, i_s), dim = 1) i_db_pred = torch.cat((o_b, i_s), dim = 1) i_df_true = torch.cat((t_f, i_t), dim = 1) i_df_pred = torch.cat((o_f, i_t), dim = 1) o_db_pred = D1(i_db_pred) o_df_pred = D2(i_df_pred) i_vgg = torch.cat((t_f, o_f), dim = 0) out_vgg = vgg_features(i_vgg) out_g = [o_sk, o_t, o_b, o_f, mask_t] out_d = [o_db_pred, o_df_pred] g_loss, detail = build_generator_loss(out_g, out_d, out_vgg, labels) g_loss.backward() G_solver.step() g_scheduler.step() requires_grad(G, False) requires_grad(D1, True) requires_grad(D2, True) if ((step+1) % cfg.write_log_interval == 0): print('Iter: {}/{} | Gen: {} | D_bg: {} | D_fus: {}'.format(step+1, cfg.max_iter, g_loss.item(), db_loss.item(), df_loss.item())) if ((step+1) % cfg.gen_example_interval == 0): savedir = os.path.join(cfg.example_result_dir, train_name, 'iter-' + str(step+1).zfill(len(str(cfg.max_iter)))) with torch.no_grad(): inp = example_iter.next() o_sk, o_t, o_b, o_f = G(inp) if not os.path.exists(save_dir): os.makedirs(save_dir) o_sk = skimage.img_as_ubyte(o_sk) o_t = skimage.img_as_ubyte(o_t + 1) o_b = skimage.img_as_ubyte(o_b + 1) o_f = skimage.img_as_ubyte(o_f + 1) io.imsave(os.path.join(save_dir, name + 'o_f.png'), o_f) io.imsave(os.path.join(save_dir, name + 'o_sk.png'), o_sk) io.imsave(os.path.join(save_dir, name + 'o_t.png'), o_t) io.imsave(os.path.join(save_dir, name + 'o_b.png'), o_b)
if cuda: label_real = label_real.cuda() label_fake = label_fake.cuda() label_gen_real = label_gen_real.cuda() label_gen_fake = label_gen_fake.cuda() dis_loss = criterion(real.view(-1, 1), label_real) + criterion( fake.view(-1, 1), label_fake) gen_loss = criterion(fake.view(-1, 1), label_gen_real) + criterion( real.view(-1, 1), label_gen_fake) return dis_loss, gen_loss if __name__ == '__main__': A_tr, A_te, B_tr, B_te = data_read('./data', 'horse', 'zebra', IMAGE_SIZE, BATCH_SIZE, True) generatorA = Generator() generatorB = Generator() discriminatorA = Discriminator() discriminatorB = Discriminator() if CUDA: generatorA = generatorA.cuda() generatorB = generatorB.cuda() discriminatorA = discriminatorA.cuda() discriminatorB = discriminatorB.cuda() gen_params = chain(generatorA.parameters(), generatorB.parameters()) dis_params = chain(discriminatorA.parameters(), discriminatorB.parameters()) optim_gen = optim.Adam(gen_params, lr=LR, betas=(0.5, 0.999), weight_decay=1e-5)
download=True, transform=transform) testset = torchvision.datasets.MNIST(root=download_path, train=False, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH * num_gpus, shuffle=True, num_workers=64) testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH * num_gpus, shuffle=False, num_workers=64) G = Generator() C = Critic() if torch.cuda.is_available(): G = G.cuda() C = C.cuda() if torch.cuda.device_count() > 1: G = torch.nn.DataParallel(G) C = torch.nn.DataParallel(C) if cp > 0: load_G = 'model_parameters/G' + str(cp) + '.pth' load_C = 'model_parameters/C' + str(cp) + '.pth' state_dict_G = torch.load(load_G) state_dict_C = torch.load(load_C) G.load_state_dict(state_dict_G["G_state_dict:"])
args.distributed = n_gpu > 1 if args.distributed: torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend="nccl", init_method="env://") synchronize() args.latent = 512 args.n_mlp = 8 args.start_iter = 0 generator = Generator( args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier).to(device) discriminator = Discriminator( args.size, channel_multiplier=args.channel_multiplier).to(device) g_ema = Generator(args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier).to(device) g_ema.eval() accumulate(g_ema, generator, 0) g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1) d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1) g_optim = optim.Adam(
utils.plot_lines(points=dset.p, title='Weight of each gaussian', path='{}/gaussian_weight.png'.format(exp_dir)) sample_points = dset.sample(100) utils.plot_scatter(points=sample_points, centers=dset.centers, title='Sampled data points', path='{}/samples.png'.format(exp_dir)) prefix = "unrolled_steps-{}-prior_std-{:.2f}".format( config.unrolled_steps, np.std(dset.p)) print("Save file with prefix", prefix) G = Generator(input_size=config.g_inp, hidden_size=config.g_hid, output_size=config.g_out).cuda() G._apply(lambda t: t.detach().checkpoint()) D = Discriminator(input_size=config.d_inp, hidden_size=config.d_hid, output_size=config.d_out).cuda() D._apply(lambda t: t.detach().checkpoint()) criterion = nn.BCELoss( ) # Binary cross entropy: http://pytorch.org/docs/nn.html#bceloss def binary_cross_entropy(x, y): loss = -(x.log() * y + (1 - x).log() * (1 - y)) return loss.mean() criterion = binary_cross_entropy
help='using GPU or CPU') parser.add_argument('--image_name', type=str, help='test low resolution image name') parser.add_argument('--weights', default='weights/netG_epoch_4_54.pth', type=str, help='generator model epoch name') opt = parser.parse_args() UPSCALE_FACTOR = opt.upscale_factor TEST_MODE = True if opt.test_mode == 'GPU' else False IMAGE_NAME = opt.image_name weights_path = opt.weights model = Generator(UPSCALE_FACTOR).eval() if TEST_MODE: model.cuda() model.load_state_dict(torch.load(weights_path)) else: model.load_state_dict( torch.load(weights_path, map_location=lambda storage, loc: storage)) image = Image.open(IMAGE_NAME) image = Variable(ToTensor()(image), volatile=True).unsqueeze(0) if TEST_MODE: image = image.cuda() start = time.clock() out = model(image) elapsed = (time.clock() - start)
####### Define data loader and models ########## train_sampler = RandomSampler( data_source=dataset, replacement=True, num_samples=int(1e100), # make the dataloader "infinite" ) dataloader = DataLoader( dataset, batch_size=cfg.batch_size, sampler=train_sampler, pin_memory=True, #shuffle=True, num_workers=4, ) G = Generator( z_dim=cfg.z_dim, n_feat=cfg.n_feat, ) if cfg.minibatch_discrimination: logging.info("Using minibatch discrimination") D = Discriminator_MD(batch_size=cfg.batch_size, n_feat=cfg.n_feat) else: #logging.info("Using minibatch discr") D = Discriminator(n_feat=cfg.n_feat) G.apply(weights_init) D.apply(weights_init) G = G.to(cfg.device) D = D.to(cfg.device)
transform = transforms.Compose([ transforms.Resize(resize), transforms.CenterCrop(resize), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) imgs = [] for imgfile in args.files: img = transform(Image.open(imgfile).convert("RGB")) imgs.append(img) imgs = torch.stack(imgs, 0).to(device) g_ema = Generator(args.size, 512, 8) g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False) g_ema.eval() g_ema = g_ema.to(device) with torch.no_grad(): noise_sample = torch.randn(n_mean_latent, 512, device=device) latent_out = g_ema.style(noise_sample) latent_mean = latent_out.mean(0) latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent)**0.5 percept = lpips.PerceptualLoss(model="net-lin", net="vgg", use_gpu=device.startswith("cuda"))
# custom weights initialization called on netG and netD def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv') != -1: m.weight.data.normal_(0.0, 0.02) elif classname.find('BatchNorm') != -1: m.weight.data.normal_(1.0, 0.02) m.bias.data.fill_(0) '''---------------------------------------INSTANTIATE MODELS-------------------------------------''' from model import Generator, Discriminator netG = Generator(ngpu).to(device) netG.apply(weights_init) if args.netG != '': netG.load_state_dict(torch.load(args.netG)) # print(netG) netD = Discriminator(ngpu).to(device) netD.apply(weights_init) if args.netD != '': netD.load_state_dict(torch.load(args.netD)) # print(netD) if gpu: netD.to(device) netG.to(device) '''----------------------------------------LOSS & OPTIMIZER--------------------------------------'''