def test_multi(self): """Translate images using StarGAN trained on multiple datasets.""" # Load the trained generator. self.restore_model(self.test_iters) with torch.no_grad(): for i, (x_real, c_org) in enumerate(self.celeba_loader): # Prepare input images and target domain labels. x_real = x_real.to(self.device) c_celeba_list = self.create_labels(c_org, self.c_dim, 'CelebA', self.selected_attrs) c_rafd_list = self.create_labels(c_org, self.c2_dim, 'RaFD') zero_celeba = torch.zeros(x_real.size(0), self.c_dim).to(self.device) # Zero vector for CelebA. zero_rafd = torch.zeros(x_real.size(0), self.c2_dim).to(self.device) # Zero vector for RaFD. mask_celeba = self.label2onehot(torch.zeros(x_real.size(0)), 2).to(self.device) # Mask vector: [1, 0]. mask_rafd = self.label2onehot(torch.ones(x_real.size(0)), 2).to(self.device) # Mask vector: [0, 1]. # Translate images. x_fake_list = [x_real] for c_celeba in c_celeba_list: c_trg = torch.cat([c_celeba, zero_rafd, mask_celeba], dim=1) x_fake_list.append(self.G(x_real, c_trg)) for c_rafd in c_rafd_list: c_trg = torch.cat([zero_celeba, c_rafd, mask_rafd], dim=1) x_fake_list.append(self.G(x_real, c_trg)) # Save the translated images. x_concat = torch.cat(x_fake_list, dim=3) result_path = os.path.join(self.result_dir, '{}-images.jpg'.format(i+1)) save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0) print('Saved real and fake images into {}...'.format(result_path))
def test(self): """Translate images using StarGAN trained on a single dataset.""" # Load the trained generator. self.restore_model(self.test_iters) # Set data loader. if self.dataset == 'CelebA': data_loader = self.celeba_loader elif self.dataset == 'RaFD': data_loader = self.rafd_loader with torch.no_grad(): for i, (x_real, c_org) in enumerate(data_loader): # Prepare input images and target domain labels. x_real = x_real.to(self.device) c_trg_list = self.create_labels(c_org, self.c_dim, self.dataset, self.selected_attrs) # Translate images. x_fake_list = [x_real] for c_trg in c_trg_list: x_fake_list.append(self.G(x_real, c_trg)) # Save the translated images. x_concat = torch.cat(x_fake_list, dim=3) result_path = os.path.join(self.result_dir, '{}-images.jpg'.format(i+1)) save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0) print('Saved real and fake images into {}...'.format(result_path))
def save_img_results(imgs_tcpu, fake_imgs, num_imgs, count, image_dir, summary_writer): num = cfg.TRAIN.VIS_COUNT # The range of real_img (i.e., self.imgs_tcpu[i][0:num]) # is changed to [0, 1] by function vutils.save_image real_img = imgs_tcpu[-1][0:num] vutils.save_image( real_img, '%s/real_samples.png' % (image_dir), normalize=True) real_img_set = vutils.make_grid(real_img).numpy() real_img_set = np.transpose(real_img_set, (1, 2, 0)) real_img_set = real_img_set * 255 real_img_set = real_img_set.astype(np.uint8) sup_real_img = summary.image('real_img', real_img_set) summary_writer.add_summary(sup_real_img, count) for i in range(num_imgs): fake_img = fake_imgs[i][0:num] # The range of fake_img.data (i.e., self.fake_imgs[i][0:num]) # is still [-1. 1]... vutils.save_image( fake_img.data, '%s/count_%09d_fake_samples%d.png' % (image_dir, count, i), normalize=True) fake_img_set = vutils.make_grid(fake_img.data).cpu().numpy() fake_img_set = np.transpose(fake_img_set, (1, 2, 0)) fake_img_set = (fake_img_set + 1) * 255 / 2 fake_img_set = fake_img_set.astype(np.uint8) sup_fake_img = summary.image('fake_img%d' % i, fake_img_set) summary_writer.add_summary(sup_fake_img, count) summary_writer.flush()
def sampleTrue(dataset, imageSize, dataroot, sampleSize, batchSize, saveFolder, workers=4): print('sampling real images ...') saveFolder = saveFolder + '0/' dataset = make_dataset(dataset, dataroot, imageSize) dataloader = torch.utils.data.DataLoader( dataset, shuffle=True, batch_size=batchSize, num_workers=int(workers)) if not os.path.exists(saveFolder): try: os.makedirs(saveFolder) except OSError: pass iter = 0 for i, data in enumerate(dataloader, 0): img, _ = data for j in range(0, len(img)): vutils.save_image(img[j].mul(0.5).add( 0.5), saveFolder + giveName(iter) + ".png") iter += 1 if iter >= sampleSize: break if iter >= sampleSize: break
def plot_rec(x, netEC, netEP, netD): x_c = x[0] x_p = x[np.random.randint(1, opt.max_step)] h_c = netEC(x_c) h_p = netEP(x_p) # print('h_c shape: ', h_c.shape) # print('h p shape: ', h_p.shape) rec = netD([h_c, h_p]) x_c, x_p, rec = x_c.data, x_p.data, rec.data fname = '%s/rec/rec_test.png' % (opt.log_dir) comparison = None for i in range(len(x_c)): if comparison is None: comparison = torch.stack([x_c[i], x_p[i], rec[i]]) else: new_comparison = torch.stack([x_c[i], x_p[i], rec[i]]) comparison = torch.cat([comparison, new_comparison]) print('comparison: ', comparison.shape) # row_sz = 5 # nplot = 20 # for i in range(0, nplot - row_sz, row_sz): # row = [[xc, xp, xr] for xc, xp, xr in zip(x_c[i:i + row_sz], x_p[i:i + row_sz], rec[i:i + row_sz])] # print('row: ', row) # to_plot.append(list(itertools.chain(*row))) # print(len(to_plot[0])) # utils.save_tensors_image(fname, comparison) if not os.path.exists(os.path.dirname(fname)): os.makedirs(os.path.dirname(fname)) save_image(comparison.cpu(), fname, nrow=3)
def test(self): """Facial attribute transfer on CelebA or facial expression synthesis on RaFD.""" # Load trained parameters G_path = os.path.join(self.model_save_path, '{}_G.pth'.format(self.test_model)) self.G.load_state_dict(torch.load(G_path)) self.G.eval() if self.dataset == 'CelebA': data_loader = self.celebA_loader else: data_loader = self.rafd_loader for i, (real_x, org_c) in enumerate(data_loader): real_x = self.to_var(real_x, volatile=True) if self.dataset == 'CelebA': target_c_list = self.make_celeb_labels(org_c) else: target_c_list = [] for j in range(self.c_dim): target_c = self.one_hot(torch.ones(real_x.size(0)) * j, self.c_dim) target_c_list.append(self.to_var(target_c, volatile=True)) # Start translations fake_image_list = [real_x] for target_c in target_c_list: fake_image_list.append(self.G(real_x, target_c)) fake_images = torch.cat(fake_image_list, dim=3) save_path = os.path.join(self.result_path, '{}_fake.png'.format(i+1)) save_image(self.denorm(fake_images.data), save_path, nrow=1, padding=0) print('Translated test images and saved into "{}"..!'.format(save_path))
def saver(state): if state[torchbearer.BATCH] == 0: data = state[torchbearer.X] recon_batch = state[torchbearer.Y_PRED] comparison = torch.cat([data[:num_images], recon_batch.view(128, 1, 28, 28)[:num_images]]) save_image(comparison.cpu(), str(folder) + 'reconstruction_' + str(state[torchbearer.EPOCH]) + '.png', nrow=num_images)
def sample_image(n_row, batches_done): """Saves a grid of generated digits ranging from 0 to n_classes""" # Sample noise z = Variable(FloatTensor(np.random.normal(0, 1, (n_row**2, opt.latent_dim)))) # Get labels ranging from 0 to n_classes for n rows labels = np.array([num for _ in range(n_row) for num in range(n_row)]) labels = Variable(LongTensor(labels)) gen_imgs = generator(z, labels) save_image(gen_imgs.data, 'images/%d.png' % batches_done, nrow=n_row, normalize=True)
def _train(self, epoch): """Perform the actual train.""" # put model into train mode self.d_model.train() # TODO: why? cp_loader = deepcopy(self.train_loader) if self.verbose: progress_bar = tqdm(total=len(cp_loader), desc='Current Epoch', file=sys.stdout, leave=False, ncols=75, position=0, unit=' Batch') else: progress_bar = None real_label = 1 fake_label = 0 for batch_idx, inputs in enumerate(cp_loader): # Update Discriminator network maximize log(D(x)) + log(1 - D(G(z))) # train with real self.optimizer_d.zero_grad() inputs = inputs.to(self.device) batch_size = inputs.size(0) outputs = self.d_model(inputs) label = torch.full((batch_size,), real_label, device=self.device) loss_d_real = self.loss_function(outputs, label) loss_d_real.backward() # train with fake noise = torch.randn((batch_size, self.g_model.nz, 1, 1,), device=self.device) fake_outputs = self.g_model(noise) label.fill_(fake_label) outputs = self.d_model(fake_outputs.detach()) loss_g_fake = self.loss_function(outputs, label) loss_g_fake.backward() self.optimizer_d.step() # (2) Update G network: maximize log(D(G(z))) self.g_model.zero_grad() label.fill_(real_label) outputs = self.d_model(fake_outputs) loss_g = self.loss_function(outputs, label) loss_g.backward() self.optimizer_g.step() if self.verbose: if batch_idx % 10 == 0: progress_bar.update(10) if self.out_f is not None and batch_idx % 100 == 0: fake = self.g_model(self.sample_noise) vutils.save_image( fake.detach(), '%s/fake_samples_epoch_%03d.png' % (self.out_f, epoch), normalize=True) if self.verbose: progress_bar.close()
def sample_images(batches_done): """Saves a generated sample from the test set""" imgs = next(iter(val_dataloader)) real_A = Variable(imgs['A'].type(Tensor)) fake_B = G_AB(real_A) real_B = Variable(imgs['B'].type(Tensor)) fake_A = G_BA(real_B) img_sample = torch.cat((real_A.data, fake_B.data, real_B.data, fake_A.data), 0) save_image(img_sample, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def save_image(img): post = transforms.Compose([transforms.Lambda(lambda x: x.mul_(1./255)), transforms.Normalize(mean=[-0.40760392, -0.45795686, -0.48501961], std=[1,1,1]), transforms.Lambda(lambda x: x[torch.LongTensor([2,1,0])]), #turn to RGB ]) img = post(img) img = img.clamp_(0,1) vutils.save_image(img, '%s/transfer.png' % (opt.outf), normalize=True) return
def reconstruction_loss(self, images, input, size_average=True): # Get the lengths of capsule outputs. v_mag = torch.sqrt((input**2).sum(dim=2)) # Get index of longest capsule output. _, v_max_index = v_mag.max(dim=1) v_max_index = v_max_index.data # Use just the winning capsule's representation (and zeros for other capsules) to reconstruct input image. batch_size = input.size(0) all_masked = [None] * batch_size for batch_idx in range(batch_size): # Get one sample from the batch. input_batch = input[batch_idx] # Copy only the maximum capsule index from this batch sample. # This masks out (leaves as zero) the other capsules in this sample. batch_masked = Variable(torch.zeros(input_batch.size())).cuda() batch_masked[v_max_index[batch_idx]] = input_batch[v_max_index[batch_idx]] all_masked[batch_idx] = batch_masked # Stack masked capsules over the batch dimension. masked = torch.stack(all_masked, dim=0) # Reconstruct input image. masked = masked.view(input.size(0), -1) output = self.relu(self.reconstruct0(masked)) output = self.relu(self.reconstruct1(output)) output = self.sigmoid(self.reconstruct2(output)) output = output.view(-1, self.image_channels, self.image_height, self.image_width) # Save reconstructed images occasionally. if self.reconstructed_image_count % 10 == 0: if output.size(1) == 2: # handle two-channel images zeros = torch.zeros(output.size(0), 1, output.size(2), output.size(3)) output_image = torch.cat([zeros, output.data.cpu()], dim=1) else: # assume RGB or grayscale output_image = output.data.cpu() vutils.save_image(output_image, "reconstruction.png") self.reconstructed_image_count += 1 # The reconstruction loss is the sum squared difference between the input image and reconstructed image. # Multiplied by a small number so it doesn't dominate the margin (class) loss. error = (output - images).view(output.size(0), -1) error = error**2 error = torch.sum(error, dim=1) * 0.0005 # Average over batch if size_average: error = error.mean() return error
def save_images(netG, fixed_noise, outputDir, epoch): ''' Generates a batch of images from the given 'noise'. Saves 64 of the generated samples to 'outputDir' system path. Inputs are the network (netG), a 'noise' input, system path to which images will be saved (outputDir) and current 'epoch'. ''' noise = Variable(fixed_noise) netG.eval() fake = netG(noise) netG.train() vutils.save_image( fake.data[0:64, :, :, :], '%s/fake_samples_epoch_%03d.png' % (outputDir, epoch), nrow=8)
def reconstruct_test(self,epoch): for i,batch in enumerate(self.test_dataloader): images = batch['image'] images = images.float() bumps = batch['bump'] bumps = bumps.float() masks = batch['mask'] masks = masks.float() images = Variable(images.cuda()) recon_mask, recon = self.Gnet.forward(images) output = torch.cat((masks,recon_mask.data.cpu(),bumps,recon.data.cpu()),dim=3) utils.save_image(output, net.outpath + '/'+str(epoch)+'.'+str(i)+'.jpg',nrow=4, normalize=True)
def sample_images(batches_done): """Saves a generated sample from the test set""" imgs = next(iter(val_dataloader)) X1 = Variable(imgs['A'].type(Tensor)) X2 = Variable(imgs['B'].type(Tensor)) _, Z1 = E1(X1) _, Z2 = E2(X2) fake_X1 = G1(Z2) fake_X2 = G2(Z1) img_sample = torch.cat((X1.data, fake_X2.data, X2.data, fake_X1.data), 0) save_image(img_sample, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def save(self, source, iteration): save_dir = os.path.join(self.model_dir, "gen_images") if os.path.exists(save_dir) == False: os.mkdir(save_dir) images_file = os.path.join(save_dir, "image_{}.png".format(iteration)) if self.cuda: source = source.cuda() source = Variable(source) outputs = self.gen_model(source) vutils.save_image(outputs.cpu().data, images_file, normalize=True)
def pp_interp(net, alpha): """ Only works with model_resnet_preproc.py as your architecture!!! """ conv2d = net.d.preproc deconv2d = nn.ConvTranspose2d(16, 3, 3, stride=1, padding=1) deconv2d = deconv2d.cuda() deconv2d.weight = conv2d.weight gz1 = net.sample(bs=128) gz2 = net.sample(bs=128) #alpha = net.sample_lambda(gz1.size(0)) gz_mix = alpha*gz1 + (1.-alpha)*gz2 save_image(gz1*0.5 + 0.5, filename="gz1.png") save_image(gz2*0.5 + 0.5, filename="gz2.png") save_image(gz_mix*0.5 + 0.5, filename="gz_mix.png") # Ok, do the mixup in hidden space. gz1_h = conv2d(gz1) gz2_h = conv2d(gz2) #alpha = 0.05 gz_mix_h = alpha*gz1_h + (1.-alpha)*gz2_h gz_mix_h_dec = deconv2d(gz_mix_h) save_image(gz_mix_h_dec*0.5 + 0.5, filename="gz_mix_h_dec.png") print(conv2d.weight == deconv2d.weight) import pdb pdb.set_trace()
def closure(): optimizer.zero_grad() out = resnet(generated) style_loss = [GramMSELoss().cuda()(out[i],style_target[i])*style_weight[i] for i in range(len(style_target))] content_loss = nn.MSELoss().cuda()(out[content_layer_num],content_target) total_loss = 1000 * sum(style_loss) + sum(content_loss) total_loss.backward() if iteration[0] % 100 == 0: print(total_loss) v_utils.save_image(image_postprocess(generated.data),"./gen_{}.png".format(iteration[0])) iteration[0] += 1 return total_loss
def generate(self, input_sample=None): if input_sample is None: input_sample = torch.randn(self.gen_training_result[1], self.nz, 1, 1, device=self.device) if not isinstance(input_sample, torch.Tensor) and \ isinstance(input_sample, np.ndarray): input_sample = torch.from_numpy(input_sample) if not isinstance(input_sample, torch.Tensor) and \ not isinstance(input_sample, np.ndarray): raise TypeError("Input should be a torch.tensor or a numpy.ndarray") self.net_g.eval() with torch.no_grad(): input_sample = input_sample.to(self.device) generated_fake = self.net_g(input_sample) vutils.save_image(generated_fake.detach(), '%s/evaluation.png' % self.gen_training_result[0], normalize=True)
def test(self, epoch): self.set_train(is_train=False) test_loss = 0 for i, (x, _) in enumerate(self.test_loader): with torch.no_grad(): recon_x = self.model_eval(x)[0] test_loss += self.compute_loss_and_gradient(x) if i == 0: n = min(x.size(0), 8) comparison = torch.cat([x[:n], recon_x.reshape(self.args.batch_size, 1, 28, 28)[:n]]) save_image(comparison.detach().cpu(), os.path.join(OUTPUT_DIR, 'reconstruction_' + str(epoch) + '.png'), nrow=n) test_loss /= len(self.test_loader.dataset) print('Test set loss: {:.4f}'.format(test_loss))
def test(epoch): model.eval() test_loss = 0 with torch.no_grad(): for i, (data, _) in enumerate(test_loader): data = data.to(device) recon_batch, mu, logvar = model(data) test_loss += loss_function(recon_batch, data, mu, logvar).item() if i == 0: n = min(data.size(0), 8) comparison = torch.cat([data[:n], recon_batch.view(args.batch_size, 1, 28, 28)[:n]]) save_image(comparison.cpu(), 'results/reconstruction_' + str(epoch) + '.png', nrow=n) test_loss /= len(test_loader.dataset) print('====> Test set loss: {:.4f}'.format(test_loss))
def test(epoch): model.eval() test_loss = 0 for i, (data, _) in enumerate(test_loader): if args.cuda: data = data.cuda() data = Variable(data, volatile=True) recon_batch, mu, logvar = model(data) test_loss += loss_function(recon_batch, data, mu, logvar).data[0] if i == 0: n = min(data.size(0), 8) comparison = torch.cat([data[:n], recon_batch.view(args.batch_size, 1, 28, 28)[:n]]) save_image(comparison.data.cpu(), 'results/reconstruction_' + str(epoch) + '.png', nrow=n) test_loss /= len(test_loader.dataset) print('====> Test set loss: {:.4f}'.format(test_loss))
def __call__(self, image): self.iteration += 1 if self.noiseLevel > 0.0: a = image.view(-1) numNoiseBits = int(a.shape[0] * self.noiseLevel) permutedIndices = np.random.permutation(a.shape[0]) a[permutedIndices[0:numNoiseBits // 2]] = self.highValue a[permutedIndices[numNoiseBits // 2:numNoiseBits]] = self.lowValue # Save a subset of the images for debugging if self.logDir is not None: if np.random.random() <= self.logProbability: outfile = os.path.join(self.logDir, "im_noise_" + str(int(self.noiseLevel * 100)) + "_" + str(self.iteration).rjust(6, '0') + ".png") save_image(image, outfile) return image
def plot_analogy(x, netEC, netEP, netD): x_c = x[0] h_c = netEC(x_c) nrow = opt.batch_size row_sz = opt.max_step to_plot = [] zeros = torch.zeros(opt.channels, opt.image_width, opt.image_width) # to_plot.append(zeros) for i in range(nrow): to_plot.append(x[0][i]) to_plot = torch.stack(to_plot) for j in range(0, row_sz): h_p = netEP(x[j]) for i in range(nrow): h_p[i] = h_p[0] rec = netD([h_c, Variable(h_p)]) print('rec shape: ', rec.shape) to_plot = torch.cat([to_plot, rec]) originals = [] for i in range(row_sz): originals.append(x[i][0]) originals = [x[0][0]] + originals print('original: ', len(originals)) print('len(to_plot): ', len(to_plot)) plt_list = [] for i in range(len(to_plot)): if i % nrow == 0: print('int(i / nrow): ', int(i / nrow)) plt_list.append(originals[int(i / nrow)]) plt_list.append(to_plot[i]) to_plot = torch.stack(plt_list) fname = '%s/rec/analogy_test.png' % (opt.log_dir) if not os.path.exists(os.path.dirname(fname)): os.makedirs(os.path.dirname(fname)) save_image(to_plot, fname, nrow=(nrow+1))
def sample_images(batches_done): """Saves a generated sample from the validation set""" imgs = next(iter(val_dataloader)) img_samples = None for img_A, img_B in zip(imgs['A'], imgs['B']): # Repeat input image by number of channels real_A = img_A.view(1, *img_A.shape).repeat(8, 1, 1, 1) real_A = Variable(real_A.type(Tensor)) # Get interpolated noise [-1, 1] sampled_z = np.repeat(np.linspace(-1, 1, 8)[:, np.newaxis], opt.latent_dim, 1) sampled_z = Variable(Tensor(sampled_z)) # Generator samples fake_B = generator(real_A, sampled_z) # Concatenate samples horisontally fake_B = torch.cat([x for x in fake_B.data.cpu()], -1) img_sample = torch.cat((img_A, fake_B), -1) img_sample = img_sample.view(1, *img_sample.shape) # Cocatenate with previous samples vertically img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2) save_image(img_samples, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def sample_images(batches_done): """Saves a generated sample from the validation set""" imgs = next(iter(val_dataloader)) img_samples = None for img1, img2 in zip(imgs['A'], imgs['B']): # Create copies of image X1 = img1.unsqueeze(0).repeat(opt.style_dim, 1, 1, 1) X1 = Variable(X1.type(Tensor)) # Get interpolated style codes s_code = np.repeat(np.linspace(-1, 1, opt.style_dim)[:, np.newaxis], opt.style_dim, 1) s_code = Variable(Tensor(s_code)) # Generate samples c_code_1, _ = Enc1(X1) X12 = Dec2(c_code_1, s_code) # Concatenate samples horisontally X12 = torch.cat([x for x in X12.data.cpu()], -1) img_sample = torch.cat((img1, X12), -1).unsqueeze(0) # Concatenate with previous samples vertically img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2) save_image(img_samples, 'images/%s/%s.png' % (opt.dataset_name, batches_done), nrow=5, normalize=True)
def test(self): self.load_checkpoint() self.G.to(self.device) tqdm_loader = tqdm(self.data_loader.test_loader, total=self.data_loader.test_iterations, desc='Testing at checkpoint {}'.format(self.config.checkpoint)) self.G.eval() with torch.no_grad(): for i, (x_real, c_org) in enumerate(tqdm_loader): x_real = x_real.to(self.device) c_trg_list = self.create_labels(c_org, self.config.attrs) x_fake_list = [x_real] for c_trg in c_trg_list: attr_diff = c_trg - c_org x_fake_list.append(self.G(x_real, attr_diff.to(self.device))) x_concat = torch.cat(x_fake_list, dim=3) result_path = os.path.join(self.config.result_dir, 'sample_{}.jpg'.format(i + 1)) save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0)
def stylization(p_wct, content_image_path, style_image_path, content_seg_path, style_seg_path, output_image_path, cuda): # Load image cont_img = Image.open(content_image_path).convert('RGB') styl_img = Image.open(style_image_path).convert('RGB') try: cont_seg = Image.open(content_seg_path) styl_seg = Image.open(style_seg_path) except: cont_seg = [] styl_seg = [] cont_img = transforms.ToTensor()(cont_img).unsqueeze(0) styl_img = transforms.ToTensor()(styl_img).unsqueeze(0) if cuda: cont_img = cont_img.cuda(0) styl_img = styl_img.cuda(0) p_wct.cuda(0) cont_img = Variable(cont_img, volatile=True) styl_img = Variable(styl_img, volatile=True) cont_seg = np.asarray(cont_seg) styl_seg = np.asarray(styl_seg) with Timer("Elapsed time in stylization: %f"): stylized_img = p_wct.transform(cont_img, styl_img, cont_seg, styl_seg) utils.save_image(stylized_img.data.cpu().float(), output_image_path, nrow=1) with Timer("Elapsed time in propagation: %f"): out_img = p_pro.process(output_image_path, content_image_path) out_img.save(output_image_path) if not cuda: print("NotImplemented: The CPU version of smooth filter has not been implemented currently.") return with Timer("Elapsed time in post processing: %f"): out_img = smooth_filter(output_image_path, content_image_path, f_radius=15, f_edge=1e-1) out_img.save(output_image_path)
def sampleFake(netG, nz, sampleSize, batchSize, saveFolder): print('sampling fake images ...') saveFolder = saveFolder + '0/' try: os.makedirs(saveFolder) except OSError: pass noise = torch.FloatTensor(batchSize, nz, 1, 1).cuda() iter = 0 for i in range(0, 1 + sampleSize // batchSize): noise.data.normal_(0, 1) fake = netG(noise) for j in range(0, len(fake.data)): if iter < sampleSize: vutils.save_image(fake.data[j].mul(0.5).add( 0.5), saveFolder + giveName(iter) + ".png") iter += 1 if iter >= sampleSize: break
def test(self): """ :return: """ self.G.eval() for i, (input, target) in enumerate(self.dataloader): input_A = Variable(input.cuda()) input_B = Variable(target.cuda()) # make downsampled input self.real_imgs_A = [input_A] for _ in range(self.num_D - 1): real_img = self.downsample(self.real_imgs_A[-1]) self.real_imgs_A.append(real_img) self.real_imgs_A.reverse() self.real_imgs_B = [input_B] for _ in range(self.num_D - 1): real_img = self.downsample(self.real_imgs_B[-1]) self.real_imgs_B.append(real_img) self.real_imgs_B.reverse() # (1) generate fake images self.fake_imgs = self.G(input_A) vutils.save_image(input_A.data, '%s/real_samples_A_%03d.png' % (self.config['outf_test'], i), normalize=True, padding=0) vutils.save_image(input_B.data, '%s/real_samples_B_%03d.png' % (self.config['outf_test'], i), normalize=True, padding=0) vutils.save_image(self.fake_imgs[0].data, '%s/fake_samples_0_%3d.png' % (self.config['outf_test'], i), normalize=True, padding=0) vutils.save_image(self.fake_imgs[1].data, '%s/fake_samples_1_%03d.png' % (self.config['outf_test'], i), normalize=True, padding=0) vutils.save_image(self.fake_imgs[2].data, '%s/fake_samples_2_%03d.png' % (self.config['outf_test'], i), normalize=True, padding=0)
# For KL divergence, see Appendix B in VAE paper or http://yunjey47.tistory.com/43 reconst_loss = F.binary_cross_entropy(x_reconst, x, size_average=False) kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) # Backprop and optimize loss = reconst_loss + kl_div optimizer.zero_grad() loss.backward() optimizer.step() if (i + 1) % 10 == 0: print( "Epoch[{}/{}], Step [{}/{}], Reconst Loss: {:.4f}, KL Div: {:.4f}" .format(epoch + 1, num_epochs, i + 1, len(data_loader), reconst_loss.item(), kl_div.item())) with torch.no_grad(): # Save the sampled images z = torch.randn(batch_size, z_dim).to(device) out = model.decode(z).view(-1, 1, 28, 28) save_image( out, os.path.join(sample_dir, 'sampled-{}.png'.format(epoch + 1))) # Save the reconstructed images out, _, _ = model(x) x_concat = torch.cat([x.view(-1, 1, 28, 28), out.view(-1, 1, 28, 28)], dim=3) save_image( x_concat, os.path.join(sample_dir, 'reconst-{}.png'.format(epoch + 1)))
# ---------------------- # Train Discriminators # ---------------------- optimizer_D.zero_grad() # Determine validity of real and generated images validity1_real, validity2_real = coupled_discriminators(imgs1, imgs2) validity1_fake, validity2_fake = coupled_discriminators(gen_imgs1.detach(), gen_imgs2.detach()) d_loss = ( adversarial_loss(validity1_real, valid) + adversarial_loss(validity1_fake, fake) + adversarial_loss(validity2_real, valid) + adversarial_loss(validity2_fake, fake) ) / 4 d_loss.backward() optimizer_D.step() print( "[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" % (epoch, opt.n_epochs, i, len(dataloader1), d_loss.item(), g_loss.item()) ) batches_done = epoch * len(dataloader1) + i if batches_done % opt.sample_interval == 0: gen_imgs = torch.cat((gen_imgs1.data, gen_imgs2.data), 0) save_image(gen_imgs, "images/%d.png" % batches_done, nrow=8, normalize=True)
def test(self): """ Test GANomaly model. Args: dataloader ([type]): Dataloader for the test set Raises: IOError: Model weights not found. """ with torch.no_grad(): # Load the weights of netg and netd. if self.opt.load_weights: path = "./output/{}/{}/train/weights/netG.pth".format(self.name.lower(), self.opt.dataset) pretrained_dict = torch.load(path)['state_dict'] try: self.netg.load_state_dict(pretrained_dict) except IOError: raise IOError("netG weights not found") print(' Loaded weights.') self.opt.phase = 'test' #self.opt.showProcess.setValue(80) # Create big error tensor for the test set. self.an_scores = torch.zeros(size=(len(self.dataloader['test'].dataset),), dtype=torch.float32, device=self.device) self.gt_labels = torch.zeros(size=(len(self.dataloader['test'].dataset),), dtype=torch.long, device=self.device) self.latent_i = torch.zeros(size=(len(self.dataloader['test'].dataset), self.opt.nz), dtype=torch.float32, device=self.device) self.latent_o = torch.zeros(size=(len(self.dataloader['test'].dataset), self.opt.nz), dtype=torch.float32, device=self.device) # print(" Testing model %s." % self.name) self.times = [] self.total_steps = 0 epoch_iter = 0 for i, data in enumerate(self.dataloader['test'], 0): self.total_steps += self.opt.batchsize epoch_iter += self.opt.batchsize time_i = time.time() self.set_input(data) self.fake, latent_i, latent_o = self.netg(self.input) error = torch.mean(torch.pow((latent_i-latent_o), 2), dim=1) time_o = time.time() self.an_scores[i*self.opt.batchsize : i*self.opt.batchsize+error.size(0)] = error.reshape(error.size(0)) self.gt_labels[i*self.opt.batchsize : i*self.opt.batchsize+error.size(0)] = self.gt.reshape(error.size(0)) self.latent_i [i*self.opt.batchsize : i*self.opt.batchsize+error.size(0), :] = latent_i.reshape(error.size(0), self.opt.nz) self.latent_o [i*self.opt.batchsize : i*self.opt.batchsize+error.size(0), :] = latent_o.reshape(error.size(0), self.opt.nz) self.times.append(time_o - time_i) # Save test images. if self.opt.save_test_images: dst = os.path.join(self.opt.outf, self.opt.name, 'test', 'images') if not os.path.isdir(dst): os.makedirs(dst) real, fake, _ = self.get_current_images() vutils.save_image(real, '%s/real_%03d.eps' % (dst, i+1), normalize=True) vutils.save_image(fake, '%s/fake_%03d.eps' % (dst, i+1), normalize=True) # Measure inference time. self.times = np.array(self.times) self.times = np.mean(self.times[:100] * 1000) # Scale error vector between [0, 1] print(torch.min(self.an_scores)) print(torch.max(self.an_scores)) maxNUM = torch.max(self.an_scores) minNUM = torch.min(self.an_scores) self.an_scores = (self.an_scores - torch.min(self.an_scores)) / (torch.max(self.an_scores) - torch.min(self.an_scores)) # auc, eer = roc(self.gt_labels, self.an_scores) # -------------- 处理阈值 ------------------ print('-------------- 处理阈值 ------------------') print(len(self.gt_labels)) plt.ion() scores = {} ##plt.ion() # Create data frame for scores and labels. scores['scores'] = self.an_scores scores['labels'] = self.gt_labels hist = pd.DataFrame.from_dict(scores) #hist.to_csv("histogram.csv") # Filter normal and abnormal scores. abn_scr = hist.loc[hist.labels == 1]['scores'] nrm_scr = hist.loc[hist.labels == 0]['scores'] # Create figure and plot the distribution. ##fig, axes = plt.subplots(figsize=(4, 4)) b = [] c = [] # for i in range(1000): # b.append(nrm_scr[i]) # for j in range(1000, 3011): # c.append(abn_scr[j]) print('asasddda') print(len(nrm_scr)) print(len(abn_scr)) for i in nrm_scr: b.append(i) for j in abn_scr: c.append(j) ##sns.distplot(nrm_scr, label=r'Normal Scores', color='r', bins=100, hist=True) ##sns.distplot(abn_scr, label=r'Abnormal Scores', color='b', bins=100, hist=True) nrm = np.zeros((50), dtype=np.int) minfix = 0.4 abn = np.zeros((50), dtype=np.int) abmin = 30 for k in np.arange(0, 1, 0.02): kint = int(k * 50) for j in range(len(nrm_scr)): if b[j] >= k and b[j] < (k + 0.02): nrm[kint] = nrm[kint] + 1 for j in range(len(abn_scr)): if c[j] >= k and c[j] < (k + 0.02): abn[kint] = abn[kint] + 1 print(nrm) print(abn) # startInd = 3 # for k in range(0,20): # if abs(nrm[k] - abn[k]) <= 3: # continue # else: # startInd = k # max_dist = (len(nrm) + len(abn))*0.28 # for k in range(startInd, 20): # if abs(nrm[k] - abn[k]) < 5: # #max_dist = abs(nrm[k] - abn[k]) # minfix = round((k / 20) + 0.02, 3) # break # for k in range(3, 17): # # print(nrm[k]) # # print(abn[k]) # # print('----') # if abs(nrm[k]-abn[k]) > abmin and not (nrm[k] == 0 and abn[k] == 0): # abmin = abs(nrm[k] - abn[k]) # minfix = round((k / 20) + 0.02, 3) max_dist = (len(nrm) + len(abn)) * 0.25 for k in range(0,50): num1 = np.sum(nrm[0:k]) num2 = np.sum(abn[k::]) if (num1 + num2) >= max_dist: minfix = round((k / 50) + 0.05, 3) max_dist = num1+num2 proline = minfix print(proline) print(self.gt_labels[0:20]) print(self.an_scores[0:20]) print('------------- 处理阈值 END --------------') # ------------- 处理阈值 END -------------- auc = evaluate(self.gt_labels, self.an_scores, metric=self.opt.metric) performance = OrderedDict([('Avg Run Time (ms/batch)', self.times), ('AUC', auc)]) if self.opt.display_id > 0 and self.opt.phase == 'test': counter_ratio = float(epoch_iter) / len(self.dataloader['test'].dataset) self.visualizer.plot_performance(self.epoch, counter_ratio, performance) # --- 写入文件 --- dict_info = {} dict_info['minVal'] = float(minNUM.item()) dict_info['maxVal'] = float(maxNUM.item()) dict_info['proline'] = float(proline) dict_info['auc'] = float(auc) dict_info['Avg Run Time (ms/batch)'] = float(self.times) #self.opt.showText.append(str(performance)); #self.opt.showProcess.setValue(100) return performance, dict_info
x_reconst, mu, log_var = model(x) # Compute reconstruction loss and kl divergence # For KL divergence, see Appendix B in VAE paper or http://yunjey47.tistory.com/43 reconst_loss = F.binary_cross_entropy(x_reconst, x, reduction='sum') kl_div = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) loss = reconst_loss + kl_div optimizer.zero_grad() loss.backward() optimizer.step() total_reconst_loss += reconst_loss.item() total_kl_loss += kl_div.item() total_loss += loss.item() print( "Epoch[{}/{}]], Total Loss: {:.4f}, Reconst Loss: {:.4f}, KL Div: {:.4f}" .format(epoch + 1, num_epochs, total_loss / len(mnist), total_reconst_loss / len(mnist), total_kl_loss / len(mnist))) with torch.no_grad(): # Save the sampled images z = torch.randn(batch_size, z_dim).to(device) out = model.decode(z).view(-1, 1, 28, 28) save_image( out, os.path.join(sample_dir, 'sampled-{}.png'.format(epoch + 1))) # Save the reconstructed images save_image( x_reconst.view(-1, 1, 28, 28), os.path.join(sample_dir, 'reconst-{}.png'.format(epoch + 1)))
import torch from utils1 import * from torchvision.utils import save_image import torch.nn as nn from sagan_models import Generator_SA, Discriminator_SA from igan_models import Generator_INV, Discriminator_INV from dcgan_models import Generator_DC, Discriminator_DC from gan_models import Generator_MLP, Discriminator_MLP path = 'models/mnist_sagan_5.5/' listdir = os.listdir(path) for i, pth in enumerate(listdir): if "_G" in pth: epoch = pth.split("_G")[0] print(i) model = torch.load(os.path.join(path, pth)) # print(model) torch.manual_seed(0) # print(torch.randn(1,128)) if not os.path.exists(os.path.join(path, 'img/' + epoch)): os.makedirs(os.path.join(path, 'img/' + epoch)) for i in range(1000): fixed_z = tensor2var(torch.randn(1, 128)) fake_images, _, _ = model(fixed_z) fake_images_new = fake_images[0].view(1, 3, 64, 64) save_image( denorm(fake_images_new.data), os.path.join(path, 'img/' + epoch + '/' + str(i) + '.jpg'))
d_loss.backward() d_opt.step() z = tc.randn(mini_batch, latent_sz).view(-1, latent_sz, 1, 1).to(device) fake_images = G(z) D_result = D(fake_images).squeeze() g_loss = loss_func(D_result, real_label) G.zero_grad() g_loss.backward() g_opt.step() if step % 200 == 0: print( 'epoch {}/{}, step {}, d_loss {:.4f}, g_loss {:.4f}, Real_score {:.2f}, Fake_score {:.2f}' .format(ep, epoch_sz, step, d_loss.item(), g_loss.item(), D(images).mean().item(), D(fake_images).mean().item())) if ep + 1 == 1: out = images out = image_range(out) save_image(out, os.path.join(result_path, 'real_img.png')) out = fake_images out = image_range(out) save_image(out, os.path.join(result_path, 'fake_img {}.png'.format(ep))) tc.save(G.state_dict(), 'G.ckpt') tc.save(D.state_dict(), 'D.ckpt')
def train(self): print(f'''Starting training: Epochs: {self.epochs} Batch size: {self.bs} Learning rate: {self.lr} Checkpoints: {self.save_cp} Device: {self.device.type} ''') if self.save_cp: try: if not os.path.exists(os.path.join(self.outdir, 'ckpt')): os.makedirs(os.path.join(self.outdir, 'ckpt')) print('Created checkpoint directory') except OSError: pass optimizerG = optim.Adam(self.netG.parameters(), lr=self.lr, betas=(0.9, 0.999)) optimizerD = optim.Adam(self.netD.parameters(), lr=self.lr, betas=(0.9, 0.999)) criterion = nn.BCELoss() criterionL1 = torch.nn.L1Loss().to(self.device) criterionL2 = torch.nn.MSELoss().to(self.device) criterionGAN = GANLoss('vanilla').to(self.device) # criterionGAN = nn.BCEWithLogitsLoss().to(self.device) # torch.autograd.set_detect_anomaly(True) for epoch in range(self.start_epoch, self.epochs): print(f'D is trained {max(1, 5-epoch)} times') for i, data in enumerate(self.dataloader): for _ in range(1): #max(1, 5-epoch) self.netD.zero_grad() self.set_requires_grad(self.netD, True) ############################ # Loss_D: L_D = -(log(D(y) + log(1 - D(G(x)))) ########################### # forward y = data[0].to(self.device) # 320 x 180 x = self.gen_x(y).to(self.device) # 320 x 180 Gx = self.netG(x) # 320 x 180 # train with fake fake_pair = torch.cat((x, Gx), 1) D_Gx = self.netD(fake_pair.detach()) errD_fake = criterionGAN(D_Gx, False) # train with real real_pair = torch.cat((x, y), 1) D_y = self.netD(real_pair) errD_real = criterionGAN(D_y, True) # backprop errD = (errD_real + errD_fake) * 0.5 errD.backward() optimizerD.step() self.netG.zero_grad() self.set_requires_grad( self.netD, False) # D requires no gradients when optimizing G ############################ # Loss_G_GAN: L_G = -log(D(G(x)) # Fake the D ########################### D_Gx = self.netD(fake_pair) errG_GAN = criterionGAN(D_Gx, True) ############################ # Loss_G_C: L_C = ||y - G(x)||_1 ########################### errG_C = criterionL2(Gx, y) * 50 # errG_CM = criterionL2((1-self.mask)*Gx, (1-self.mask)*y) * 50 # backprop errG = errG_GAN + errG_C errG.backward() optimizerG.step() if i % 100 == 0: print( f'[{epoch}/{self.epochs}][{i}/{len(self.dataloader)}] LossD_real: {errD_real.item():.4f} LossD_fake: {errD_fake.item():.4f} LossG_GAN: {errG_GAN.item():.4f} LossG_C: {errG_C.item():.4f}' ) # Log print( f'[{epoch}/{self.epochs}][{i}/{len(self.dataloader)}] LossD: {errD.item():.4f} LossG: {errG.item():.4f}' ) if self.save_cp: torch.save(self.netG.state_dict(), f'{self.outdir}/ckpt/G_epoch{epoch}.pth') torch.save(self.netD.state_dict(), f'{self.outdir}/ckpt/D_epoch{epoch}.pth') vutils.save_image(Gx.detach(), '%s/impainted_samples_epoch_%03d.png' % (self.outdir, epoch), normalize=True)
# model = load_local_or_remote_file(model_path) state_dict = torch.load(model_path) model.load_state_dict(state_dict) model.to(ptu.device) model.eval() for color in ["grey"]: reference_path = "demos/door_demos_v3/demo_v3_%s_0.pkl"%color traj = np.load("demos/door_demos_v3/demo_v3_%s_0.pkl"%color, allow_pickle=True)[0] goal_image_flat = traj["observations"][-1]["image_observation"] goal_image = goal_image_flat.reshape(1, 3, 500, 300).transpose([0, 1, 3, 2]) / 255.0 # goal_image = goal_image[:, ::-1, :, :].copy() # flip bgr goal_image = goal_image[:, :, 60:300, 30:470] goal_image_pt = ptu.from_numpy(goal_image) save_image(goal_image_pt.data.cpu(), 'goal.png', nrow=1) goal_latent = model.encode(goal_image_pt).detach().cpu().numpy().flatten() initial_image_flat = traj["observations"][0]["image_observation"] initial_image = initial_image_flat.reshape(1, 3, 500, 300).transpose([0, 1, 3, 2]) / 255.0 # initial_image = initial_image[:, ::-1, :, :].copy() # flip bgr initial_image = initial_image[:, :, 60:300, 30:470] initial_image_pt = ptu.from_numpy(initial_image) save_image(initial_image_pt.data.cpu(), 'initial.png', nrow=1) initial_latent = model.encode(initial_image_pt).detach().cpu().numpy().flatten() print("Finished initial_latent") reward_params = dict( goal_latent=goal_latent, initial_latent=initial_latent, goal_image=goal_image_flat, initial_image=initial_image_flat,
def train(): dir_name, save_path = getSavePath() netG, netD = networks.getGD_SN(args.structure, args.dataset, args.image_size, args.num_features, dim_z=args.input_dim, bottleneck=args.bottleneck) if args.ema_trick: ema_netG_9999 = copy.deepcopy(netG) if args.reload > 0: netG.load_state_dict(torch.load(save_path + 'G_epoch{}.pth'.format(args.reload))) netD.load_state_dict(torch.load(save_path + 'D_epoch{}.pth'.format(args.reload))) if args.ema_trick: ema_netG_9999.load_state_dict( torch.load(save_path + 'emaG0.9999_epoch{}.pth'.format(args.reload), map_location=torch.device('cpu'))) netG.cuda() netD.cuda() g_optimizer = torch.optim.Adam(netG.parameters(), lr=args.g_lr, betas=(args.beta1, args.beta2)) d_optimizer = torch.optim.Adam(netD.parameters(), lr=args.d_lr, betas=(args.beta1, args.beta2)) g_losses, d_losses = [], [] grad_normD, grad_normG = [], [] loader = datasets.getDataLoader(args.dataset, args.image_size, batch_size=args.batch_size) data_iter = iter(loader) for i in range(1, args.num_iters+1): if i >= args.lr_decay_start: utils.decay_lr(g_optimizer, args.num_iters, args.lr_decay_start, args.g_lr) utils.decay_lr(d_optimizer, args.num_iters, args.lr_decay_start, args.d_lr) if i <= args.reload: continue if i == 1: torch.save(netG.state_dict(), save_path + 'G_epoch0.pth') torch.save(netD.state_dict(), save_path + 'D_epoch0.pth') # G-step for _ in range(args.g_freq): try: x = next(data_iter)[0].cuda().float() except StopIteration: data_iter = iter(loader) x = next(data_iter)[0].cuda().float() z = torch.randn(args.batch_size, args.input_dim, device=device) g_optimizer.zero_grad() x_hat = netG(z) y_hat = netD(x_hat) y = netD(x) g_loss = get_gloss(args.losstype, y_hat, y) g_losses.append(g_loss.item()) g_loss.backward() g_optimizer.step() grad_normG.append(utils.getGradNorm(netG)) if args.ema_trick: utils.soft_copy_param(ema_netG_9999, netG, 0.9999) for _ in range(args.d_freq): try: x = next(data_iter)[0].cuda().float() except StopIteration: data_iter = iter(loader) x = next(data_iter)[0].cuda().float() z = torch.randn(args.batch_size, args.input_dim, device=device) d_optimizer.zero_grad() x_hat = netG(z).detach() y_hat = netD(x_hat) y = netD(x) d_loss = get_dloss(args.losstype, y_hat, y) d_losses.append(d_loss.item()) d_loss.backward() d_optimizer.step() grad_normD.append(utils.getGradNorm(netD)) netD.proj.weight.data = F.normalize(netD.proj.weight.data, dim=1) if i % args.print_freq == 0: print('Iteration: {}; G-Loss: {}; D-Loss: {};'.format(i, g_loss, d_loss)) if i == 1: save_image((x / 2. + 0.5), os.path.join(dir_name, 'real.pdf')) if i == 1 or i % args.plot_freq == 0: plot_x = netG(torch.randn(args.batch_size, args.input_dim, device=device)).data plot_x = plot_x / 2. + 0.5 save_image(plot_x, os.path.join(dir_name, 'fake_images-{}.pdf'.format(i + 1))) utils.plot_losses(g_losses, d_losses, grad_normG, grad_normD, dir_name) if i % args.save_freq == 0: torch.save(netG.state_dict(), save_path + 'G_epoch{}.pth'.format(i)) torch.save(netD.state_dict(), save_path + 'D_epoch{}.pth'.format(i)) if args.ema_trick: torch.save(ema_netG_9999.state_dict(), save_path + 'emaG0.9999_epoch{}.pth'.format(i))
def save_images(self, images, shape, filename): # Additional dimensions are only passed to the shape instance when > 1 dimensions = 1 if len(shape) == 3 else shape[3] img_view = images.view(images.size(0), dimensions, shape[1], shape[2]) save_image(denorm(img_view.data), filename)
d_B_optimizer.zero_grad() loss_B_real = criterion_gan(real_B, valid(real_B)) loss_B_fake = criterion_gan(fake_B.detach(), fake(fake_B)) loss_discriminator_B = (loss_B_real + loss_B_fake) / 2 loss_discriminator_B = Variable(loss_discriminator_B, requires_grad=True) loss_discriminator_B.backward() d_B_optimizer.step() loss_d = (loss_discriminator_A + loss_discriminator_B) / 2 loss = loss_g + loss_d if i % 50 == 0: print( 'Epoch [{}/{}], Step [{}/{}], Loss:{:.4f}, G-Loss: {:.4f} , G-R-Loss: {:.4f} , G-I-Loss: {:.4f} , G-Adv-Loss: {:.4f}, D-Loss: {:.4f}, D-A-Loss: {:.4f}, D-B-Loss: {:.4f} ' .format(iteration + 1, args.epoch, i, int(len(train_loader)), loss.item(), loss_g, loss_recon, loss_identity, loss_gen, loss_d, loss_discriminator_A, loss_discriminator_B)) if i % 20 == 0: plot_list = [] for image in [real_A[0], real_B[0], fake_A[0], fake_B[0]]: image = image.detach().numpy() plot_list.append(image) save_image(torch.tensor(plot_list), 'generated-images/' + str(iteration + 1) + '-' + str(i) + '.png', nrow=4)
def train(self): dis_criterion = nn.BCELoss() aux_criterion = nn.NLLLoss() # add class loss input = torch.FloatTensor(self.batch_size, 3, self.image_size, self.image_size) noise = torch.FloatTensor(self.batch_size, self.nz, 1, 1) fixed_noise = torch.FloatTensor(self.batch_size, self.nz, 1, 1).normal_(0, 1) dis_label = torch.FloatTensor(self.batch_size) aux_label = torch.LongTensor(self.batch_size) # add class label real_label = 1 fake_label = 0 if self.cuda: dis_criterion.cuda() aux_criterion.cuda() input, dis_label, aux_label = input.cuda(), dis_label.cuda( ), aux_label.cuda() # add class label noise, fixed_noise = noise.cuda(), fixed_noise.cuda() inputv = Variable(input) noisev = Variable(noise) fixed_noisev = Variable(fixed_noise) dis_labelv = Variable(dis_label) aux_labelv = Variable(aux_label) # setup optimizer optimizerD = optim.Adam(self.netD.parameters(), lr=self.lr, betas=(self.beta1, 0.999)) optimizerG = optim.Adam(self.netG.parameters(), lr=self.lr, betas=(self.beta1, 0.999)) for epoch in range(self.niter): for i, data in enumerate(self.data_loader, 0): ############################ # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ########################### for p in self.netD.parameters(): # reset requires_grad p.requires_grad = True # they are set to False below in netG update # train with real self.netD.zero_grad() real_cpu, c_label = data # add c_label batch_size = real_cpu.size(0) if self.cuda: real_cpu = real_cpu.cuda() c_label = c_label.cuda() # add c label inputv.data.resize_as_(real_cpu).copy_(real_cpu) dis_labelv.data.resize_(batch_size).fill_(real_label) aux_labelv.data.resize_(batch_size).copy_(c_label) dis_out, aux_out = self.netD(inputv) dis_errD_real = dis_criterion(dis_out, dis_labelv) aux_errD_real = aux_criterion(aux_out, aux_labelv) errD_real = dis_errD_real + aux_errD_real errD_real.backward() D_x = dis_out.data.mean() # train with fake noisev.data.resize_(batch_size, self.nz, 1, 1).normal_(0, 1) c_label = np.random.randint(0, self.nl, batch_size) noisev_ = np.random.normal(0, 1, (batch_size, self.nz)) class_onehot = np.zeros((batch_size, self.nl)) class_onehot[np.arange(batch_size), c_label] = 1 noisev_[np.arange(batch_size), :self.nl] = class_onehot[ np.arange(batch_size)] noisev_ = (torch.from_numpy(noisev_)) noisev.data.copy_(noisev_.view(batch_size, self.nz, 1, 1)) aux_labelv.data.resize_(batch_size).copy_( torch.from_numpy(c_label)) fake = self.netG(noisev) dis_labelv.data.fill_(fake_label) dis_out, aux_out = self.netD(fake.detach()) dis_errD_fake = dis_criterion(dis_out, dis_labelv) aux_errD_fake = aux_criterion(aux_out, aux_labelv) errD_fake = dis_errD_fake + aux_errD_fake errD_fake.backward() D_G_z1 = dis_out.data.mean() errD = errD_real + errD_fake optimizerD.step() ############################ # (2) Update G network: maximize log(D(G(z))) ########################### for p in self.netD.parameters(): p.requires_grad = False # to avoid computation self.netG.zero_grad() dis_labelv.data.fill_( real_label) # fake labels are real for generator cost dis_out, aux_out = self.netD(fake) dis_errG = dis_criterion(dis_out, dis_labelv) aux_errG = aux_criterion(aux_out, aux_labelv) errG = dis_errG + aux_errG errG.backward() D_G_z2 = dis_out.data.mean() optimizerG.step() print( '[%d/%d][%d/%d] Loss_D_real: %.4f Loss_D_fake: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f' % (epoch + 1, self.niter, i + 1, len( self.data_loader), errD_real.data[0], errD_fake.data[0], errG.data[0], D_x, D_G_z1, D_G_z2)) if i % 100 == 0: vutils.save_image(real_cpu, '%s/real_samples.png' % self.outf, normalize=True) fake = self.netG(fixed_noisev) vutils.save_image(fake.data, '%s/fake_samples_epoch_%03d.png' % (self.outf, epoch + 1), normalize=True) # do checkpointing torch.save(self.netG.state_dict(), '%s/netG_epoch_%03d.pth' % (self.outf, epoch + 1)) torch.save(self.netD.state_dict(), '%s/netD_epoch_%03d.pth' % (self.outf, epoch + 1))
# compute the average loss avg_loss_G.update(errG.item(), batch_size) avg_loss_D.update(errD.item(), batch_size) avg_loss_A.update(accuracy, batch_size) avg_loss_M.update(mi.item(), batch_size) print('[%d/%d][%d/%d] Loss_D: %.4f (%.4f) Loss_G: %.4f (%.4f) D(x): %.4f D(G(z)): %.4f / %.4f Acc: %.4f (%.4f) MI: %.4f (%.4f)' % (epoch, opt.niter, i, len(dataloader), errD.item(), avg_loss_D.avg, errG.item(), avg_loss_G.avg, D_x, D_G_z1, D_G_z2, accuracy, avg_loss_A.avg, mi.item(), avg_loss_M.avg)) if i % 100 == 0: vutils.save_image( utils.normalize(real_cpu), '%s/real_samples.png' % opt.outf) # print('Label for eval = {}'.format(eval_label)) fake = netG(eval_noise, eval_label) vutils.save_image( utils.normalize(fake.data), '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch) ) # update eval_label if opt.visualize_class_label >= 0 and opt.label_rotation: eval_label_const = (eval_label_const + 1) % num_classes eval_label.data.fill_(eval_label_const) # compute metrics is_mean, is_std, fid = get_metrics(sampler, num_inception_images=opt.num_inception_images, num_splits=10, prints=True, use_torch=False)
target = torch.full((b_size, ), real_label).to(device) output = netD(fake_img_tensor) D_G_z2 = output.mean().item() errG = criterion(output, target) errG.backward() optimizerG.step() # Save Losses for plotting later G_losses.append(errG.item()) D_losses.append(errD.item()) if i % 100 == 0: print( f'[{epoch+1}/{num_epochs}][{i}/{len(dataloader)}] Loss_D: {errD.item():.4f} Loss_G: {errG.item():.4f}\tD(x): {D_x:.4f}\tD(G(z)): {D_G_z1:.4f} / {D_G_z2:.4f}' ) vutils.save_image(real_img, '%s/real_img%03d.jpg' % ("./celeba/load_weight_progress2", epoch + 1), normalize=True) with torch.no_grad(): fake = netG(noise).detach() vutils.save_image(fake.data, '%s/fake_img_epoch_%03d.jpg' % ("./celeba/load_weight_progress", epoch + 1), normalize=True) with torch.no_grad(): fake = netG(fixed_noise).detach() vutils.save_image(fake.data, '%s/fake_img_version_epoch_%03d.jpg' % ("./celeba/load_weight_results2", epoch + 1), normalize=True)
def train(self): self.train_hist = {} self.train_hist['D_loss'] = [] self.train_hist['G_loss'] = [] self.train_hist['per_epoch_time'] = [] self.real_labels = torch.ones(self.batch_size, 1).to(device=self.device) self.fake_labels = torch.zeros(self.batch_size, 1).to(device=self.device) #Train print('Training Starts.') for epoch in range(self.loaded_epoch, self.epochs): epoch_start_time = time.time() print('Epoch [%4d/%4d]:' % ((epoch + 1), self.epochs)) for iter, (x, _) in enumerate(self.data_loader): #_ means discard something; y is discarded. if iter == self.data_loader.dataset.__len__() // self.batch_size: break x = x.view(-1, self.data_dim).to(device=self.device) #(-1, 28*28) flatten z = torch.randn((self.batch_size, self.input_dim)).to(device=self.device) # 1. Training the Discriminator # 1-1. Compute D loss outputs = self.D(x) D_real_loss = self.BCE_loss(outputs, self.real_labels) real_score = outputs outputs = self.D(self.G(z)) D_fake_loss = self.BCE_loss(outputs, self.fake_labels) D_loss = 0.5 * (D_real_loss + D_fake_loss) # Loss of Discriminator fake_score = outputs # 1-2. Backpropagate and Optimize self.train_hist['D_loss'].append(D_loss.item()) # tensor.item: returns the value of this tensor as a python data. self.reset_grad() D_loss.backward() # Calculate gradients. self.D_optimizer.step() # Performs a single optimization step. # 2. Training the Generator # 2.1 Compute G loss fake_images = self.G(z) outputs = self.D(fake_images) G_loss = self.BCE_loss(outputs, self.real_labels) # 2.2 Backpropagate and Optimize self.train_hist['G_loss'].append(G_loss.item()) self.reset_grad() G_loss.backward() self.G_optimizer.step() if ((iter + 1) % 200) == 0: print(" Iter [{}/{}]: D_loss: {:.4f}, G_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}" .format(iter + 1, self.data_loader.dataset.__len__() // self.batch_size, D_loss.item(), G_loss.item(), real_score.mean().item(), fake_score.mean().item())) self.train_hist['per_epoch_time'].append(time.time() - epoch_start_time) avg_epoch_time = int(np.mean(self.train_hist['per_epoch_time'])) print(" Avg. 1 epoch time: [%s] / Est. remaining time: [%s]" % (str(datetime.timedelta(seconds = avg_epoch_time)), str(datetime.timedelta(seconds = (self.epochs - epoch - 1)*avg_epoch_time)))) if (epoch + 1) % self.saving_epoch_interval == 0: self.save(epoch) # Save sampled images fake_images = (fake_images + 1) / 2 fake_images = fake_images.reshape(self.G(z).size(0), 1, 28, 28) save_image(fake_images, os.path.join(self.sample_dir, 'fake_epoch_{}.png'.format(epoch+1))) print("Training Finish.")
def save_data(batch, folder, scene_id, batch_idx, save_images, save_txt_depth, save_binary_depth, save_cam, save_sem_images, save_txt_semantics, save_binary_semantics): # Create subfolder using ID of the scene full_path = os.path.join(folder, scene_id) print("\nSaving under directory:", full_path) try: os.mkdir(full_path) except FileExistsError: pass # Each value in batch dict: List of size n (different views) # NOTE: n is default 2, focusing on image pairs for now. img_batch0, img_batch1 = batch["images"] if save_images else (None, None) depth_batch0, depth_batch1 = batch[ "depths"] if save_txt_depth or save_binary_depth else (None, None) cam_batch0, cam_batch1 = batch["cameras"] if save_cam else (None, None) semantic_batch0, semantic_batch1 = batch[ "semantics"] if save_sem_images or save_txt_semantics or save_binary_semantics else ( None, None) file_prefix = os.path.join(full_path, scene_id) cam_file_content = "{:<12} = {}';\n" sample_batch_size = batch["images"][0].shape[0] start_idx = batch_idx * sample_batch_size num_views = 2 exts = [] if save_images: exts.append("png") if save_txt_depth: exts.append("depth") if save_binary_depth: exts.append("depth.npy") if save_cam: exts.append("txt") if save_sem_images: exts.append("seg.png") if save_txt_semantics: exts.append("semantic") if save_binary_semantics: exts.append("semantic.npy") files_per_view = len(exts) files_per_sample = num_views * files_per_view exts = "{" + ",".join(exts) + "}" for sample_idx in range(sample_batch_size): curr_file_idx = str(start_idx + sample_idx) template = file_prefix + "_" + curr_file_idx + "_{pair_id}.{ext}" # Save RGB images (scene_idx_pairid.png) if save_images: rgb_img0, rgb_img1 = img_batch0[sample_idx].cpu( ), img_batch1[sample_idx].cpu() save_image(rgb_img0, template.format(pair_id=0, ext='png')) save_image(rgb_img1, template.format(pair_id=1, ext='png')) # Save depth information if save_txt_depth or save_binary_depth: depth0, depth1 = depth_batch0[sample_idx].squeeze(0).cpu().numpy(),\ depth_batch1[sample_idx].squeeze(0).cpu().numpy() # Save depth information as text file (scene_idx_pairid.depth) if save_txt_depth: np.savetxt(template.format(pair_id=0, ext='depth'), depth0.ravel(), fmt='%.5f', delimiter=' ', newline=' ') np.savetxt(template.format(pair_id=1, ext='depth'), depth1.ravel(), fmt='%.5f', delimiter=' ', newline=' ') # Save depth information as binary file (scene_idx_pairid.depth.npy) if save_binary_depth: np.save(template.format(pair_id=0, ext='depth.npy'), depth0) np.save(template.format(pair_id=1, ext='depth.npy'), depth1) # Save camera parameters (scene_idx_pairid.txt) # NOTE: According to SynSin implementation of get_camera_matrices (@camera_transformations.py): # P: World->Cam, Pinv: Cam->World if save_cam: P0, K0, Pinv0, Kinv0 = cam_batch0["P"][sample_idx],\ cam_batch0["K"][sample_idx],\ cam_batch0["Pinv"][sample_idx],\ cam_batch0["Kinv"][sample_idx] cam_pos, cam_up, cam_dir = split_RT(Pinv0.cpu().numpy()) info = cam_file_content.format("cam_pos", cam_pos) info += cam_file_content.format("cam_dir", cam_dir) info += cam_file_content.format("cam_up", cam_up) with open(template.format(pair_id=0, ext='txt'), 'w+') as f: f.write(info) P1, K1, Pinv1, Kinv1 = cam_batch1["P"][sample_idx],\ cam_batch1["K"][sample_idx],\ cam_batch1["Pinv"][sample_idx],\ cam_batch1["Kinv"][sample_idx] cam_pos, cam_up, cam_dir = split_RT(Pinv1.cpu().numpy()) info = cam_file_content.format("cam_pos", cam_pos) info += cam_file_content.format("cam_dir", cam_dir) info += cam_file_content.format("cam_up", cam_up) with open(template.format(pair_id=1, ext='txt'), 'w+') as f: f.write(info) # Save semantic information in the form of int IDs if save_sem_images or save_txt_semantics or save_binary_semantics: semantic0, semantic1 = semantic_batch0[sample_idx].squeeze(0).cpu().numpy(),\ semantic_batch1[sample_idx].squeeze(0).cpu().numpy() if save_sem_images: sem_img0 = get_semantic_image(semantic0) sem_img1 = get_semantic_image(semantic1) sem_img0.save(template.format(pair_id=0, ext='seg.png')) sem_img1.save(template.format(pair_id=1, ext='seg.png')) # Save semantic information as text file (scene_idx_pairid.semantic) if save_txt_semantics: np.savetxt(template.format(pair_id=0, ext='semantic'), semantic0.ravel(), fmt='%d', delimiter=' ', newline=' ') np.savetxt(template.format(pair_id=1, ext='semantic'), semantic1.ravel(), fmt='%d', delimiter=' ', newline=' ') # Save semantic information as binary file (scene_idx_pairid.semantic.npy) if save_binary_semantics: np.save(template.format(pair_id=0, ext='semantic.npy'), semantic0) np.save(template.format(pair_id=1, ext='semantic.npy'), semantic1) print("Files created: {}_{}_{{0,1}}.{}".format(scene_id, curr_file_idx, exts)) print("Saving completed. Number of files created under {}: {}\n".format( full_path, files_per_sample * sample_batch_size))
def main(opts): # Create the data loader dataset = DatasetFromFolder(root, 'train') loader = DataLoader(dataset, batch_size=opts.batch_size, shuffle=True) # Create the model start_epoch = 0 G = StyleGenerator() D = StyleDiscriminator() # Load the pre-trained weight if os.path.exists(opts.checkpoint): print("Load checkpoint") state = torch.load(opts.checkpoint) G.load_state_dict(state['G']) D.load_state_dict(state['D']) start_epoch = state['start_epoch'] else: print( "Pre-trained weight cannot load successfully, train from scratch!") # Multi-GPU support if torch.cuda.device_count() > 1: print("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs") G = nn.DataParallel(G) D = nn.DataParallel(D) G.to(opts.device) D.to(opts.device) # Create the criterion, optimizer and scheduler optim_D = optim.Adam(D.parameters(), lr=0.00001, betas=(0.5, 0.999)) optim_G = optim.Adam(G.parameters(), lr=0.00001, betas=(0.5, 0.999)) scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99) scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99) # Train fix_z = torch.randn([opts.batch_size, 512]).to(opts.device) softplus = nn.Softplus() Loss_D_list = [0.0] Loss_G_list = [0.0] for ep in range(start_epoch, opts.epoch): bar = tqdm(loader) loss_D_list = [] loss_G_list = [] for i, (real_img, ) in enumerate(bar): # ======================================================================================================= # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) # ======================================================================================================= # Compute adversarial loss toward discriminator D.zero_grad() real_img = real_img.to(opts.device) real_logit = D(real_img) fake_img = G(torch.randn([real_img.size(0), 512]).to(opts.device)) fake_logit = D(fake_img.detach()) d_loss = softplus(fake_logit).mean() d_loss = d_loss + softplus(-real_logit).mean() if opts.r1_gamma != 0.0: r1_penalty = R1Penalty(real_img.detach(), D) d_loss = d_loss + r1_penalty * (opts.r1_gamma * 0.5) if opts.r2_gamma != 0.0: r2_penalty = R2Penalty(fake_img.detach(), D) d_loss = d_loss + r2_penalty * (opts.r2_gamma * 0.5) loss_D_list.append(d_loss.item()) # Update discriminator d_loss.backward() optim_D.step() # ======================================================================================================= # (2) Update G network: maximize log(D(G(z))) # ======================================================================================================= if i % CRITIC_ITER == 0: G.zero_grad() fake_logit = D(fake_img) g_loss = softplus(-fake_logit).mean() loss_G_list.append(g_loss.item()) # Update generator g_loss.backward() optim_G.step() # Output training stats bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format( ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1])) # Save the result Loss_G_list.append(np.mean(loss_G_list)) Loss_D_list.append(np.mean(loss_D_list)) # Check how the generator is doing by saving G's output on fixed_noise with torch.no_grad(): fake_img = G(fix_z).detach().cpu() save_image(fake_img, os.path.join(opts.det, 'images', str(ep) + '.png'), nrow=4, normalize=True) # Save model state = { 'G': G.state_dict(), 'D': D.state_dict(), 'Loss_G': Loss_G_list, 'Loss_D': Loss_D_list, 'start_epoch': ep, } torch.save(state, os.path.join(opts.det, 'models', 'latest.pth')) scheduler_D.step() scheduler_G.step() # Plot the total loss curve Loss_D_list = Loss_D_list[1:] Loss_G_list = Loss_G_list[1:] plotLossCurve(opts, Loss_D_list, Loss_G_list)
recon_loss.data.item())) if (epoch + 1) % 1 == 0: batch_size = 104 test_iter = iter(test_loader) test_data = next(test_iter) z_real = encoder(Variable(test_data[0]).cuda()) reconst = decoder(torch.randn_like(z_real)).cpu().view( batch_size, 1, 28, 28) if not os.path.isdir('./data/reconst_images'): os.makedirs('data/reconst_images') save_image( test_data[0].view(batch_size, 1, 28, 28), './data/reconst_images/wae_gan_clamp_input_%d.png' % (epoch + 1)) save_image( reconst.data, './data/reconst_images/wae_gan_clamp_images_%d.png' % (epoch + 1)) filename = "wae_gan_mnist_clamp" torch.save(encoder.state_dict(), filename + "_enc") torch.save(decoder.state_dict(), filename + "_dec") def load_model(filename): encoder = Encoder() decoder = Decoder() encoder.load_state_dict(torch.load(filename + "_enc"))
def save_pic(x_rec): resultsample = x_rec * 0.5 + 0.5 resultsample = resultsample.cpu() save_image(resultsample, 'sample_%i_lr.png' % i, nrow=16)
input = Variable(input) if args.use_cuda: input = input.cuda() optimizer.zero_grad() nll, kld = model.loss(input, criterion, eval=False) loss = nll + kld loss.backward() optimizer.step() if iteration % 10 == 0: print('epoch {}, iteration {}, loss {}'.format( epoch, iteration, loss.cpu().data.numpy()[0])) print(F.softmax(model.p_c_logits, dim=0).cpu().data.numpy()) sampling = model.sample(args.use_cuda) vutils.save_image( sampling, 'sampling/vae_{}.png'.format(epoch * len(dataloader) + iteration)) samplings = [f for f in listdir('sampling')] samplings = [ imageio.imread('sampling/' + path) for path in samplings for _ in range(2) ] imageio.mimsave('sampling/movie.gif', samplings)
avg_rec_loss = 0 avg_kld = 0 for i, (img1, text1, len_text1, img2, text2, len_text2) in tqdm(enumerate(train_loader)): img1, text_feat1, text_feat1_mismatch, img2, text_feat2, text_feat2_mismatch = \ preprecessing(img1, text1, len_text1, img2, text2, len_text2, text_encoder) ########################################################## #### one for discrimator #### img1_D = img1 img2_D = img2 ########################################################## if epoch % 20 == 0 and i == 0: save_image(img1.data, './result/origin_epoch_%d.jpg' % (epoch + 1)) save_image(img2.data, './result/target_epoch_%d.jpg' % (epoch + 1)) ONES = Variable(torch.ones(img1.size(0))) # size(0) is the number of the minibatch samples ZEROS = Variable(torch.zeros(img1.size(0))) if config.is_cuda: ONES, ZEROS = ONES.cuda(), ZEROS.cuda() ############################################## ### UPDATE DISCRIMINATOR ### ############################################## # ------------- img1, text1 pairs ------------ D.zero_grad() # real image with matching text real_logit = D(img1_D, text_feat1)
size = 256 input_A = torch.ones([1, 3, size, size], dtype=torch.float).to(device) input_B = torch.ones([1, 3, size, size], dtype=torch.float).to(device) transforms_ = [ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ] data_root = "/home/kuan/workspace/muke/pytorch/09/datasets/apple2orange" dataloader = DataLoader(ImageDataset(data_root, transforms_), batch_size=1, shuffle=False, num_workers=8) if not os.path.exists("outputs/A"): os.makedirs("outputs/A") if not os.path.exists("outputs/B"): os.makedirs("outputs/B") for i, batch in enumerate(dataloader): real_A = torch.tensor(input_A.copy_(batch['A']), dtype=torch.float).to(device) real_B = torch.tensor(input_B.copy_(batch['B']), dtype=torch.float).to(device) fake_B = 0.5 * (netG_A2B(real_A).data + 1.0) fake_A = 0.5 * (netG_B2A(real_B).data + 1.0) save_image(fake_A, "outputs/A/{}.png".format(i)) save_image(fake_B, "outputs/B/{}.png".format(i)) print(i)
], 0).cuda() hint = torch.cat((real_vim * mask, mask), 1) if flag: # fix samples writer.add_image( 'target imgs', vutils.make_grid(real_cim.mul(0.5).add(0.5), nrow=16)) writer.add_image( 'sketch imgs', vutils.make_grid(real_sim.mul(0.5).add(0.5), nrow=16)) writer.add_image( 'hint', vutils.make_grid((real_vim * mask).mul(0.5).add(0.5), nrow=16)) vutils.save_image( real_cim.mul(0.5).add(0.5), '%s/color_samples' % opt.outf + '.png') vutils.save_image( real_sim.mul(0.5).add(0.5), '%s/blur_samples' % opt.outf + '.png') fixed_sketch.resize_as_(real_sim).copy_(real_sim) fixed_hint.resize_as_(hint).copy_(hint) flag -= 1 fake = netG(Variable(real_sim), Variable(hint)) if gen_iterations < opt.baseGeni: contentLoss = criterion_L2( netF((fake.mul(0.5) - Variable(saber)) / Variable(diver)), netF(Variable((real_cim.mul(0.5) - saber) / diver)))
#====================================================================== for step in tqdm(range(args.n_samplings), desc="Samplings"): model_G.eval() model_D.eval() # seed 値の変更 np.random.seed(seed) torch.manual_seed(seed) # 入力ノイズ z の再生成 input_noize_z = torch.randn(size=(args.batch_size, args.n_input_noize_z, 1, 1)).to(device) #==================================================== # 生成器 G の 推論処理 #==================================================== with torch.no_grad(): # G(z) : 生成器から出力される偽物画像 G_z = model_G(input_noize_z) #--------------------- # 出力画像の生成&保存 #--------------------- save_image(tensor=G_z, filename=os.path.join(args.results_dir, args.exper_name) + "/fake_image_seed{}_batchAll.png".format(seed)) seed += 1 n_print -= 1 print("Finished Test Loop.")
#save_image(x, './mlp_img/k_values.eps', nrow=10) # Decode k values in order node_cluster = np.zeros(128).astype(np.float32) for idx in range(10): count = 0 for i in range(K): if c_label[i] == idx and count < 5: node_cluster = np.c_[node_cluster, c_cluster[i, :]] count += 1 node_cluster = np.transpose(node_cluster) node_cluster = np.delete(node_cluster, (0), axis=0) x = Variable(torch.from_numpy(node_cluster)).cuda() x = to_img(model[1](x).cpu().data) save_image(x, './mlp_img/k_values_order.eps', nrow=10) #------------------------------------------------- # Find images that mapped to cluster point #------------------------------------------------- # Find images mapped to each cluster for j in range(K): count = 0 cluster_i = np.zeros(128).astype(np.float32) for i in range(dataset_size): if cl[i] == j and count < 8: cluster_i = np.c_[cluster_i, Feat[i]] count += 1 cluster_i = np.transpose(cluster_i) cluster_i = np.delete(cluster_i, (0), axis=0) x_cluster_i = Variable(torch.from_numpy(cluster_i)).cuda()
for batch_idx, data in enumerate(dataloader): img, _ = data num_img = img.size(0) img = img.view(num_img, -1) img = Variable(img) if torch.cuda.is_available(): img = img.cuda() # forward pass recon_out, mu, logvar = model(img) loss = loss_function(recon_out, img, mu, logvar) # backward optimizer.zero_grad() loss.backward() optimizer.step() train_loss += loss.data if batch_idx % 100 == 0: print('Epoch:{} [{}/{} ({:.0f}%)]\tloss:{:.6f}'.format( epoch + 1, batch_idx * len(img), len(dataloader.dataset), 100.0 * batch_idx / len(dataloader), loss.data[0] / len(img))) print('Epoch {} Finished, Average loss:{:.4f}'.format( epoch + 1, train_loss / len(dataloader.dataset))) if epoch % 10 == 0: pic = to_img(recon_out.cpu().data) save_image(pic, './vae_img/images_{}.png'.format(epoch)) torch.save(model.state_dict(), './variational_autoencoder.pth')
g = g.to(device) z = np.random.RandomState(0).randn(n_sample, 512).astype("float32") with torch.no_grad(): img_pt, _ = g( [torch.from_numpy(z).to(device)], truncation=0.5, truncation_latent=latent_avg.to(device), randomize_noise=False, ) Gs_kwargs = dnnlib.EasyDict() Gs_kwargs.randomize_noise = False img_tf = g_ema.run(z, None, **Gs_kwargs) img_tf = torch.from_numpy(img_tf).to(device) img_diff = ((img_pt + 1) / 2).clamp(0.0, 1.0) - ( (img_tf.to(device) + 1) / 2).clamp(0.0, 1.0) img_concat = torch.cat((img_tf, img_pt, img_diff), dim=0) print(img_diff.abs().max()) utils.save_image(img_concat, name + ".png", nrow=n_sample, normalize=True, range=(-1, 1))
def sample(cfg, logger): torch.cuda.set_device(0) model = SoftIntroVAEModelTL( startf=cfg.MODEL.START_CHANNEL_COUNT, layer_count=cfg.MODEL.LAYER_COUNT, maxf=cfg.MODEL.MAX_CHANNEL_COUNT, latent_size=cfg.MODEL.LATENT_SPACE_SIZE, dlatent_avg_beta=cfg.MODEL.DLATENT_AVG_BETA, style_mixing_prob=cfg.MODEL.STYLE_MIXING_PROB, mapping_layers=cfg.MODEL.MAPPING_LAYERS, channels=cfg.MODEL.CHANNELS, generator=cfg.MODEL.GENERATOR, encoder=cfg.MODEL.ENCODER, beta_kl=cfg.MODEL.BETA_KL, beta_rec=cfg.MODEL.BETA_REC, beta_neg=cfg.MODEL.BETA_NEG[cfg.MODEL.LAYER_COUNT - 1], scale=cfg.MODEL.SCALE) model.cuda(0) model.eval() model.requires_grad_(False) decoder = model.decoder encoder = model.encoder mapping_tl = model.mapping_tl mapping_fl = model.mapping_fl dlatent_avg = model.dlatent_avg logger.info("Trainable parameters decoder:") print(count_parameters(decoder)) logger.info("Trainable parameters encoder:") print(count_parameters(encoder)) arguments = dict() arguments["iteration"] = 0 model_dict = { 'discriminator_s': encoder, 'generator_s': decoder, 'mapping_tl_s': mapping_tl, 'mapping_fl_s': mapping_fl, 'dlatent_avg': dlatent_avg } checkpointer = Checkpointer(cfg, model_dict, {}, logger=logger, save=False) extra_checkpoint_data = checkpointer.load() model.eval() layer_count = cfg.MODEL.LAYER_COUNT def encode(x): z, mu, _ = model.encode(x, layer_count - 1, 1) styles = model.mapping_fl(mu) return styles def decode(x): return model.decoder(x, layer_count - 1, 1, noise=True) path = cfg.DATASET.SAMPLES_PATH im_size = 2**(cfg.MODEL.LAYER_COUNT + 1) paths = list(os.listdir(path)) paths = sorted(paths) random.seed(5) random.shuffle(paths) def move_to(list, item, new_index): list.remove(item) list.insert(new_index, item) def make(paths): src = [] for filename in paths: img = np.asarray(Image.open(path + '/' + filename)) if img.shape[2] == 4: img = img[:, :, :3] im = img.transpose((2, 0, 1)) x = torch.tensor(np.asarray(im, dtype=np.float32), requires_grad=True).cuda() / 127.5 - 1. if x.shape[0] == 4: x = x[:3] factor = x.shape[2] // im_size if factor != 1: x = torch.nn.functional.avg_pool2d(x[None, ...], factor, factor)[0] assert x.shape[2] == im_size src.append(x) with torch.no_grad(): reconstructions = [] for s in src: latents = encode(s[None, ...]) reconstructions.append(decode(latents).cpu().detach().numpy()) return src, reconstructions def chunker_list(seq, size): return list((seq[i::size] for i in range(size))) final = chunker_list(paths, 4) path0, path1, path2, path3 = final path0.reverse() path1.reverse() path2.reverse() path3.reverse() src0, rec0 = make(path0) src1, rec1 = make(path1) src2, rec2 = make(path2) src3, rec3 = make(path3) initial_resolution = im_size lods_down = 1 padding_step = 4 width = 0 height = 0 current_padding = 0 final_resolution = initial_resolution for _ in range(lods_down): final_resolution /= 2 for i in range(lods_down + 1): width += current_padding * 2**(lods_down - i) height += current_padding * 2**(lods_down - i) current_padding += padding_step width += 2**(lods_down + 1) * final_resolution height += (lods_down + 1) * initial_resolution width = int(width) height = int(height) def make_part(current_padding, src, rec): canvas = np.ones([3, height + 20, width + 10]) padd = 0 initial_padding = current_padding height_padding = 0 for i in range(lods_down + 1): for x in range(2**i): for y in range(2**i): try: ims = src.pop() imr = rec.pop()[0] ims = ims.cpu().detach().numpy() imr = imr res = int(initial_resolution / 2**i) ims = resize(ims, (3, initial_resolution / 2**i, initial_resolution / 2**i)) imr = resize(imr, (3, initial_resolution / 2**i, initial_resolution / 2**i)) place( canvas, ims, current_padding + x * (2 * res + current_padding), i * initial_resolution + height_padding + y * (res + current_padding)) place( canvas, imr, current_padding + res + x * (2 * res + current_padding), i * initial_resolution + height_padding + y * (res + current_padding)) except IndexError: return canvas height_padding += initial_padding * 2 current_padding -= padding_step padd += padding_step return canvas canvas = [ make_part(current_padding, src0, rec0), make_part(current_padding, src1, rec1), make_part(current_padding, src2, rec2), make_part(current_padding, src3, rec3) ] canvas = np.concatenate(canvas, axis=2) path = './make_figures/output' os.makedirs(path, exist_ok=True) os.makedirs(os.path.join(path, cfg.NAME), exist_ok=True) print('Saving image') save_path = './make_figures/output/%s/reconstructions_multiresolution.png' % cfg.NAME os.makedirs(os.path.dirname(save_path), exist_ok=True) save_image(torch.Tensor(canvas), save_path)
npImg1 = cv2.imread("HRS_99H8_4801_s144.bmp") img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0) img2 = torch.randn(img1.size(), requires_grad = True) print(img1.size()) print(img2.size()) # device = torch.device('cuda') # if torch.cuda.is_available(): # img1 = img1.to(device) # img2 = img2.to(device) criterion = ML1Loss(img1.size()[2], img1.size()[3]) optimizer = optim.Adam([img2], lr=0.01) count = 0 loss = 1000 while loss > 100: optimizer.zero_grad() loss = criterion(img2, img1, ["HRS_99H8_4801_s144.bmp"]) count += 1 # print("{0}: {1}".format(count,loss)) loss.backward() optimizer.step() utils.save_image(img2,"img2.png",normalize=True, range=(0,255)) # utils.save_image(img1,"img1.png",normalize=True, range=(0,255))
def train(self): self.create_model() losses = {'perceptual': [], 'content': []} optimizer = torch.optim.Adam(itertools.chain( self.attention1.parameters(), self.encoder.parameters(), self.decoder.parameters()), lr=self.config.learning_rate) criterion = torch.nn.L1Loss() criterion_p = torch.nn.MSELoss(reduction='mean') styles = iter(self.style_loader) self.encoder.train() self.decoder.train() self.reporter.writeInfo("Start to train the model") for e in range(1, self.config.epoch_size + 1): for i, (content, target) in enumerate(self.train_loader): try: style, target = next(styles) except: styles = iter(self.style_loader) style, target = next(styles) content = content.cuda() style = style.cuda() fea_c = self.encoder(content) fea_s = self.encoder(style) out_feature, attention_map = self.attention1(fea_c, fea_s) out_content = self.decoder(out_feature) c1, c2, c3, _ = self.D(content) h1, h2, h3, _ = self.D(out_content) s1, s2, s3, _ = self.D(style) loss_content = torch.norm(c3 - h3, p=2) loss_perceptual = 0 for t in range(3): loss_perceptual += criterion( gram_matrix(eval('s' + str(t + 1))), gram_matrix(eval('h' + str(t + 1)))) loss = loss_content * self.config.content_weight + loss_perceptual * self.config.style_weight optimizer.zero_grad() loss.backward() optimizer.step() if i % self.config.log_interval == 0: now = datetime.datetime.now() otherStyleTime = now.strftime("%Y-%m-%d %H:%M:%S") print(otherStyleTime) print('epoch: ', e, ' iter: ', i) print( 'attention scartters: ', torch.std(attention_map.argmax(-1).float(), 1).mean().cpu()) print(attention_map.shape) self.attention1.hard = True self.attention1.eval() self.decoder.eval() tosave, perc, cont = self.eval() save_image( denorm(tosave), self.image_dir + '/epoch_{}-iter_{}.png'.format(e, i)) print("image saved to " + self.image_dir + '/epoch_{}-iter_{}.png'.format(e, i)) print('content loss:', cont) print('perceptual loss:', perc) self.reporter.writeTrainLog( e, i, f''' attention scartters: {torch.std(attention_map.argmax(-1).float(), 1).mean().cpu()}\n content loss: {cont}\n perceptual loss: {perc} ''') losses['perceptual'].append(perc) losses['content'].append(cont) self.plot_loss_curve(losses) self.attention1.hard = False self.attention1.train() self.decoder.train() torch.save( { 'layer1': self.attention1.state_dict(), 'encoder': self.encoder.state_dict(), 'decoder': self.decoder.state_dict() }, f'{self.model_state_dir}/epoch_{e}-iter_{i}.pth') print('model saved.')