def __init__(self, device): self.device = device self.netG_A = self.__init_weights(Generator(3, use_dropout=False).to(self.device)) self.netG_B = self.__init_weights(Generator(3, use_dropout=False).to(self.device)) self.netD_A = self.__init_weights(Discriminator(3).to(self.device)) self.netD_B = self.__init_weights(Discriminator(3).to(self.device)) self.criterion_gan = nn.MSELoss() self.criterion_cycle = nn.L1Loss() self.criterion_idt = nn.L1Loss() self.optimizer_G = torch.optim.Adam(itertools.chain( self.netG_A.parameters(), self.netG_B.parameters()), lr=0.0002, betas=(0.5, 0.999)) self.optimizer_D = torch.optim.Adam(itertools.chain( self.netD_A.parameters(), self.netD_B.parameters()), lr=0.0002, betas=(0.5, 0.999)) self.optimizers = [self.optimizer_G, self.optimizer_D] self.fake_A_pool = ImagePool(50) self.fake_B_pool = ImagePool(50) self.loss_names = ['D_A', 'G_A', 'cycle_A', 'idt_A', 'D_B', 'G_B', 'cycle_B', 'idt_B'] self.model_names = ['G_A', 'G_B', 'D_A', 'D_B'] self.lambda_A = 10 self.lambda_B = 10 self.lambda_idt = 0.5 self.save_dir = './models'
def main(args): # print args print_args(args) # download and extract dataset get_cifar10_data(args.data_path) data_dirs = extract_cifar10_images(args.data_path) dataset = Cifar10Dataset(root_dir=data_dirs["test"], mirror=False, random_seed=1) print("test dataset len: {}".format(len(dataset))) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) generator_bn = Generator("batch") generator_sn = Generator("batch") generator_bn.load_state_dict( torch.load("checkpoints/checkpoint_ep0_gen.pt", map_location="cpu")) generator_sn.load_state_dict( torch.load("checkpoints_spectral/checkpoint_ep10_gen.pt", map_location="cpu")) if not os.path.exists(args.save_path): os.makedirs(args.save_path) for idx, sample in enumerate(data_loader): img_l, real_img_lab = sample[:, 0:1, :, :], sample fake_img_ab_bn = generator_bn(img_l).detach() fake_img_lab_bn = torch.cat([img_l, fake_img_ab_bn], dim=1) fake_img_ab_sn = generator_sn(img_l).detach() fake_img_lab_sn = torch.cat([img_l, fake_img_ab_sn], dim=1) print("sample {}/{}".format(idx + 1, len(data_loader) + 1)) save_test_sample(real_img_lab, fake_img_lab_bn, fake_img_lab_sn, osp.join(args.save_path, "test_sample_{}.png".format(idx)), show=True)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--learning-rate', '-lr', type=float, default=1e-3) parser.add_argument('--epochs', type=int, default=20) parser.add_argument('--no-cuda', action='store_true') parser.add_argument('--data-parallel', action='store_true') parser.add_argument('--num-d-iterations', type=int, default=1) args = parser.parse_args() args.cuda = torch.cuda.is_available() and not args.no_cuda print(args) device = torch.device('cuda' if args.cuda else 'cpu') net_g = Generator(ch=128).to(device) net_d = Discriminator(ch=128).to(device) optim_g = optim.Adam( net_g.parameters(), lr=args.learning_rate, betas=(0.5, 0.999)) optim_d = optim.Adam( net_d.parameters(), lr=args.learning_rate, betas=(0.5, 0.999)) dataloader = get_cat_dataloader() trainer = Trainer(net_g, net_d, optim_g, optim_d, dataloader, device, args.num_d_iterations) os.makedirs('samples', exist_ok=True) trainer.train(args.epochs)
def init_trainer(self): # networks self.G = Generator(nc=self.nc, nz=self.nz, size=self.size) self.D = Discriminator(nc=self.nc, nz=self.nz, size=self.size) self.G_EMA = copy.deepcopy(self.G) # move to GPU self.G = nn.DataParallel(self.G, device_ids=self.device_ids).to(self.device) self.D = nn.DataParallel(self.D, device_ids=self.device_ids).to(self.device) self.G_EMA = self.G_EMA.to('cpu') # keep this model on CPU to save GPU memory for param in self.G_EMA.parameters(): param.requires_grad_(False) # turn off grad because G_EMA will only be used for inference # optimizers self.opt_G = optim.Adam(self.G.parameters(), lr=self.lr, betas=(0,0.99), eps=1e-8, weight_decay=0.) self.opt_D = optim.Adam(self.D.parameters(), lr=self.lr, betas=(0,0.99), eps=1e-8, weight_decay=0.) # data loader self.transform = transforms.Compose([ RatioCenterCrop(1.), transforms.Resize((300,300), Image.ANTIALIAS), transforms.RandomCrop((self.size,self.size)), RandomRotate(), transforms.RandomVerticalFlip(), transforms.RandomHorizontalFlip(), transforms.ToTensor() ]) self.dataset = ISIC_GAN('train_gan.csv', transform=self.transform) self.dataloader = torch.utils.data.DataLoader(self.dataset, batch_size=self.batch_size, shuffle=True, num_workers=8, worker_init_fn=_worker_init_fn_(), drop_last=True) # tickers (used for fading in) self.tickers = self.unit_epoch * self.num_aug * len(self.dataloader)
def train(): RDN = Generator("RDN") D = Discriminator("discriminator") HR = tf.placeholder(tf.float32, [None, 96, 96, 3]) LR = tf.placeholder(tf.float32, [None, 24, 24, 3]) SR = RDN(LR) fake_logits = D(SR, LR) real_logits = D(HR, LR) D_loss, G_loss = Hinge_Loss(fake_logits, real_logits) G_loss += MSE(SR, HR) * LAMBDA itr = tf.Variable(MAX_ITERATION, dtype=tf.int32, trainable=False) learning_rate = tf.Variable(2e-4, trainable=False) op_sub = tf.assign_sub(itr, 1) D_opt = tf.train.AdamOptimizer(learning_rate, beta1=0., beta2=0.9).minimize(D_loss, var_list=D.var_list()) with tf.control_dependencies([op_sub]): G_opt = tf.train.AdamOptimizer(learning_rate, beta1=0., beta2=0.9).minimize(G_loss, var_list=RDN.var_list()) sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver() while True: HR_data, LR_data = read_crop_data(TRAINING_SET_PATH, BATCH_SIZE, [96, 96, 3], 4) sess.run(D_opt, feed_dict={HR: HR_data, LR: LR_data}) [_, iteration] = sess.run([G_opt, itr], feed_dict={HR: HR_data, LR: LR_data}) iteration = MAX_ITERATION - iteration if iteration < MAX_ITERATION // 2: learning_rate = learning_rate * (iteration * 2 / MAX_ITERATION) if iteration % 10 == 0: [D_LOSS, G_LOSS, LEARNING_RATE, img] = sess.run([D_loss, G_loss, learning_rate, SR], feed_dict={HR: HR_data, LR: LR_data}) output = (np.concatenate((HR_data[0, :, :, :], img[0, :, :, :]), axis=1) + 1) * 127.5 Image.fromarray(np.uint8(output)).save(RESULTS+str(iteration)+".jpg") print("Iteration: %d, D_loss: %f, G_loss: %f, LearningRate: %f"%(iteration, D_LOSS, G_LOSS, LEARNING_RATE)) if iteration % 500 == 0: saver.save(sess, SAVE_MODEL + "model.ckpt")
def main(args): """ Main function for the script :param args: parsed command line arguments :return: None """ print("Creating generator object ...") # create the generator object gen = th.nn.DataParallel( Generator(depth=args.depth, latent_size=args.latent_size)) print("Loading the generator weights from:", args.generator_file) # load the weights into it gen.load_state_dict(th.load(args.generator_file, map_location=str(device))) # path for saving the files: save_path = args.out_dir print("Generating scale synchronized images ...") for img_num in tqdm(range(1, args.num_samples + 1)): # generate the images: with th.no_grad(): point = th.randn(1, args.latent_size) point = (point / point.norm()) * (args.latent_size**0.5) ss_image = gen(point, depth=args.out_depth, alpha=1) # color adjust the generated image: ss_image = adjust_dynamic_range(ss_image) # save the ss_image in the directory imsave(os.path.join(save_path, str(img_num) + ".png"), ss_image.squeeze(0).permute(1, 2, 0).cpu()) print("Generated %d images at %s" % (args.num_samples, save_path))
def __init__(self, image_size, input_channels, hidden_channels, output_channels, latent_dimension, lr, device, clamp=0.01, gp_weight=10): self.image_size = image_size self.input_channels = input_channels self.hidden_chanels = hidden_channels self.output_channels = output_channels self.latent_dimension = latent_dimension self.device = device self.clamp = clamp self.gp_weight = gp_weight self.critic = Critic(image_size, hidden_channels, input_channels).to(device) self.generator = Generator( image_size, latent_dimension, hidden_channels, output_channels).to(device) self.critic.apply(self.weights_init) self.generator.apply(self.weights_init) self.optimizer_critic = torch.optim.RMSprop( self.critic.parameters(), lr) self.optimizer_gen = torch.optim.RMSprop( self.generator.parameters(), lr) self.optimizer_critic = torch.optim.Adam( self.critic.parameters(), lr, betas=(0, 0.9)) self.optimizer_gen = torch.optim.Adam( self.generator.parameters(), lr, betas=(0, 0.9)) self.critic_losses = [] self.gen_losses = [] self.losses = []
def create_generator(self): kernels_gen_encoder = [ (32, 1, 0), # [batch, 512, 512, ch] => [batch, 512, 512, 32] (32, 2, 0), # [batch, 512, 512, 32] => [batch, 256, 256, 32] (64, 2, 0), # [batch, 256, 256, 32] => [batch, 128, 128, 64] (128, 2, 0), # [batch, 128, 128, 64] => [batch, 64, 64, 128] (256, 2, 0), # [batch, 64, 64, 128] => [batch, 32, 32, 256] (512, 2, 0), # [batch, 32, 32, 256] => [batch, 16, 16, 512] (512, 2, 0), # [batch, 16, 16, 512] => [batch, 8, 8, 512] (512, 2, 0), # [batch, 8, 8, 512] => [batch, 4, 4, 512] (512, 2, 0) # [batch, 4, 4, 512] => [batch, 2, 2, 512] ] kernels_gen_decoder = [ (512, 2, 0), # [batch, 2, 2, 512] => [batch, 4, 4, 512] (512, 2, 0), # [batch, 4, 4, 512] => [batch, 8, 8, 512] (512, 2, 0), # [batch, 8, 8, 512] => [batch, 16, 16, 512] (256, 2, 0), # [batch, 16, 16, 512] => [batch, 32, 32, 256] (128, 2, 0), # [batch, 32, 32, 256] => [batch, 64, 64, 128] (64, 2, 0), # [batch, 64, 64, 128] => [batch, 128, 128, 64] (32, 2, 0), # [batch, 128, 128, 64] => [batch, 256, 256, 64] (32, 2, 0) # [batch, 256, 256, 64] => [batch, 512, 512, 32] ] return Generator('gen', kernels_gen_encoder, kernels_gen_decoder, training=self.options.training)
def train(path): imgs = get_training_imgs(path) nums = len(imgs) Gs = [] Ds = [] fixed_Zs = [] sigmas = [] ch = 16 for i in range(nums): if i % 4 == 0: ch = ch * 2 G = Generator(ch) D = Discriminator(ch) G.to("cuda:0") D.to("cuda:0") if i > 0: try: G.load_state_dict(G_.state_dict()) D.load_state_dict(D_.state_dict()) del G_, D_ except: pass Gs.append(G) Ds.append(D) print(".............Total Scale: %d, current scale: %d............."%(nums, i+1)) G_, D_ = train_single_scale(Gs, Ds, imgs[:i+1], sigmas, fixed_Zs) state_dict = {} state_dict["Gs"] = Gs state_dict["sigmas"] = sigmas state_dict["imgs"] = imgs torch.save(state_dict, path[:-3]+"pth")
def get_models(latent_dim, model_dim, device, output_dim, channels, init=True): generator = Generator(latent_dim, model_dim, channels).to(device) critic = Critic(model_dim, output_dim, channels).to(device) if init: generator.apply(__weights_init_normal) critic.apply(__weights_init_normal) return generator, critic
def __init__(self, args): self.z_dim = args.z_dim self.decay_rate = args.decay_rate self.learning_rate = args.learning_rate self.model_name = args.model_name self.batch_size = args.batch_size #initialize networks self.Generator = Generator(self.z_dim).cuda() self.Encoder = Encoder(self.z_dim).cuda() self.Discriminator = Discriminator().cuda() #set optimizers for all networks self.optimizer_G_E = torch.optim.Adam( list(self.Generator.parameters()) + list(self.Encoder.parameters()), lr=self.learning_rate, betas=(0.5, 0.999)) self.optimizer_D = torch.optim.Adam(self.Discriminator.parameters(), lr=self.learning_rate, betas=(0.5, 0.999)) #initialize network weights self.Generator.apply(weights_init) self.Encoder.apply(weights_init) self.Discriminator.apply(weights_init)
def main(): opt = parse.parse_args() print(opt) img_shape = (opt.channels, opt.img_size, opt.img_size) cuda = True if torch.cuda.is_available() else False generator = Generator(opt.n_classes, opt.latent_dim, img_shape) print(generator) if cuda: generator = generator.cuda() torch.cuda.set_device(opt.gpu_ids) os.makedirs('data/mnist', exist_ok=True) dataloader = DataLoader(datasets.MNIST('data/mnist', train=False, transform=transforms.Compose([ transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize( mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]), download=True), batch_size=opt.batch_size, shuffle=False) FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor generator.eval() generator.load_state_dict(torch.load(opt.model_dir)) for i, (image, labels) in enumerate(dataloader): batch_size = image.shape[0] z = Variable( FloatTensor( np.random.normal(loc=0.0, scale=1.0, size=(opt.batch_size, opt.latent_dim)))) real_image = Variable(image.type(FloatTensor)) labels = Variable(labels.type(LongTensor)) gen_image = generator(z, labels) print("[Batch %d/%d]" % (i, len(dataloader))) gen_image = gen_image.view(gen_image.size(0), *img_shape) real_images = real_image.view(real_image.size(0), *img_shape) save_image( gen_image.data, 'images_4_nocat_MSE_ls_lc/test/fake/%d_generated_image.png' % (i), nrow=1, normalize=True) save_image(real_images.data, 'images_4_nocat_MSE_ls_lc/test/real/%d_real_image.png' % i, nrow=1, normalize=True)
def load_gan(model_path): gan = Generator() model = torch.load(model_path)['gan'] #pretrained_state_dict = torch.load(model_path)['gan'] #model_state_dict = gan.state_dict() #model = dict_unite(pretrained_state_dict, model_state_dict) gan.load_state_dict(model) return gan
def test(config): """ This function tests a trained neuronal network. In this case the generator of the GAN. Args: config : A dictionary with the extracted configuration out of the config.ini file and the commandlineparser. Following dictionary-values are necessary: {'image_dim_x': intended image x-dimension, 'image_dim_y': intended image y-dimension, 'image_dim': image_dim_x * image_dim_y, 'batch_size': size of the batch per iteration 'datatype': datatype of the dataset (MNIST | CIFAR10 | ImageFolder), 'device': device which should execute the test (gpu | cpu), 'z_dim': Hyperparameter for the Generator, 'logpathfake': logpath for tensorboard SummaryWriter (--logdir)} """ #load Generator gen = Generator(config['z_dim'], config['image_dim']).to(config['device']) gen.load_model(config) #load Dataset print("Load dataset...") loader = load_Dataset(config) #initialize tensorboard summarywriter writer_fake = SummaryWriter(config['logpathfake']) writer_real = SummaryWriter(config['logpathreal']) trained_iterations = gen.training_iterations step_gen = gen.training_iterations #Testing trained Generator print("Testing...") for batch_idx, (real, _) in enumerate(loader): real = real.view(-1, config['image_dim']).to(config['device']) batch_size = real.shape[0] if batch_idx == 0: with torch.no_grad(): noise = torch.randn(config['batch_size'], config['z_dim']).to(config['device']) fake = gen(noise).reshape(-1, 1, config['image_dim_x'], config['image_dim_y']) data = real.reshape(-1, 1, config['image_dim_x'], config['image_dim_y']) img_grid_fake = torchvision.utils.make_grid(fake, normalize=True) img_grid_real = torchvision.utils.make_grid(data, normalize=True) writer_fake.add_image( "Mnist generated fake images out of test", img_grid_fake, global_step=trained_iterations) writer_real.add_image("Mnist reference Images", img_grid_real, global_step=0)
def generate_process(opt): print("generate_process") # ----- Device Setting ----- device = torch.device("cpu") # ----- Output Setting ----- img_output_dir = opt.output_dir if not os.path.exists(img_output_dir): os.mkdir(img_output_dir) if not os.path.exists(os.path.join(img_output_dir, "images")): os.mkdir(os.path.join(img_output_dir, "images")) if opt.save_latent is True: if not os.path.exists(os.path.join(img_output_dir, "latents")): os.mkdir(os.path.join(img_output_dir, "latents")) print("Output :", img_output_dir) # ----- Model Loading ----- print("Use model :", opt.model) model_g = Generator() model_g.load_state_dict(torch.load(opt.model, map_location="cpu")) model_g.to(device) model_g.eval() if opt.mode == "normal": latents = [torch.randn(size=(opt.width * opt.height, opt.latent_size, 1, 1)) for i in range(opt.n_img)] # latents = [torch.randn(size=(opt.width * opt.height, opt.latent_size, 1, 1)] elif opt.mode == "use_latent": assert opt.latent_dir != "None", "latent source directory is not set" latent_paths = [os.path.join(opt.latent_dir, name) for name in os.listdir(opt.latent_dir)] latents = [torch.from_numpy(np.load(path)) for path in latent_paths] elif opt.mode == "inter": latent_start = torch.from_numpy(np.load(opt.start_latent)) latent_end = torch.from_numpy(np.load(opt.end_latent)) alphas = [float(n / opt.latent_num) for n in range(opt.n_img)] latents = [alpha * latent_end + (1 - alpha) * latent_start for alpha in alphas] print("Generate image num :", len(latents)) # ----- Generate Step ----- print("Start Generate Process") for index,latent in tqdm(enumerate(latents)): img_path = os.path.join(img_output_dir, "images", str(index + 1) + ".png") generate(latent, model_g, opt.width, opt.height, img_path) if opt.save_latent: latent_path = os.path.join(img_output_dir, "latents", str(index + 1) + ".npy") np.save(latent_path, latent.numpy()) print("Finish Generate Process")
def init_test(args): """Create the data loader and the generators for testing purposes.""" # create loader dataset = Cifar10Dataset.get_datasets_from_scratch(args.data_path)['test'] print('Test dataset len: {}'.format(len(dataset))) data_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) # check CUDA availability and set device device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print('Use GPU: {}'.format(str(device) != 'cpu')) # download the weights for the generators if not os.path.exists('batchnorm_ep200_weigths_gen.pt'): print('Downloading model weights for generator with BN...') os.system( 'wget https://www.dropbox.com/s/r33ndl969q83gik/batchnorm_ep200_weigths_gen.pt' ) if not os.path.exists('spectralnorm_ep100_weights_gen.pt'): print('Downloading model weights for generator with SN...') os.system( 'wget https://www.dropbox.com/s/tccxduyqp3dj5dg/spectralnorm_ep100_weights_gen.pt' ) # load generator that was trained with batch norm generator_bn = Generator(normalization_type='batch').to(device) # load generator that was trained with spectral norm generator_sn = Generator(normalization_type='batch').to(device) # load the weights generator_bn.load_state_dict( torch.load('batchnorm_ep200_weigths_gen.pt', map_location=device)) generator_sn.load_state_dict( torch.load('spectralnorm_ep100_weights_gen.pt', map_location=device)) # make save dir, if needed if not os.path.exists(args.save_path): os.makedirs(args.save_path) return device, data_loader, generator_bn, generator_sn
def make_generated_dataset(): generator = Generator(args.dim, args.latent_dim, args.n_pixels, args.bn_g) generator.load_state_dict(torch.load(args.model_dir, map_location=device)) generator.eval() with torch.no_grad(): samples = generator(args.n_samples) samples = ((samples + 1.0) * (255 / 2.0)).numpy().astype('uint8') for i in range(args.n_samples): img = np.transpose(samples[i], (1, 2, 0)) imageio.imwrite( '{}{}{}'.format(args.gen_save_dir, i + 1, args.file_type), img)
def load_gen(model_load_path, use_bn=False, last_act='tanh', use_wscale=True, use_he_backward=True, resolution_list=[4, 8, 16, 32, 64, 128], channel_list=[512, 512, 256, 128, 64, 32]): """Load generator and grow it """ with nn.parameter_scope("generator"): _ = nn.load_parameters(model_load_path) gen = Generator(use_bn=use_bn, last_act=last_act, use_wscale=use_wscale, use_he_backward=use_he_backward) for i in range(len(resolution_list)): gen.grow(resolution_list[i], channel_list[i]) return gen
def __init__(self, dataset_dir, generator_channels, discriminator_channels, nz, style_depth, lrs, betas, eps, phase_iter, weights_halflife, batch_size, n_cpu, opt_level): self.nz = nz self.dataloader = Dataloader(dataset_dir, batch_size, phase_iter * 2, n_cpu) self.generator = Generator(generator_channels, nz, style_depth).cuda() self.generator_ema = Generator(generator_channels, nz, style_depth).cuda() self.generator_ema.load_state_dict(copy.deepcopy(self.generator.state_dict())) self.discriminator = Discriminator(discriminator_channels).cuda() self.tb = tensorboard.tf_recorder('StyleGAN') self.phase_iter = phase_iter self.lrs = lrs self.betas = betas self.weights_halflife = weights_halflife self.opt_level = opt_level self.ema = None torch.backends.cuda.benchmark = True
def train(): generator = Generator() discriminator = Discriminator() generator.to("cuda:0") discriminator.to("cuda:0") Opt_D = optim.Adam(discriminator.parameters(), lr=2e-4, betas=(0.5, 0.999)) Opt_G = optim.Adam(generator.parameters(), lr=2e-4, betas=(0.5, 0.999)) data = np.concatenate((sio.loadmat("D:/cifar10/data_batch_1.mat")["data"], sio.loadmat("D:/cifar10/data_batch_2.mat")["data"], sio.loadmat("D:/cifar10/data_batch_3.mat")["data"], sio.loadmat("D:/cifar10/data_batch_4.mat")["data"], sio.loadmat("D:/cifar10/data_batch_5.mat")["data"])) nums = data.shape[0] for i in range(100000): rand_idx = np.random.choice(range(nums), batchsize) batch = np.reshape(data[rand_idx], [batchsize, 3, 32, 32]) batch = torch.tensor(batch / 127.5 - 1, dtype=torch.float32).to("cuda:0") for j in range(n_cri): z = torch.randn(batchsize, 128).to("cuda:0") fake_img = generator(z).detach() fake_logits = discriminator(fake_img) real_logits = discriminator(batch) D_loss = torch.mean( torch.max(torch.zeros_like(real_logits), 1. - real_logits)) + torch.mean( torch.max(torch.zeros_like(fake_logits), 1. + fake_logits)) Opt_D.zero_grad() D_loss.backward() Opt_D.step() z = torch.randn(batchsize, 128).to("cuda:0") fake_img = generator(z) fake_logits = discriminator(fake_img) G_loss = -torch.mean(fake_logits) Opt_G.zero_grad() G_loss.backward() Opt_G.step() if i % 100 == 0: img = (fake_img[0] + 1) * 127.5 Image.fromarray( np.uint8( np.transpose(img.cpu().detach().numpy(), axes=[1, 2, 0]))).save("./results/" + str(i) + ".jpg") print("Iteration: %d, D_loss: %f, G_loss: %f" % (i, D_loss, G_loss)) if i % 1000 == 0: torch.save(generator, "generator.pth") torch.save(discriminator, "discriminator.pth")
def __init__(self, hyperparameters): super(STGANtrainer, self).__init__() self.hyperparameters = hyperparameters self.gen = Generator(5, 5, 4, 13) self.dis = Discriminator(5, 64, 13) self.dis_opt = Adam(self.dis.parameters(), lr=self.hyperparameters['lr_dis'], betas=self.hyperparameters['lr_beta']) self.gen_opt = Adam(self.gen.parameters(), lr=self.hyperparameters['lr_gen'], betas=self.hyperparameters['lr_beta']) self.dis_attr_opt = Adam(self.gen.parameters(), lr=0.5 * self.hyperparameters['lr_gen'], betas=self.hyperparameters['lr_beta'])
def init(): for dirname in os.listdir(state_dict_path): discriminator = Discriminator() discriminator.load_state_dict( torch.load(state_dict_path + '/' + dirname + '/discriminator')) generator = Generator() generator.load_state_dict( torch.load(state_dict_path + '/' + dirname + '/generator')) network = { 'id': dirname, 'discriminator': discriminator, 'generator': generator, } networks.append(network) network_dict[dirname] = network print('Model #' + dirname + ' loaded.')
def __init__(self, generator_path, discriminator_path): self.device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') # Path for loading and saving model weights self.generator_path = generator_path self.discriminator_path = discriminator_path # Get training and testing data train_set = PegasusDataset('data', train=True, download=True, transform=torchvision.transforms.Compose( [torchvision.transforms.ToTensor()])) test_set = PegasusDataset('data', train=False, download=True, transform=torchvision.transforms.Compose( [torchvision.transforms.ToTensor()])) self.train_loader = torch.utils.data.DataLoader(train_set, shuffle=True, batch_size=BATCH_SIZE, drop_last=True) self.test_loader = torch.utils.data.DataLoader(test_set, shuffle=True, batch_size=BATCH_SIZE, drop_last=True) # Create generator and discriminator self.G = Generator().to(self.device) self.D = Discriminator().to(self.device) #self.loadModels(self.generator_path, self.discriminator_path) # initialise the optimiser self.optimiser_G = torch.optim.Adam(self.G.parameters(), lr=0.0002, betas=(0.5, 0.99)) self.optimiser_D = torch.optim.Adam(self.D.parameters(), lr=0.0002, betas=(0.5, 0.99)) self.bce_loss = nn.BCELoss() self.dgRatio = 0
def __init__(self, in_size: int, ts_size: int = 100, latent_dim: int = 20, lr: float = 0.0005, weight_decay: float = 1e-6, iterations_critic: int = 5, gamma: float = 10, weighted: bool = True, use_gru=False): super(TadGAN, self).__init__() self.in_size = in_size self.latent_dim = latent_dim self.lr = lr self.weight_decay = weight_decay self.iterations_critic = iterations_critic self.gamma = gamma self.weighted = weighted self.hparams = { 'lr': self.lr, 'weight_decay': self.weight_decay, 'iterations_critic': self.iterations_critic, 'gamma': self.gamma } self.encoder = Encoder(in_size, ts_size=ts_size, out_size=self.latent_dim, batch_first=True, use_gru=use_gru) self.generator = Generator(use_gru=use_gru) self.critic_x = CriticX(in_size=in_size) self.critic_z = CriticZ() self.encoder.apply(init_weights) self.generator.apply(init_weights) self.critic_x.apply(init_weights) self.critic_z.apply(init_weights) if self.logger is not None: self.logger.log_hyperparams(self.hparams) self.y_hat = [] self.index = [] self.critic = []
def create_generator(self): kernels_gen_encoder = [ (64, 1, 0), # [batch, 32, 32, ch] => [batch, 32, 32, 64] (128, 2, 0), # [batch, 32, 32, 64] => [batch, 16, 16, 128] (256, 2, 0), # [batch, 16, 16, 128] => [batch, 8, 8, 256] (512, 2, 0), # [batch, 8, 8, 256] => [batch, 4, 4, 512] (512, 2, 0), # [batch, 4, 4, 512] => [batch, 2, 2, 512] ] kernels_gen_decoder = [ (512, 2, 0.5), # [batch, 2, 2, 512] => [batch, 4, 4, 512] (256, 2, 0.5), # [batch, 4, 4, 512] => [batch, 8, 8, 256] (128, 2, 0), # [batch, 8, 8, 256] => [batch, 16, 16, 128] (64, 2, 0), # [batch, 16, 16, 128] => [batch, 32, 32, 64] ] return Generator('gen', kernels_gen_encoder, kernels_gen_decoder)
def generate(): train_phase = tf.placeholder(tf.bool) z = tf.placeholder(tf.float32, [None, 128]) G = Generator("generator") fake_img = G(z, train_phase) sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver( tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "generator")) saver.restore(sess, "./save_para/.\\model.ckpt") Z = np.random.standard_normal([NUMS_GEN, 128]) FAKE_IMG = sess.run(fake_img, feed_dict={z: Z, train_phase: False}) if not os.path.exists("./generate"): os.mkdir("./generate") for i in range(NUMS_GEN): Image.fromarray(np.uint8( (FAKE_IMG[i] + 1) * 127.5)).save("./generate/" + str(i) + ".jpg")
def main(): z_dim = 100 device = -1 #CPU batch_size = 1 model = Generator(z_dim) model.to_gpu() chainer.serializers.load_npz('result-dcgan/gen_snapshot_epoch-200.npz', model) model.to_cpu() x, _ = model.generate_noise(device, batch_size) y = model(x) graph = ChainerConverter().convert([x], [y]) exec_info = generate_descriptor("webassembly", graph) exec_info.save("./model")
def __init__(self, dataset_dir, log_dir, generator_channels, discriminator_channels, nz, style_depth, lrs, betas, eps, phase_iter, batch_size, n_cpu, opt_level): self.nz = nz self.dataloader = Dataloader(dataset_dir, batch_size, phase_iter * 2, n_cpu) self.generator = cuda( DataParallel(Generator(generator_channels, nz, style_depth))) self.discriminator = cuda( DataParallel(Discriminator(discriminator_channels))) self.tb = tensorboard.tf_recorder('StyleGAN', log_dir) self.phase_iter = phase_iter self.lrs = lrs self.betas = betas self.opt_level = opt_level
def _build_graph(self): with tf.device('/CPU: 0'): self.image = self._prepare_data() gen_args = self.args['generator'] gen_args['batch_size'] = self.batch_size self.generator = Generator('Generator', gen_args, self.graph, self.training, scope_prefix= self.name, log_tensorboard=self.log_tensorboard, log_params=self.log_params) self.gen_image = self.generator.image dis_args = self.args['discriminator'] self.real_discriminator = Discriminator('Discriminator', dis_args, self.graph, self.image, False, self.training, scope_prefix= self.name, log_tensorboard=self.log_tensorboard, log_params=self.log_params) self.fake_discriminator = Discriminator('Discriminator', dis_args, self.graph, self.gen_image, False, self.training, scope_prefix=self.name, log_tensorboard=False, log_params=False, reuse=True) self.gen_loss = self._generator_loss() self.dis_loss = self._discriminator_loss() self.gen_opt_op, _, _ = self.generator._optimization_op(self.gen_loss) self.dis_opt_op, _, _ = self.real_discriminator._optimization_op(self.dis_loss) with tf.device('/CPU: 0'): self._log_train_info()
def __init__(self, arg, device): self.device = device # network self.nc = arg.nc self.nz = arg.nz self.init_size = arg.init_size self.size = arg.size self.G = Generator(nc=self.nc, nz=self.nz, size=self.size) # pre-growing total_stages = int(math.log2(self.size / self.init_size)) + 1 for i in range(total_stages - 1): self.G.grow_network() self.G.flush_network() for param in self.G.parameters(): param.requires_grad_(False) # load checkpoint checkpoint = torch.load('checkpoint.tar') self.G.load_state_dict(checkpoint['G_EMA_state_dict']) self.G = self.G.to(self.device) self.G.eval()