def train(epoch): critic.train() generator.train() for batch_idx, (input_data, label) in enumerate(train_loader): if batch_idx > k: break critic_optimizer.zero_grad() input_data = ((input_data - 0.5) * 2).to(args.device) noise = torch.rand(len(input_data), latent_size) noise = noise.to(args.device) fake = generator(noise) g_loss = l * model.gradient_penalty(input_data, fake, critic) f_loss = critic(fake.detach()).mean() t_loss = -critic(input_data).mean() critic_loss = f_loss + t_loss + g_loss critic_loss.backward() critic_optimizer.step() writer.add_scalars('critic_loss_detail', { 'g loss': g_loss, 'f loss': f_loss, 't loss': t_loss }, epoch) writer.add_scalar('critic loss', critic_loss, epoch) generator_optimizer.zero_grad() noise = torch.rand(args.batch_size, latent_size) noise = noise.to(args.device) fake = generator(noise) generator_loss = -critic(fake).mean() generator_loss.backward() generator_optimizer.step() writer.add_scalar('gen loss', generator_loss, epoch) return critic_loss, generator_loss
def train(): if len(args.gpu_idx) > 1: multi_gpu = True gpu_list = [int(i) for i in args.gpu_idx.split(',')] else: multi_gpu = False writer = SummaryWriter(log_dir=log_path) if args.dataset == 'cifar10': dataset = datasets.CIFAR10(args.data_dir, download=True, transform=transforms.Compose([ transforms.Resize(args.imsize), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) ])) elif args.dataset == 'mnist': dataset = datasets.MNIST(args.data_dir, train=True, download=True, transform=transforms.Compose([ transforms.Scale(args.imsize), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) args.in_dim = 1 args.out_dim = 1 else: dataset = datasets.ImageFolder(args.data_dir, transform=transforms.Compose([ transforms.Resize(args.imsize), transforms.ToTensor(), transforms.Normalize( (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])) dataloader = torch.utils.data.DataLoader( dataset , batch_size = args.batch , \ shuffle = True , num_workers = args.worker) #device = torch.device() generator = model.Generator(args) discriminator = model.Discriminator(args) generator.apply(weights_init) discriminator.apply(weights_init) gan_criterion = nn.BCELoss() aux_criterion = nn.CrossEntropyLoss() #input_noise = torch.from_numpy( np.random.normal(0,1,[args.batch , args.dim_embed]) ) #input_label = torch.from_numpy( np.random.randint(0,args.num_class, [args.batch,1 ]) ) if args.l_smooth: # training strategy stated in improved GAN real_label = 0.9 fake_label = 0.1 else: real_label = 1.0 fake_label = 0.0 step = 0 if args.gpu: # acutally do nothing? because bce and cce don't have paramters gan_criterion = gan_criterion.cuda() aux_criterion = aux_criterion.cuda() generator = generator.cuda() discriminator = discriminator.cuda() if multi_gpu: print('multi gpu') generator = nn.DataParallel(generator, device_ids=gpu_list) discriminator = nn.DataParallel(discriminator, device_ids=gpu_list) opt_d = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) opt_g = optim.Adam(generator.parameters(), lr=args.lr, betas=(args.beta1, 0.999)) if os.path.isfile(os.path.join(ckpt_path, args.run_name + '.ckpt')): print('found ckpt file' + os.path.join(ckpt_path, args.run_name + '.ckpt')) ckpt = torch.load(os.path.join(ckpt_path, args.run_name + '.ckpt')) generator.load_state_dict(ckpt['generator']) discriminator.load_state_dict(ckpt['discriminator']) opt_d.load_state_dict(ckpt['opt_d']) opt_g.load_state_dict(ckpt['opt_g']) step = ckpt['step'] for i in range(args.epoch): for j, data in enumerate(dataloader): images, labels = data[0], data[1] batch = images.shape[0] input_noise = torch.from_numpy( np.random.normal(0, 1, [batch, args.dim_embed]).astype(np.float32)) input_label = torch.from_numpy( np.random.randint(0, args.num_class, [batch])) input_noise = np.random.normal( 0, 1, [batch, args.dim_embed]).astype(np.float32) class_onehot = np.zeros((batch, args.num_class)) class_onehot[np.arange(batch), input_label] = 1 input_noise[np.arange(batch), :args.num_class] = class_onehot[ np.arange(batch)] input_noise = torch.from_numpy(input_noise) real_target = torch.full((batch, 1), real_label) fake_target = torch.full((batch, 1), fake_label) aux_target = torch.autograd.Variable(labels) if args.gpu: input_noise = input_noise.cuda() input_label = input_label.cuda() images = images.cuda() labels = labels.cuda() real_target = real_target.cuda() fake_target = fake_target.cuda() aux_target = aux_target.cuda() # train generator # 好像不call也沒關係 opt_g.zero_grad() fake = generator(input_noise, input_label) gan_out_g, aux_out_g = discriminator(fake) if args.wgan: gan_loss_g = -torch.mean(gan_out_g) else: gan_loss_g = gan_criterion(gan_out_g, real_target) aux_loss_g = aux_criterion(aux_out_g, input_label) g_loss = (gan_loss_g + args.aux_weight * aux_loss_g) * 0.5 g_loss.backward() #opt_g.step() opt_g.step() # train discriminator with real samples opt_d.zero_grad() gan_out_r, aux_out_r = discriminator(images) if args.wgan: gan_loss_r = -torch.mean(gan_out_r) else: gan_loss_r = gan_criterion(gan_out_r, real_target) aux_loss_r = aux_criterion(aux_out_r, aux_target) d_real_loss = (gan_loss_r + args.aux_weight * aux_loss_r) / 2.0 # train discriminator with fake samples #fake = generator(input_noise , input_label).detach() gan_out_f, aux_out_f = discriminator(fake.detach()) if args.wgan: gan_loss_f = torch.mean(gan_out_f) else: gan_loss_f = gan_criterion(gan_out_f, fake_target) aux_loss_f = aux_criterion(aux_out_f, input_label) d_fake_loss = (gan_loss_f + args.aux_weight * aux_loss_f) / 2.0 if args.wgan and args.gp: gp = model.gradient_penalty(discriminator, images, fake, args.num_class, args.gpu) d_loss = 0.5 * (d_real_loss + d_fake_loss) + args.gp_weight * gp else: d_loss = (d_real_loss + d_fake_loss) / 2.0 d_loss.backward() opt_d.step() step = step + 1 if step % 100 == 0: writer.add_scalar('losses/g_loss', g_loss, step) writer.add_scalar('losses/d_loss', d_loss, step) grid = vutils.make_grid(fake.detach(), normalize=True) writer.add_image('generated', grid, step) if args.wgan and not args.gp: for p in discriminator.parameters(): p.data.clamp_(-args.clip, args.clip) if step % args.save_freq == 0: torch.save( { 'step': step, 'generator': generator.state_dict(), 'discriminator': discriminator.state_dict(), 'opt_d': opt_d.state_dict(), 'opt_g': opt_g.state_dict() }, os.path.join(ckpt_path, args.run_name + '.ckpt')) if step % args.sample_freq == 0: sample_generator(generator, input_noise, input_label, step) pred = np.concatenate( [aux_out_r.data.cpu().numpy(), aux_out_f.data.cpu().numpy()], axis=0) gt = np.concatenate( [labels.data.cpu().numpy(), input_label.data.cpu().numpy()], axis=0) d_acc = np.mean(np.argmax(pred, axis=1) == gt) print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f, acc: %d%%] [G loss: %f]" % \ (i, args.epoch , j, len(dataloader), d_loss.item(), 100.0 * d_acc, g_loss.item()))
step = ep * len(train_loader) + i + 1 D.train() G.train() # train D x = x.to(device) z = torch.randn(batch_size, z_dim).to(device) c = torch.tensor(np.eye(c_dim)[c_dense.cpu().numpy()], dtype=z.dtype).to(device) x_f = G(z, c).detach() x_gan_logit = D(x, c) x_f_gan_logit = D(x_f, c) d_x_gan_loss, d_x_f_gan_loss = d_loss_fn(x_gan_logit, x_f_gan_logit) gp = model.gradient_penalty(D, x, x_f, mode=gp_mode) d_loss = d_x_gan_loss + d_x_f_gan_loss + gp * gp_coef D.zero_grad() d_loss.backward() d_optimizer.step() writer.add_scalar('D/d_gan_loss', (d_x_gan_loss + d_x_f_gan_loss).data.cpu().numpy(), global_step=step) writer.add_scalar('D/gp', gp.data.cpu().numpy(), global_step=step) # train G if step % n_d == 0: z = torch.randn(batch_size, z_dim).to(device)
xr = torch.from_numpy(xr).cuda() # (32, 1, 24, 24) yr = torch.from_numpy(yr).cuda() # (32, 5) # xr.requires_grad_() # yr = torch.tensor(yr, dtype=torch.float) predr = D(xr, yr) # (32, 1) lossr = -predr.mean() zx = torch.randn((32, 100), dtype=torch.float).cuda() zx = torch.cat([zx, yr], dim=1) # (32, 105) # zx = torch.tensor(zx, dtype=torch.float) xf = G(zx).detach() # xf (32, 1, 24, 24) predf = D(xf, yr) lossf = predf.mean() gp = gradient_penalty(D, xr, xf.detach(), yr) loss_D = lossr + lossf + 10 * gp optim_D.zero_grad() loss_D.backward() optim_D.step() z = torch.randn((32, 100), dtype=torch.float).cuda() x, y = next(data_iter) # y = torch.tensor(y, dtype=torch.float) y = torch.from_numpy(y).cuda() zx = torch.cat([z, y], dim=1) # zx = torch.tensor(zx, dtype=torch.float) xf = G(zx) predf = D(xf, y)
# d loss d_r_loss, d_f_loss = d_loss_fn(r_logit, f_logit) d_f_tree_losses = tree_loss_fn(f_c_logit, c, mask) if att != '': d_r_tree_losses = tree_loss_fn(r_c_logit, c, mask) start = 1 if half_acgan else 0 d_tree_loss = sum([ d_f_tree_losses[i] * lambdas[i] * layer_mask[i] for i in range(start, len(lambdas)) ]) d_tree_loss += d_r_tree_losses[0] * lambdas[0] * layer_mask[0] else: d_tree_loss = sum( [d_f_tree_losses[i] * lambdas[i] for i in range(len(lambdas))]) gp = model.gradient_penalty(D, real, fake, gp_mode) d_loss = d_r_loss + d_f_loss + d_tree_loss + gp * 10.0 # g loss g_f_loss = g_loss_fn(f_logit) g_f_tree_losses = tree_loss_fn(f_c_logit, c, mask) g_tree_loss = sum([ g_f_tree_losses[i] * lambdas[i] * layer_mask[i] for i in range(len(lambdas)) ]) g_loss = g_f_loss + g_tree_loss # optims d_step = optim(learning_rate=lr_d).minimize( d_loss, var_list=tl.trainable_variables(includes='D')) g_step = optim(learning_rate=lr_g).minimize(
discriminate_image_img_s, discriminate_image_attr_s = D(s_decoder_ouput) print("D(image_tensor) - img = ", discriminate_image_img) print("D(image_tensor) - attr = ", discriminate_image_attr) print("D(s_decoder_output) - img = ", discriminate_image_img_s) print("D(s_decoder_output) - attr = ", discriminate_image_attr_s) # -------------------------------- CALCULATE LOSSES ---------------------------------------------------------- # Calculate discriminator loss wd = tf.reduce_mean(discriminate_image_img) - tf.reduce_mean( discriminate_image_img_s) discriminator_loss_gan = -wd # Determine the loss function given the original image and Generator output gradient_loss = model.gradient_penalty(D, image_tensor, s_decoder_ouput) # Discriminator loss for image attr image_attr_loss = tf.losses.sigmoid_cross_entropy(label_tensor, discriminate_image_attr) # Final loss for discriminator discriminator_loss = discriminator_loss_gan + gradient_loss * 10.0 + image_attr_loss # Calculate generator loss discriminator_label_loss = -tf.reduce_mean(discriminate_image_attr_s) # The losses the one the do any other way that label_decoder_loss = tf.losses.sigmoid_cross_entropy( shuffled_label_tensor, discriminate_image_attr_s) image_decoder_loss = tf.losses.absolute_difference(image_tensor,