def main() -> None: tokenizer = Tokenizer(args.vocab_file) vocabulary_size = len(tokenizer) searcher = BeamSearch(tokenizer.eos_index, beam_size=args.search_width) model = VAE( num_embeddings=len(tokenizer), dim_embedding=args.dim_embedding, dim_hidden=args.dim_hidden, dim_latent=args.dim_latent, num_layers=args.num_layers, bidirectional=args.bidirectional, dropout=0., word_dropout=0., dropped_index=tokenizer.unk_index, ).to(device) model.load_state_dict(torch.load(args.checkpoint_file, map_location=device)) model.eval() sentence1 = input('Please input sentence1: ') sentence2 = input('Please input sentence2: ') s1 = [tokenizer.bos_index ] + tokenizer.encode(sentence1) + [tokenizer.eos_index] s2 = [tokenizer.bos_index ] + tokenizer.encode(sentence2) + [tokenizer.eos_index] z1, _ = model.encode( torch.tensor([s1]).to(device), torch.tensor([len(s1)]).to(device)) z2, _ = model.encode( torch.tensor([s2]).to(device), torch.tensor([len(s2)]).to(device)) print("\nGenerate intermediate sentences") print(" %s" % sentence1) for r in range(1, 10): z = (1 - 0.1 * r) * z1 + 0.1 * r * z2 hidden = model.fc_hidden(z) hidden = hidden.view(1, -1, model.dim_hidden).transpose(0, 1).contiguous() start_predictions = torch.zeros(1, device=device).fill_( tokenizer.bos_index).long() start_state = {'hidden': hidden.permute(1, 0, 2)} predictions, log_probabilities = searcher.search( start_predictions, start_state, model.step) tokens = predictions[0, 0] tokens = tokens[tokens != tokenizer.eos_index].tolist() print("[%d:%d] %s" % (10 - r, r, tokenizer.decode(tokens))) print(" %s" % sentence2)
def main() -> None: tokenizer = Tokenizer(args.vocab_file) vocabulary_size = len(tokenizer) dataset = SentenceDataset(args.input_file, tokenizer=tokenizer.encode) loader = DataLoader(dataset, args.batch_size, shuffle=False, collate_fn=dataset.collate_fn, drop_last=False) searcher = BeamSearch(tokenizer.eos_index, beam_size=args.search_width) model = VAE( num_embeddings=len(tokenizer), dim_embedding=args.dim_embedding, dim_hidden=args.dim_hidden, dim_latent=args.dim_latent, num_layers=args.num_layers, bidirectional=args.bidirectional, dropout=0., word_dropout=0., dropped_index=tokenizer.unk_index, ).to(device) model.load_state_dict(torch.load(args.checkpoint_file, map_location=device)) model.eval() print('Generating sentence...') all_hypotheses = [] with torch.no_grad(): for s in tqdm(loader): s = s.to(device) length = torch.sum(s != tokenizer.pad_index, dim=-1) bsz = s.shape[0] mean, logvar = model.encode(s, length) # z = model.reparameterize(mean, logvar) z = mean hidden = model.fc_hidden(z) hidden = hidden.view(bsz, -1, model.dim_hidden).transpose(0, 1).contiguous() start_predictions = torch.zeros(bsz, device=device).fill_( tokenizer.bos_index).long() start_state = {'hidden': hidden.permute(1, 0, 2)} predictions, log_probabilities = searcher.search( start_predictions, start_state, model.step) for preds in predictions: tokens = preds[0] tokens = tokens[tokens != tokenizer.eos_index].tolist() all_hypotheses.append(tokenizer.decode(tokens)) print('Done') with open(args.output_file, 'w') as f: f.write('\n'.join(all_hypotheses))
plt.savefig('./{}/loss.png'.format(log_dir)) # visualize latent space test_loader = torch.utils.data.DataLoader(datasets.MNIST( data_dir, train=False, download=True, transform=transforms.ToTensor()), batch_size=10000, shuffle=True) x, labels = iter(test_loader).next() x = x.view(10000, -1) if use_gpu: x = Variable(x).cuda() else: x = Variable(x) z = model.encode(x) mu, logvar = z if use_gpu: mu, logvar = mu.cpu().data.numpy(), logvar.cpu().data.numpy() else: mu, logvar = mu.data.numpy(), logvar.data.numpy() plt.figure(figsize=(10, 10)) plt.scatter(mu[:, 0], mu[:, 1], marker='.', c=labels.numpy(), cmap=plt.cm.jet) plt.colorbar() plt.grid() plt.savefig('./{}/latent_space.png'.format(log_dir)) np.save('loss_list.npy', np.array(loss_list)) np.save('test_loss_list.npy', np.array(test_loss_list)) torch.save(model.state_dict(), 'model_weights.pth')
x_ul = [] for i in range(args.batch_size_ul): xi = tf.pad(x_ul_raw[i,:,:,:], [[2,2],[2,2],[0,0]]) xi = tf.random_crop(xi, [32,32,3]) xi = tf.image.random_flip_left_right(xi) x_ul.append(xi) x_ul = tf.stack(x_ul, axis=0) else: x = x_raw x_ul = x_ul_raw vae = VAE(args.latent_dim) net = Net() out = net.classifier('net', x, keep_prob=args.keep_prob, is_training=True, update_batch_stats=True) out_ul = net.classifier('net', x_ul, keep_prob=args.keep_prob, is_training=True, update_batch_stats=False) mu, logvar = vae.encode(x_ul, False) z = vae.reparamenterize(mu, logvar, False) x_recon = vae.decode(z, False) r0 = tf.zeros_like(z, name='zero_holder') x_recon_r0 = vae.decode(z+r0, False) diff2 = 0.5 * tf.reduce_sum((x_recon - x_recon_r0)**2, axis=[1,2,3]) diffJaco = tf.gradients(diff2, r0)[0] def normalizevector(r): shape = tf.shape(r) r = tf.reshape(r, [shape[0],-1]) r /= (1e-12+tf.reduce_max(tf.abs(r), axis=1, keepdims=True)) r / tf.sqrt(tf.reduce_sum(r**2, axis=1, keepdims=True)+1e-6) return tf.reshape(r, shape) # power method
def run(): dataset = AugMNISTDataset() dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=args.batch_size, num_workers=8) vae = VAE(latent_dim=args.latent_dim).to(args.device) discrim = nn.GRU(784, args.discrim_hidden, 1).to(args.device) discrim.flatten_parameters() classifier = nn.Sequential(nn.ReLU(), nn.Linear(args.discrim_hidden, args.latent_dim)).to(args.device) summary(vae, (784, )) discrim_opt = torch.optim.Adam(chain(discrim.parameters(), classifier.parameters()), lr=args.discrim_lr) enc_opt = torch.optim.Adam(vae.e_params(), lr=args.lr) dec_opt = torch.optim.Adam(vae.d_params(), lr=args.lr) step = 0 for epoch in range(args.epochs): for idx, sample in enumerate(dataloader): image = sample['image'].to(args.device).view(-1, 784) mu, logvar = vae.encode(image) z = vae.reparameterize(mu, logvar) x_hat = vae.decode(z) mse, kld = loss_function(x_hat, image, mu, logvar) vae_loss = mse + args.beta * kld vae_loss.backward() enc_grad = torch.nn.utils.clip_grad_norm(vae.e_params(), 100) enc_opt.step() enc_opt.zero_grad() if args.reg_coef: samples, labels, z = generate_subset_samples( args.batch_size // 4, args.n_samples, args.n_sample_dims, args.latent_dim, vae, args.device) discrim_out = discrim(samples.detach())[1][-1] logits = classifier(discrim_out) discrim_loss = F.binary_cross_entropy_with_logits( logits, labels) discrim_loss.backward() discrim_grad = torch.nn.utils.clip_grad_norm( chain(discrim.parameters(), classifier.parameters()), 100) discrim_opt.step() discrim_opt.zero_grad() dec_pre_grad = torch.nn.utils.clip_grad_norm( vae.d_params(), 100) discrim_out = discrim(samples)[1][-1] logits = classifier(discrim_out) discrim_loss = args.reg_coef * F.binary_cross_entropy_with_logits( logits, labels) discrim_loss.backward() discrim_opt.zero_grad() dec_post_grad = torch.nn.utils.clip_grad_norm( vae.d_params(), 100) dec_opt.step() dec_opt.zero_grad() else: dec_post_grad = torch.nn.utils.clip_grad_norm( vae.d_params(), 100) dec_opt.step() dec_opt.zero_grad() if step % 500 == 0: for k in range(args.latent_dim): z = resample_kth_z(k, args.n_show, args.latent_dim, args.device) save_samples(vae, z, f'z{k}_vae_output.png') z = torch.randn(args.n_show, args.latent_dim, device=args.device) save_samples(vae, z, 'vae_output.png') original_grid = tv.utils.make_grid( sample['image'][:args.n_show], int(args.n_show**0.5)) tv.utils.save_image(original_grid, 'images.png') print(f'step := {step}') if args.reg_coef: discrim_acc = ((logits > 0) == labels).float().mean() print(f'discrim_loss := {discrim_loss.item()}') print(f'discrim_acc := {discrim_acc}') print(f'grad/discrim := {discrim_grad}') print(f'grad/dec_pre := {dec_pre_grad}') print(f'vae_loss := {vae_loss}') print(f'mse := {mse}') print(f'kld := {kld}') print(f'grad/enc := {enc_grad}') print(f'grad/dec_post := {dec_post_grad}') step += 1
} train_loader, test_loader = load_data("Pointillism", data_load_params) model = VAE() if torch.cuda.is_available(): model.cuda() print('Using GPU') checkpoint = torch.load('./finalrun/runs/checkpoints/checkpoint2.pth.tar', map_location='cpu') model.load_state_dict(checkpoint['state_dict']) for param in model.parameters(): param.requires_grad = False model.eval() images = iter(test_loader).next() latent_vectors = model.encode(images[0])[0].data.numpy() categories = images[1].data.numpy() plt.figure(1) plt.scatter(latent_vectors[:, 0], latent_vectors[:, 1], c=[color[i] for i in categories]) print("Made first figure") # mean = np.mean(latent_vectors, axis=0) # std = np.std(latent_vectors, axis=0) # print(mean.shape) # print(std.shape) # normalized_latent_vectors = (latent_vectors-mean)/std # print(normalized_latent_vectors.shape) # # latent_df = pd.DataFrame(latent_vectors)
def main(_): if FLAGS.exp_name is None: FLAGS.exp_name = "vae_{0}_{1}_{2}_{3}".format(FLAGS.gaps, FLAGS.arch, FLAGS.latent_dim, FLAGS.beta) image_shape = [ int(i) for i in FLAGS.image_shape.strip('()[]{}').split(',') ] dirs = init_directories(FLAGS.exp_name, FLAGS.output_dir) dirs['data'] = '../../data' if FLAGS.data_dir is None else FLAGS.data_dir dirs['codes'] = os.path.join(dirs['data'], 'codes/') create_directories(dirs, FLAGS.train, FLAGS.save_codes) z_dist = Gaussian(FLAGS.latent_dim) output_dim = reduce(mul, image_shape, 1) output_dist = Gaussian(output_dim) run_config = tf.ConfigProto(allow_soft_placement=True) run_config.gpu_options.allow_growth = True run_config.gpu_options.per_process_gpu_memory_fraction = 0.9 sess = tf.Session(config=run_config) vae = VAE( session=sess, output_dist=output_dist, z_dist=z_dist, arch=FLAGS.arch, batch_size=FLAGS.batch_size, image_shape=image_shape, exp_name=FLAGS.exp_name, dirs=dirs, beta=FLAGS.beta, gaps=FLAGS.gaps, vis_reconst=FLAGS.visualize_reconstruct, vis_disent=FLAGS.visualize_disentangle, n_disentangle_samples=FLAGS.n_disentangle_samples, ) if FLAGS.train: data_manager = TeapotsDataManager(dirs['data'], FLAGS.batch_size, image_shape, shuffle=True, gaps=FLAGS.gaps, file_ext=FLAGS.file_ext, train_fract=0.8, inf=True) vae.train_iter, vae.dev_iter, vae.test_iter = data_manager.get_iterators( ) n_iters_per_epoch = data_manager.n_train // data_manager.batch_size FLAGS.stats_interval = int(FLAGS.stats_interval * n_iters_per_epoch) FLAGS.ckpt_interval = int(FLAGS.ckpt_interval * n_iters_per_epoch) n_iters = int(FLAGS.epochs * n_iters_per_epoch) vae.train(n_iters, n_iters_per_epoch, FLAGS.stats_interval, FLAGS.ckpt_interval) if FLAGS.save_codes: b_size = 500 #large batch, forward prop only data_manager = TeapotsDataManager(dirs['data'], b_size, image_shape, shuffle=False, gaps=False, file_ext=FLAGS.file_ext, train_fract=1., inf=False) data_manager.set_divisor_batch_size() vae.train_iter, vae.dev_iter, vae.test_iter = data_manager.get_iterators( ) vae.session.run(tf.global_variables_initializer()) saved_step = vae.load() assert saved_step > 1, "A trained model is needed to encode the data!" codes = [] for batch_num, (img_batch, _) in enumerate(vae.train_iter): code = vae.encode(img_batch) #[batch_size, reg_latent_dim] codes.append(code) if batch_num < 5 or batch_num % 100 == 0: print(("Batch number {0}".format(batch_num))) codes = np.vstack(codes) filename = os.path.join(dirs['codes'], "codes_" + FLAGS.exp_name) np.save(filename, codes) print(("Codes saved to: {0}".format(filename)))
class AudioToBodyDynamics(object): """ Defines a wrapper class for training and evaluating a model. Inputs: args (argparse object): model settings generator (tuple DataLoader): a tuple of at least one DataLoader """ def __init__(self, args, generator, freestyle=False): # TODO super(AudioToBodyDynamics, self).__init__() self.device = args.device self.log_frequency = args.log_frequency self.is_freestyle_mode = freestyle self.generator = generator self.model_name = args.model_name self.ident = args.ident self.model_name = args.model_name input_dim, output_dim = generator[0].dataset.getDimsPerBatch() model_options = { 'seq_len': args.seq_len, 'device': args.device, 'dropout': args.dp, 'batch_size': args.batch_size, 'hidden_dim': args.hidden_size, 'input_dim': input_dim, 'output_dim': output_dim, 'trainable_init': args.trainable_init } if args.model_name == "AudioToJointsThree": from model import AudioToJointsThree self.model = AudioToJointsThree(model_options).cuda(args.device) elif args.model_name == 'AudioToJointsNonlinear': from model import AudioToJointsNonlinear self.model = AudioToJointsNonlinear(model_options).cuda( args.device) elif args.model_name == "AudioToJoints": from model import AudioToJoints self.model = AudioToJoints(model_options).cuda(args.device) elif args.model_name == 'JointsToJoints': from model import JointsToJoints self.model = JointsToJoints(model_options).cuda( args.device).double() elif args.model_name == 'LSTMToDense': from model import LSTMToDense self.model = LSTMToDense(model_options).cuda(args.device).double() elif args.model_name == 'AudioToJointsSeq2Seq': from model import AudioToJointsSeq2Seq self.model = AudioToJointsSeq2Seq(model_options).cuda( args.device).double() elif args.model_name == 'MDNRNN': from model import MDNRNN self.model = MDNRNN(model_options).cuda(args.device).double() elif args.model_name == 'VAE': from model import VAE self.model = VAE(model_options).cuda(args.device).double() # construct the model self.optim = optim.Adam(self.model.parameters(), lr=args.lr) # Load checkpoint model if self.is_freestyle_mode: path = f"{model_dir}{args.model_name}_{str(args.ident)}.pth" print(path) self.loadModelCheckpoint(path) # general loss function def buildLoss(self, predictions, targets): square_diff = (predictions - targets)**2 out = torch.sum(square_diff, -1, keepdim=True) return torch.mean(out) def mdn_loss(self, y, pi, mu, sigma): m = torch.distributions.Normal(loc=mu, scale=sigma) loss = torch.exp(m.log_prob(y)) loss = torch.sum(loss * pi, dim=2) loss = -torch.log(loss) return torch.mean(loss) # Loss function from https://github.com/pytorch/examples/blob/master/vae/main.py, # Appendix B of https://github.com/pytorch/examples/blob/master/vae/main.py def vae_loss(self, targets, recon_targets, mu, logvar): BCE = nn.functional.binary_cross_entropy(recon_targets, targets, reduction='sum') KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) return BCE + KLD def saveModel(self, state_info, path): torch.save(state_info, path) def loadModelCheckpoint(self, path): checkpoint = torch.load(path, map_location=self.device) self.model.load_state_dict(checkpoint['model_state_dict']) self.optim.load_state_dict(checkpoint['optim_state_dict']) def runNetwork(self, inputs, targets): """ Train on one given mfcc pose pair Args: inputs (array): [batch, seq_len, mfcc_features * 3] targets (array): [batch, seq_len, 19 * 2 poses] Returns: predictions, truth, loss """ def to_numpy(x): # import from gpu device to cpu, convert to numpy return x.cpu().data.numpy() inputs = Variable(torch.DoubleTensor(inputs.double()).to(self.device)) # reshape targets into (batch * seq_len, input features) targets = Variable(torch.DoubleTensor(targets).to(self.device)) if self.model_name == 'AudioToJointsSeq2Seq': predictions = self.model.forward(inputs, targets) elif self.model_name == 'VAE': predictions, mu, logvar = self.model.forward(inputs) else: predictions = self.model.forward(inputs) criterion = nn.L1Loss() if self.model_name == 'AudioToJointsSeq2Seq': loss = criterion(predictions.to(self.device), targets.to(self.device).float()) elif self.model_name == 'MDNRNN': # predictions = (pi, mu, sigma), (h, c) loss = self.mdn_loss(targets, predictions[0][0], predictions[0][1], predictions[0][2]) elif self.model_name == 'VAE': loss = self.vae_loss(targets, predictions, mu, logvar) else: loss = criterion(predictions, targets) return (to_numpy(predictions), to_numpy(targets)), loss def runEpoch(self): # given one epoch train_losses = [] #coeff_losses val_losses = [] predictions, targets = [], [] if not self.is_freestyle_mode: # train # for each data point for mfccs, poses in self.generator[0]: self.model.train() # pass train flag to model pred_targs, train_loss = self.runNetwork(mfccs, poses) self.optim.zero_grad() train_loss.backward() self.optim.step() train_loss = train_loss.data.tolist() train_losses.append(train_loss) # validation loss for mfccs, poses in self.generator[1]: self.model.eval() pred_targs, val_loss = self.runNetwork(mfccs, poses) val_loss = val_loss.data.tolist() val_losses.append(val_loss) pred = pred_targs[0].reshape( int(pred_targs[0].shape[0] * pred_targs[0].shape[1]), 19, 2) predictions.append(pred) targets.append(pred_targs[1]) # test or predict / play w/ model if self.is_freestyle_mode: for mfccs, poses in self.generator[0]: self.model.eval() # mfccs = mfccs.float() pred_targs, val_loss = self.runNetwork(mfccs, poses) val_loss = val_loss.data.tolist() val_losses.append(val_loss) pred = pred_targs[0].reshape( int(pred_targs[0].shape[0] * pred_targs[0].shape[1]), 19, 2) predictions.append(pred) targets.append(pred_targs[1]) return train_losses, val_losses, predictions, targets def trainModel(self, max_epochs, logfldr, model_dir): # TODO log.debug("Training model") epoch_losses = [] batch_losses = [] val_losses = [] i, best_loss, iters_without_improvement = 0, float('inf'), 0 best_train_loss, best_val_loss = float('inf'), float('inf') if logfldr: if logfldr[-1] != '/': logfldr += '/' filename = f'{logfldr}epoch_of_model_{str(self.ident)}.txt' state_info = { 'epoch': i, 'epoch_losses': epoch_losses, 'batch_losses': batch_losses, 'validation_losses': val_losses, 'model_state_dict': self.model.state_dict(), 'optim_state_dict': self.optim.state_dict(), } for i in range(max_epochs): if int(i / 10) == 0: if i == 0: with open(filename, 'w') as f: f.write(f"Epoch: {i} started\n") else: with open(filename, 'a+') as f: f.write(f"Epoch: {i} started\n") # save the model if model_dir: if model_dir[-1] != '/': model_dir += '/' path = f"{model_dir}{self.model_name}_{str(self.ident)}.pth" self.saveModel(state_info, path) # train_info, val_info, predictions, targets iter_train, iter_val, predictions, targets = self.runEpoch() iter_mean = np.mean(iter_train) iter_val_mean = np.mean(iter_val) # iter_val_mean = np.mean(iter_val[0]), np.mean(iter_val[1]) epoch_losses.append(iter_mean) batch_losses.extend(iter_train) val_losses.append(iter_val_mean) log.info("Epoch {} / {}".format(i, max_epochs)) log.info(f"Training Loss : {iter_mean}") log.info(f"Validation Loss : {iter_val_mean}") best_train_loss = iter_mean if iter_mean < best_train_loss else best_train_loss best_val_loss = iter_val_mean if iter_val_mean < best_val_loss else best_val_loss # Visualize VAE latent space if self.model_name == 'VAE': self.vae_plot() self.plotResults(logfldr, epoch_losses, batch_losses, val_losses) path = f"{model_dir}{self.model_name}_{str(self.ident)}.pth" self.saveModel(state_info, path) return best_train_loss, best_val_loss # plot random subset of poses in VAE latent space def vae_plot(self): z_list = torch.Tensor(1, 2) poses = [] for input, output in self.generator: for inp in input: poses.append(inp) mu, logvar = self.model.encode(input) z = self.model.reparameterize(mu, logvar) z2 = z[:, -1, :] z_list = torch.cat((z_list.double(), z2.double()), 0) indices = np.random.randint(low=1, high=z_list.shape[0], size=1000) coords = np.array([z_list[ind, :].detach().numpy() for ind in indices]) # # k-means clustering for coloring # kmeans = KMeans(n_clusters=5).fit(coords) # y_kmeans = kmeans.predict(coords) # plt.scatter(coords[:,0], coords[:,1], c=y_kmeans, cmap='viridis') # plt.show() # # # draw each mean pose # centers = kmeans.cluster_centers_ # recons = [self.model.decode(torch.from_numpy(center)).detach().numpy().reshape(19,2) for center in centers] # k-medoids clustering for coloring kmedoids = KMedoids(n_clusters=5).fit(coords) y_kmedoids = kmedoids.predict(coords) plt.scatter(coords[:, 0], coords[:, 1], c=y_kmedoids, cmap='viridis') plt.show() recons = [] for center in kmedoids.cluster_centers_: c = np.array(center) for i in range(len(coords)): if np.array_equal(c, coords[i]): recons.append(poses[indices[i] - 1].detach().numpy().reshape(19, 2)) self.draw_poses(np.array(recons)) # Takes in np array of poses that are each 19x2 arrays def draw_poses(self, poses): count = 0 shift_by = np.array([750, 800]) - poses[0][8] poses += shift_by for pose in poses: person_id = str(0) + ", " + str([0]) canvas = draw_pose_figure(person_id, pose) file_name = "images/" + f"{count:05}.jpg" cv2.imwrite(file_name, canvas) count += 1 def plotResults(self, logfldr, epoch_losses, batch_losses, val_losses): losses = [epoch_losses, batch_losses, val_losses] names = [["Epoch loss"], ["Batch loss"], ["Val loss"]] _, ax = plt.subplots(nrows=len(losses), ncols=1) for index, pair in enumerate(zip(losses, names)): data = [pair[0][j] for j in range(len(pair[0]))] ax[index].plot(data, label=pair[1]) ax[index].legend() if logfldr: if logfldr[-1] != '/': logfldr += '/' save_filename = os.path.join( logfldr, f"{self.model_name}_{str(self.ident)}_results.png") plt.savefig(save_filename) plt.close()
ntest = 1 for batch_idx, (gldimg, label) in enumerate(test_loader): iteration = min(args.batch_size, ntest) for i in xrange(iteration): test_originals.append(Variable(gldimg[i])) occluded, center = occludeimg_and_returncenter(gldimg[i]) test_imgs.append(Variable(occluded)) test_centers.append(center) lb = Variable( torch.LongTensor( [label[i]] ) ) test_labels.append(labelembed.index_select(0, lb)) break if args.model == "VAE": for i in xrange(iteration): img = test_imgs[i] mu, var = model.encode(img.view(1, 784)) center = test_centers[i] maxidx = 1000 minloss = 10000000 for j in range(maxidx): #z_mu = model.reparametrize(mu, var) z = Variable(torch.FloatTensor(1,20).normal_()) recon = model.decode(z).view(1, 28, -1) recon1 = occludeimg_with_center(recon.view(1, 28, -1), center) loss = reconstruction_function(recon1.view(-1), img.view(-1)) if loss < minloss: minloss = loss min_recon = recon compare(test_originals[i].data, img.data, min_recon.data)
_, _, log_var2 = beta_vae(img) log_vars = torch.t( torch.stack((torch.mean(log_var1, dim=0), torch.mean(log_var2, dim=0)))) bar(log_vars.numpy()) # save images periodically if epoch % 10 == 0: pic = out1.data.view(out1.size(0), 1, 28, 28) save_image(pic, './img/vae_' + str(epoch) + '_epochs.png') pic = out2.data.view(out2.size(0), 1, 28, 28) save_image(pic, './img/ß_vae_' + str(epoch) + '_epochs.png') print('DONE') print('Running t-SNE...', end='', flush=True) with torch.no_grad(): tsne = TSNE(n_components=2, random_state=0) for img, labels in visloader: img = img.view(img.size(0), -1) raw_coords = tsne.fit_transform(img) vae_coords = tsne.fit_transform(vae.encode(img).numpy()) beta_vae_coords = tsne.fit_transform(beta_vae.encode(img).numpy()) beta_scatter(raw_coords, vae_coords, beta_vae_coords, labels) break print('DONE')
def export_dict(self, src_dico, tgt_dico, emb_en, emb_it, seed): params = self.params # Export adversarial dictionaries optim_X_AE = VAE(params).cuda() optim_Y_AE = VAE(params).cuda() print('Loading pre-trained models...') optim_X_AE.load_state_dict( torch.load(self.tune_dir + '/best/seed_{}_best_X.t7'.format(seed))) optim_Y_AE.load_state_dict( torch.load(self.tune_dir + '/best/seed_{}_best_Y.t7'.format(seed))) X_Z = optim_X_AE.encode(Variable(emb_en)).data Y_Z = optim_Y_AE.encode(Variable(emb_it)).data mstart_time = timer() for method in [params.eval_method]: results = get_word_translation_accuracy(params.src_lang, src_dico[1], X_Z, params.tgt_lang, tgt_dico[1], emb_it, method=method, dico_eval='default') acc1 = results[0][1] for method in [params.eval_method]: results = get_word_translation_accuracy(params.tgt_lang, tgt_dico[1], Y_Z, params.src_lang, src_dico[1], emb_en, method=method, dico_eval='default') acc2 = results[0][1] # csls = 0 print('{} takes {:.2f}s'.format(method, timer() - mstart_time)) print('Method:{} score:{:.4f}-{:.4f}'.format(method, acc1, acc2)) print('Building dictionaries...') params.dico_build = "S2T&T2S" params.dico_method = "csls_knn_10" X_Z = X_Z / X_Z.norm(2, 1, keepdim=True).expand_as(X_Z) emb_it = emb_it / emb_it.norm(2, 1, keepdim=True).expand_as(emb_it) f_dico_induce = build_dictionary(X_Z, emb_it, params) f_dico_induce = f_dico_induce.cpu().numpy() Y_Z = Y_Z / Y_Z.norm(2, 1, keepdim=True).expand_as(Y_Z) emb_en = emb_en / emb_en.norm(2, 1, keepdim=True).expand_as(emb_en) b_dico_induce = build_dictionary(Y_Z, emb_en, params) b_dico_induce = b_dico_induce.cpu().numpy() f_dico_set = set([(a, b) for a, b in f_dico_induce]) b_dico_set = set([(b, a) for a, b in b_dico_induce]) intersect = list(f_dico_set & b_dico_set) union = list(f_dico_set | b_dico_set) with io.open( self.tune_dir + '/best/{}-{}.dict'.format(params.src_lang, params.tgt_lang), 'w', encoding='utf-8', newline='\n') as f: for item in f_dico_induce: f.write('{} {}\n'.format(src_dico[0][item[0]], tgt_dico[0][item[1]])) with io.open( self.tune_dir + '/best/{}-{}.dict'.format(params.tgt_lang, params.src_lang), 'w', encoding='utf-8', newline='\n') as f: for item in b_dico_induce: f.write('{} {}\n'.format(tgt_dico[0][item[0]], src_dico[0][item[1]])) with io.open(self.tune_dir + '/best/{}-{}.intersect'.format( params.src_lang, params.tgt_lang), 'w', encoding='utf-8', newline='\n') as f: for item in intersect: f.write('{} {}\n'.format(src_dico[0][item[0]], tgt_dico[0][item[1]])) with io.open(self.tune_dir + '/best/{}-{}.intersect'.format( params.tgt_lang, params.src_lang), 'w', encoding='utf-8', newline='\n') as f: for item in intersect: f.write('{} {}\n'.format(tgt_dico[0][item[1]], src_dico[0][item[0]])) with io.open( self.tune_dir + '/best/{}-{}.union'.format(params.src_lang, params.tgt_lang), 'w', encoding='utf-8', newline='\n') as f: for item in union: f.write('{} {}\n'.format(src_dico[0][item[0]], tgt_dico[0][item[1]])) with io.open( self.tune_dir + '/best/{}-{}.union'.format(params.tgt_lang, params.src_lang), 'w', encoding='utf-8', newline='\n') as f: for item in union: f.write('{} {}\n'.format(tgt_dico[0][item[1]], src_dico[0][item[0]]))
def main(): time_str = time.strftime("%Y%m%d-%H%M%S") print('time_str: ', time_str) exp_count = 0 if args.experiment == 'a|s': direc_name_ = '_'.join([args.env, args.experiment]) else: direc_name_ = '_'.join( [args.env, args.experiment, 'bp2VAE', str(args.bp2VAE)]) direc_name_exist = True while direc_name_exist: exp_count += 1 direc_name = '/'.join([direc_name_, str(exp_count)]) direc_name_exist = os.path.exists(direc_name) try: os.makedirs(direc_name) except OSError as e: if e.errno != errno.EEXIST: raise if args.tensorboard_dir is None: logger = Logger('/'.join([direc_name, time_str])) else: logger = Logger(args.tensorboard_dir) env = gym.make(args.env) if args.wrapper: if args.video_dir is None: args.video_dir = '/'.join([direc_name, 'videos']) env = gym.wrappers.Monitor(env, args.video_dir, force=True) print('observation_space: ', env.observation_space) print('action_space: ', env.action_space) env.seed(args.seed) torch.manual_seed(args.seed) if args.experiment == 'a|s': dim_x = env.observation_space.shape[0] elif args.experiment == 'a|z(s)' or args.experiment == 'a|z(s, s_next)' or \ args.experiment == 'a|z(a_prev, s, s_next)': dim_x = args.z_dim policy = ActorCritic(input_size=dim_x, hidden1_size=3 * dim_x, hidden2_size=6 * dim_x, action_size=env.action_space.n) if args.use_cuda: Tensor = torch.cuda.FloatTensor torch.cuda.manual_seed_all(args.seed) policy.cuda() else: Tensor = torch.FloatTensor policy_optimizer = optim.Adam(policy.parameters(), lr=args.policy_lr) if args.experiment != 'a|s': from util import ReplayBuffer, vae_loss_function dim_s = env.observation_space.shape[0] if args.experiment == 'a|z(s)' or args.experiment == 'a|z(s, s_next)': from model import VAE vae = VAE(input_size=dim_s, hidden1_size=3 * args.z_dim, hidden2_size=args.z_dim) elif args.experiment == 'a|z(a_prev, s, s_next)': from model import CVAE vae = CVAE(input_size=dim_s, class_size=1, hidden1_size=3 * args.z_dim, hidden2_size=args.z_dim) if args.use_cuda: vae.cuda() vae_optimizer = optim.Adam(vae.parameters(), lr=args.vae_lr) if args.experiment == 'a|z(s)': from util import Transition_S2S as Transition elif args.experiment == 'a|z(s, s_next)' or args.experiment == 'a|z(a_prev, s, s_next)': from util import Transition_S2SNext as Transition buffer = ReplayBuffer(args.buffer_capacity, Transition) update_vae = True if args.experiment == 'a|s': from util import Record_S elif args.experiment == 'a|z(s)': from util import Record_S2S elif args.experiment == 'a|z(s, s_next)' or args.experiment == 'a|z(a_prev, s, s_next)': from util import Record_S2SNext def train_actor_critic(n): saved_info = policy.saved_info R = 0 cum_returns_ = [] for r in policy.rewards[::-1]: R = r + args.gamma * R cum_returns_.insert(0, R) cum_returns = Tensor(cum_returns_) cum_returns = (cum_returns - cum_returns.mean()) \ / (cum_returns.std() + np.finfo(np.float32).eps) cum_returns = Variable(cum_returns, requires_grad=False).unsqueeze(1) batch_info = SavedInfo(*zip(*saved_info)) batch_log_prob = torch.cat(batch_info.log_prob) batch_value = torch.cat(batch_info.value) batch_adv = cum_returns - batch_value policy_loss = -torch.sum(batch_log_prob * batch_adv) value_loss = F.smooth_l1_loss(batch_value, cum_returns, size_average=False) policy_optimizer.zero_grad() total_loss = policy_loss + value_loss total_loss.backward() policy_optimizer.step() if args.use_cuda: logger.scalar_summary('value_loss', value_loss.data.cpu()[0], n) logger.scalar_summary('policy_loss', policy_loss.data.cpu()[0], n) all_value_loss.append(value_loss.data.cpu()[0]) all_policy_loss.append(policy_loss.data.cpu()[0]) else: logger.scalar_summary('value_loss', value_loss.data[0], n) logger.scalar_summary('policy_loss', policy_loss.data[0], n) all_value_loss.append(value_loss.data[0]) all_policy_loss.append(policy_loss.data[0]) del policy.rewards[:] del policy.saved_info[:] if args.experiment != 'a|s': def train_vae(n): train_times = (n // args.vae_update_frequency - 1) * args.vae_update_times for i in range(args.vae_update_times): train_times += 1 sample = buffer.sample(args.batch_size) batch = Transition(*zip(*sample)) state_batch = torch.cat(batch.state) if args.experiment == 'a|z(s)': recon_batch, mu, log_var = vae.forward(state_batch) mse_loss, kl_loss = vae_loss_function( recon_batch, state_batch, mu, log_var, logger, train_times, kl_discount=args.kl_weight, mode=args.experiment) elif args.experiment == 'a|z(s, s_next)' or args.experiment == 'a|z(a_prev, s, s_next)': next_state_batch = Variable(torch.cat(batch.next_state), requires_grad=False) predicted_batch, mu, log_var = vae.forward(state_batch) mse_loss, kl_loss = vae_loss_function( predicted_batch, next_state_batch, mu, log_var, logger, train_times, kl_discount=args.kl_weight, mode=args.experiment) vae_loss = mse_loss + kl_loss vae_optimizer.zero_grad() vae_loss.backward() vae_optimizer.step() logger.scalar_summary('vae_loss', vae_loss.data[0], train_times) all_vae_loss.append(vae_loss.data[0]) all_mse_loss.append(mse_loss.data[0]) all_kl_loss.append(kl_loss.data[0]) # To store cum_reward, value_loss and policy_loss from each episode all_cum_reward = [] all_last_hundred_average = [] all_value_loss = [] all_policy_loss = [] if args.experiment != 'a|s': # Store each vae_loss calculated all_vae_loss = [] all_mse_loss = [] all_kl_loss = [] for episode in count(1): done = False state_ = torch.Tensor([env.reset()]) cum_reward = 0 if args.experiment == 'a|z(a_prev, s, s_next)': action = random.randint(0, 2) state_, reward, done, info = env.step(action) cum_reward += reward state_ = torch.Tensor([np.append(state_, action)]) while not done: if args.experiment == 'a|s': state = Variable(state_, requires_grad=False) elif args.experiment == 'a|z(s)' or args.experiment == 'a|z(s, s_next)' \ or args.experiment == 'a|z(a_prev, s, s_next)': state_ = Variable(state_, requires_grad=False) mu, log_var = vae.encode(state_) if args.bp2VAE and update_vae: state = vae.reparametrize(mu, log_var) else: state = vae.reparametrize(mu, log_var).detach() action_ = policy.select_action(state) if args.use_cuda: action = action_.cpu()[0, 0] else: action = action_[0, 0] next_state_, reward, done, info = env.step(action) next_state_ = torch.Tensor([next_state_]) cum_reward += reward if args.render: env.render() policy.rewards.append(reward) if args.experiment == 'a|z(s)': buffer.push(state_) elif args.experiment == 'a|z(s, s_next)' or args.experiment == 'a|z(a_prev, s, s_next)': if not done: buffer.push(state_, next_state_) if args.experiment == 'a|z(a_prev, s, s_next)': next_state_ = torch.cat( [next_state_, torch.Tensor([action])], 1) state_ = next_state_ train_actor_critic(episode) last_hundred_average = sum(all_cum_reward[-100:]) / 100 logger.scalar_summary('cum_reward', cum_reward, episode) logger.scalar_summary('last_hundred_average', last_hundred_average, episode) all_cum_reward.append(cum_reward) all_last_hundred_average.append(last_hundred_average) if update_vae: if args.experiment != 'a|s' and episode % args.vae_update_frequency == 0: assert len(buffer) >= args.batch_size train_vae(episode) if len(all_vae_loss) > 1000: if abs( sum(all_vae_loss[-500:]) / 500 - sum(all_vae_loss[-1000:-500]) / 500) < args.vae_update_threshold: update_vae = False if episode % args.log_interval == 0: print( 'Episode {}\tLast cum return: {:5f}\t100-episodes average cum return: {:.2f}' .format(episode, cum_reward, last_hundred_average)) if episode > args.num_episodes: print("100-episodes average cum return is now {} and " "the last episode runs to {} time steps!".format( last_hundred_average, cum_reward)) env.close() torch.save(policy, '/'.join([direc_name, 'model'])) if args.experiment == 'a|s': record = Record_S( policy_loss=all_policy_loss, value_loss=all_value_loss, cum_reward=all_cum_reward, last_hundred_average=all_last_hundred_average) elif args.experiment == 'a|z(s)': record = Record_S2S( policy_loss=all_policy_loss, value_loss=all_value_loss, cum_reward=all_cum_reward, last_hundred_average=all_last_hundred_average, mse_recon_loss=all_mse_loss, kl_loss=all_kl_loss, vae_loss=all_vae_loss) elif args.experiment == 'a|z(s, s_next)' or args.experiment == 'a|z(a_prev, s, s_next)': record = Record_S2SNext( policy_loss=all_policy_loss, value_loss=all_value_loss, cum_reward=all_cum_reward, last_hundred_average=all_last_hundred_average, mse_pred_loss=all_mse_loss, kl_loss=all_kl_loss, vae_loss=all_vae_loss) pickle.dump(record, open('/'.join([direc_name, 'record']), 'wb')) break