def train(args): hp = { "batch_size": args.batch_size, "learning_rate": args.learning_rate, "dropout": args.dropout_prob, "lambda_l2_reg": args.lambda_, "nonlinearity": tf.nn.elu, "grad_clipping": args.clip, "initialization": args.init } data = Data(args.data_file) arch = [ data.sample_size ] + \ [ args.layers_width ] * args.num_layers + \ [ args.latent_dim ] # (and symmetrically back out again) v = VAE(arch, hp, log_dir=LOG_DIR) v.train(data, max_steps=args.num_steps, max_epochs=args.num_epochs, cross_validate=False, verbose=True, save=True, outdir=METAGRAPH_DIR, plots_outdir=PLOTS_DIR, plot_latent_over_time=False) print("Trained!") return v
def __init__(self, input_shape, log_dir, filters=32, kernel_size=2, latent_channels=1, beta=1.0): # initialise FreyVAE specific variables self.filters = filters self.kernel_size = kernel_size self.latent_channels = latent_channels # call parent constructor VAE.__init__(self, input_shape, log_dir, beta=beta)
def vae_visualization(file_name="vae", random_file=None): # Create folders. if not os.path.isdir(FIGURE_VAE_VISUALIZATION_DIR): os.makedirs(FIGURE_VAE_VISUALIZATION_DIR) if random_file == None: # Load random data. file_list = os.listdir(RAW_DATA_DIR) random_file = np.random.choice(file_list) random_file_name = os.path.splitext(random_file)[0] obs = np.load(RAW_DATA_DIR + random_file)["obs"] # Load models. vae = VAE() vae.build_model(is_training=False, is_assigning=False) with tf.Session(graph=vae.graph) as sess: # Load variables. saver = tf.train.Saver() saver.restore(sess, SAVE_VAE_DIR + file_name) # Compute the reconstruction. recons = sess.run(vae.output, feed_dict={vae.Input: obs / 255.0}) tf.contrib.keras.backend.clear_session() imageio.mimsave( FIGURE_VAE_VISUALIZATION_DIR + random_file_name + ".gif", [plot_obs_recons(obs[i], recons[i]) for i in range(MAX_FRAME)], fps=20)
def encoding_process(load_file_name, file_list): import tensorflow as tf file_name_list = [os.path.splitext(file)[0] for file in file_list] data_length = len(file_list) # Load models. vae = VAE() vae.build_model(is_training=False, is_assigning=False) with tf.Session(graph=vae.graph) as sess: # Load variables. saver = tf.train.Saver() saver.restore(sess, SAVE_VAE_DIR + file_name) for i in range(data_length): # Load raw data. data = np.load(RAW_DATA_DIR + file_list[i]) obs = data["obs"] action = data["action"] # Compute the mean and standard deviation rather than the encoding z. mu, sigma = sess.run([vae.mu, vae.sigma], feed_dict={vae.Input: obs / 255.0}) # Save file. np.savez_compressed(ENCODING_DATA_DIR + file_name_list[i], mu=mu, sigma=sigma, action=action)
def main(): config = get_config() logging.basicConfig( format='%(asctime)s | %(message)s', handlers=[ logging.FileHandler(os.path.join(config.log_root, config.log_name)), logging.StreamHandler()], level=logging.INFO) transformer = transforms.ToTensor() train_dataset = datasets.MNIST(config.data_root, train=True, download=True, transform=transformer) test_dataset = datasets.MNIST(config.data_root, train=False, download=True, transform=transformer) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=4, pin_memory=True) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=4, pin_memory=True) model = VAE() trainer = Trainer(model, train_loader, test_loader, Adam(model.parameters(), lr=0.0002, betas=(0.5, 0.999)), loss_function, 1) trainer.train(config.epochs, config.log_metrics_every)
def restore_model(restore_name): checkpoint_dir = os.path.join('../bsg/weights', restore_name) checkpoint_fns = os.listdir(checkpoint_dir) max_checkpoint_epoch, latest_checkpoint_idx = -1, -1 for cidx, checkpoint_fn in enumerate(checkpoint_fns): if 'best' in checkpoint_fn: latest_checkpoint_idx = cidx break checkpoint_epoch = int(checkpoint_fn.split('_')[-1].split('.')[0]) max_checkpoint_epoch = max(max_checkpoint_epoch, checkpoint_epoch) if checkpoint_epoch == max_checkpoint_epoch: latest_checkpoint_idx = cidx latest_checkpoint_fn = os.path.join(checkpoint_dir, checkpoint_fns[latest_checkpoint_idx]) print('Loading model from {}'.format(latest_checkpoint_fn)) if not torch.cuda.is_available(): checkpoint_state = torch.load( latest_checkpoint_fn, map_location=lambda storage, loc: storage) else: checkpoint_state = torch.load(latest_checkpoint_fn) vocab = checkpoint_state['vocab'] print('Previous checkpoint at epoch={}...'.format(max_checkpoint_epoch)) for k, v in checkpoint_state['losses'].items(): print('{}={}'.format(k, v)) args = argparse.ArgumentParser() for k, v in checkpoint_state['args'].items(): print('{}={}'.format(k, v)) setattr(args, k, v) vae_model = VAE(args, vocab.size()) vae_model.load_state_dict(checkpoint_state['model_state_dict']) optimizer_state = checkpoint_state['optimizer_state_dict'] return args, vae_model, vocab, optimizer_state
def __init__(self, args): self.log_path = args.log_path self.device = torch.device("cuda:0" if args.cuda else "cpu") self.img_size = args.img_size self.sample_num = args.sample_num self.transform = transforms.Compose([ transforms.ToPILImage(), transforms.Resize([64, 64], 1), transforms.ToTensor(), transforms.Normalize([.5], [.5]) ]) self.pil_transform = transforms.ToPILImage(mode="RGB") self.norm_scale = np.loadtxt(os.path.join(args.config_path, "norm_scale.txt"), dtype=np.float32, delimiter=",")[None] self.norm_min = np.loadtxt(os.path.join(args.config_path, "norm_min.txt"), dtype=np.float32, delimiter=",")[None] self.pb_list = torch.from_numpy( np.loadtxt(os.path.join(args.config_path, "pb_list.txt"), dtype=np.float32, delimiter=",")) self.kmeans = KMeans(n_clusters=2) self.kmeans.fit(self.pb_list) print("=" * 5, "Init LSTMPB", "=" * 5) self.rnn = LSTMPB(args, pb_unit=self.pb_list[5][None]) pt_file = load_model(args.model_load_path, "*/*LSTMPB*.pt") self.rnn.load_state_dict(torch.load(pt_file)) print("=" * 5, "Init VAE", "=" * 5) self.vae = VAE(img_size=args.img_size, z_dim=args.vae_z_dims) pt_file = load_model(args.model_load_path, "*/VAE*.pt") self.vae.load_state_dict(torch.load(pt_file)) self.vae.eval() print("=" * 5, "Init CVAE", "=" * 5) self.cvae = CVAE(img_size=args.img_size, z_dim=args.cvae_z_dims) pt_file = load_model(args.model_load_path, "*/*CVAE*.pt") self.cvae.load_state_dict(torch.load(pt_file)) self.cvae.eval() self.norm_mode = { "joint": [0, 1, 2, 3, 4], "visual": [5, 6, 7, 8, 9, 10, 11] } self.norm_mode[ "all"] = self.norm_mode["joint"] + self.norm_mode["visual"] self.global_step = 0 self.his_log = HistoryWindow(maxlen=args.window_size) #visualize current goal _, goal = self.vae.decoder(self.denorm(self.goal, "visual")) goal = ((goal[0] * .5 + .5) * 255).to(torch.int8) self.goal_img = self.pil_transform(goal)
def load_model(self, controller_weights=None): p = self.params self.action_utils = ActionUtils(p['env_name']) self.action_size = self.action_utils.action_size() self.vae = VAE() sys.stdout = open(os.devnull, 'w') self.vae.make_vae_shape(p['img_size'], p['img_size'], p['latent_size']) sys.stdout = sys.__stdout__ self.vae.load_model('../' + p['vae_hps']['weights_path']) # TODO: Make MDN just take in all of params. mdn_hps = p['mdn_hps'] mdn_hps['max_seq_len'] = p['max_seq_len'] mdn_hps['in_width'] = p['latent_size'] + self.action_size mdn_hps['out_width'] = p['latent_size'] mdn_hps['action_size'] = self.action_size mdn_hps['rnn_size'] = p['hidden_size'] mdn_hps['batch_size'] = 1 mdn_hps['max_seq_len'] = 1 mdn_hps['use_recurrent_dropout'] = 0 mdn_hps['training'] = 0 # self.mdn_rnn = MDNRNN(mdn_hps) # hps_inf = MDNRNN.set_hps_to_inference(hps) self.mdn_rnn = MDNRNN(mdn_hps) self.mdn_rnn.load('../' + p['mdn_hps']['weights_path']) self.controller = ControllerModel( [p['latent_size'] + p['hidden_size'], self.action_size]) if controller_weights: self.controller.load_weights(controller_weights)
def next_gen_process(id, model_info): latent_size = 32 vae_model = VAE(latent_size) checkpoint_dir = './vae_ckpt/' latest = tf.train.latest_checkpoint(checkpoint_dir) vae_model.load_weights(latest) s = np.random.randint(1e8) np.random.seed(s + id) # different seed for each process env = gym.make('CarRacing-v0', verbose=0) env = wrap_env(env, W, H, gray_scale=False) results = [] for info in model_info: m = Model() m.set_weights(info['weights']) m.copy_model(info['attr'], from_pickle=True) res = m.evaluate(env, vae_model) results.append(res) env.close() return results
def write_images(H, ema_params, viz_batch_original, viz_batch_processed, fname, logprint): rng = random.PRNGKey(H.seed_sample) ema_apply = partial( VAE(H).apply, {'params': jax_utils.unreplicate(ema_params)}) forward_get_latents = partial(ema_apply, method=VAE(H).forward_get_latents) forward_samples_set_latents = partial( ema_apply, method=VAE(H).forward_samples_set_latents) forward_uncond_samples = partial(ema_apply, method=VAE(H).forward_uncond_samples) zs = [s['z'] for s in forward_get_latents(viz_batch_processed, rng)] batches = [viz_batch_original.numpy()] mb = viz_batch_processed.shape[0] lv_points = np.floor( np.linspace(0, 1, H.num_variables_visualize + 2) * len(zs)).astype(int)[1:-1] for i in lv_points: batches.append(forward_samples_set_latents(mb, zs[:i], rng, t=0.1)) for t in [1.0, 0.9, 0.8, 0.7][:H.num_temperatures_visualize]: batches.append(forward_uncond_samples(mb, rng, t=t)) n_rows = len(batches) im = np.concatenate(batches, axis=0).reshape( (n_rows, mb, *viz_batch_processed.shape[1:])).transpose([0, 2, 1, 3, 4]).reshape([ n_rows * viz_batch_processed.shape[1], mb * viz_batch_processed.shape[2], 3 ]) logprint(f'printing samples to {fname}') Image.fromarray(im).save(fname)
def criterion(X, X_hat, mean, std): """ :param Tensor X: The original input data that was passed to the B-VAE. (N, input_shape[1], H, W) :param Tensor X_hat: The reconstructed data, the output of the B-VAE. (N, input_shape[1], H, W) :param Tensor mean: The output of the mean layer, computed with the output of the encoder. (N, z_dim) :param Tensor std: The output of the standard deviation layer, computed with the output of the encoder. (N, z_dim) :return: A dictionary containing the values of the losses computed. :rtype: dict This method computes the loss of the B-VAE using the formula: L(x, x_hat) = - E_{z ~ q_{phi}(z | x)}[log(p_{theta}(x|z))] + beta * D_{KL}[q_{phi}(z | x) || p_{theta}(x)] Intuitively, the expectation term is the Data Fidelity term, and the second term is a regularizer that makes sure the distribution of the encoder and the decoder stay close. """ # get the 2 losses data_fidelity_loss = VAE._data_fidelity_loss(X, X_hat) kl_divergence_loss = VAE._kl_divergence_loss(mean, std) # add them to compute the loss for each training example in the mini batch loss = -data_fidelity_loss + self.beta * kl_divergence_loss # place them all inside a dictionary and return it losses = {"data_fidelity": torch.mean(data_fidelity_loss), "kl-divergence": torch.mean(kl_divergence_loss), "beta_kl-divergence": self.beta * torch.mean(kl_divergence_loss), "loss": torch.mean(loss)} return losses
def train(self, generations=1500, evolutionary_leap_gens=10, disp_best=True): env = gym.make('CarRacing-v0', verbose=0) env = wrap_env(env, W, H, gray_scale=False) latent_size = 32 vae_model = VAE(latent_size) checkpoint_dir = './vae_ckpt/' latest = tf.train.latest_checkpoint(checkpoint_dir) vae_model.load_weights(latest) for g in range(1, generations + 1): print('Generation', g) st = time() if not g % evolutionary_leap_gens: print('Leap') best_model = self.evolutionary_leap() else: best_model = self.next_gen() print(best_model) if disp_best: best_model.evaluate(env, vae_model, disp=True, n=1) best_model.save_all(name='{generation:04d}'.format(generation=g) + 'result:' + str(best_model.mean_result)) print('T:', time() - st) env.close()
def train_model_art(model_dir, num_steps, batch_size=64, learning_rate=0.0005): """ Train a VAE on the art data :param model_dir: :param num_steps: :param batch_size: :param learning_rate: :return: """ with tf.Session() as sess: vae = VAE(sess=sess, model_dir=model_dir, batch_size=batch_size, learning_rate=learning_rate, height=128, width=128, cdim=1, n_z=128) # Load cifar dataset data = load_art_data( data_path="/home/kevin/deep_learning/cat-dataset/cats/CAT_00", batch_size=batch_size) # Training vae.train(data=data, num_epochs=num_steps)
def __init__(self, n_input, n_hidden, dim_z, n_output, gamma, binary=True, **kwargs): """initialize neural networks :param gamma: weight for total correlation term in loss function """ super(FactorVAE, self).__init__() self.dim_z = dim_z self.binary = binary self.gamma = gamma self.input_size = (n_input,) # VAE networks self.vae = VAE(n_input, n_hidden, dim_z, n_output, binary, **kwargs) # discriminator layers D_hidden_num = 3 D_hidden_dim = 1000 D_hidden_dims = [D_hidden_dim] * D_hidden_num D_act = nn.LeakyReLU D_act_args = {"negative_slope": 0.2, "inplace": False} D_output_dim = 2 self.discriminator = nns.create_mlp(self.dim_z, D_hidden_dims, act_layer=D_act, act_args=D_act_args, norm=True) self.discriminator = nn.Sequential( self.discriminator, nn.Linear(D_hidden_dim, D_output_dim))
def main(args): """main procedure""" # get configuration device = global_conf["device"] img_size = global_conf["image_size"] data_dir = global_conf["data_dir"] res_dir = global_conf["res_dir"] # prepare data train_loader, test_loader = prepare_data(args, dir_path=data_dir) # prepare model model = VAE(img_size[0] * img_size[1], args.n_hidden, args.dim_z, img_size[0] * img_size[1]) optimizer = optim.Adam(model.parameters(), lr=args.lr) # train and test losses = [] for epoch in range(1, args.epochs + 1): avg_loss = train(model, train_loader, epoch, optimizer, args, device, img_size) losses.append(avg_loss) test(model, test_loader, epoch, args, device, img_size, res_dir) with torch.no_grad(): sample = model.sample(64, device).cpu() save_image(sample.view(64, 1, img_size[0], img_size[1]), res_dir + '/sample_' + str(epoch) + '.png') # plot train losses plt.xlabel('Epoch') plt.ylabel('Loss') plt.plot(losses) plt.savefig(res_dir + '/loss.png')
def demo_embedding(): # get input data mnist_data = input_data.read_data_sets('../../dataset/mnist_data', one_hot=True) num_sample = mnist_data.train.num_examples batch_size = 100 network_architecture = \ dict(n_hidden_encoder_1=500, # 1st layer encoder neurons n_hidden_encoder_2=500, # 2nd layer encoder neurons n_hidden_decoder_1=500, # 1st layer decoder neurons n_hidden_decoder_2=500, # 2nd layer decoder neurons n_input=784, # MNIST data input (img shape: 28*28) n_z=2) # dimensionality of latent space # define model vae_model = VAE(network_architecture, batch_size=batch_size) # train the model train(model=vae_model, inputs=mnist_data.train, num_epoch=50, num_sample=num_sample, batch_size=batch_size) # embedding data into 2D space X_sample, Y_sample = mnist_data.test.next_batch(5000) z_mu = vae_model.embedding(X_sample) plt.figure(figsize=(8, 6)) plt.scatter(z_mu[:, 0], z_mu[:, 1], c=np.argmax(Y_sample, 1)) plt.colorbar() plt.grid() plt.show()
def objective(params): """ Objective function to be minimized: loss with respect to our hyperparameters. """ enc_kernel1 = int(params[0]) enc_kernel2 = int(params[1]) enc_kernel3 = int(params[2]) dec_kernel1 = int(params[3]) dec_kernel2 = int(params[4]) dec_kernel3 = int(params[5]) # Contact matrices are 21x21 input_dim = 441 encoder = Encoder(input_size=input_dim, latent_size=8, kernel1=enc_kernel1, kernel2=enc_kernel2, kernel3=enc_kernel3) decoder = Decoder(latent_dim=8, output_size=input_size, kernel1=dec_kernel1, kernel2=dec_kernel2, kernel3=dec_kernel3) vae = VAE(encoder, decoder) criterion = nn.MSELoss() use_cuda = args.use_cuda if use_cuda: encoder = encoder.cuda() deconder = decoder.cuda() vae = vae.cuda() criterion = criterion.cuda() optimizer = optim.Adam(vae.parameters(), lr=0.0001) epoch_loss = 0 total_loss = 0 for epoch in range(100): for i, data in enumerate(trainloader, 0): inputs = data['cont_matrix'] inputs = inputs.resize_(args.batch_size, 1, 21, 21) inputs = inputs.float() if use_cuda: inputs = inputs.cuda() inputs = Variable(inputs) optimizer.zero_grad() dec = vae(inputs) ll = latent_loss(vae.z_mean, vae.z_sigma) loss = criterion(dec, inputs) + ll loss.backward() optimizer.step() epoch_loss = loss.data[0] print(epoch, epoch_loss) total_loss += epoch_loss return total_loss
def main(): train_dataset = datasets.MNIST( root='C:/Users/user/Documents/InterestingAttempt/VAE/mnist_data/', train=True, transform=transforms.Compose([transforms.ToTensor()])) test_dataset = datasets.MNIST( root='C:/Users/user/Documents/InterestingAttempt/VAE/mnist_data/', train=False, transform=transforms.Compose([transforms.ToTensor()])) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=100, shuffle=True) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=1000, shuffle=False) criterion = nn.MSELoss() # criterion = nn.BCELoss() criterion2 = KLLoss() epoch_num = 30 lr = 1e-3 weight_decay = 1e-5 lamda = 0.01 latent_num = 2 mid_features = 256 outf = r'C:\Users\user\Documents\InterestingAttempt\VAE\logs\linear_{}_{}_{}_{}_{}'.format( latent_num, lr, lamda, weight_decay, epoch_num) if not os.path.exists(outf): os.makedirs(outf) model = VAE(28 * 28, mid_features, latent_num).cuda() optimizer = optim.Adam( model.parameters(), weight_decay=weight_decay, betas=(0.9, 0.999)) writer = SummaryWriter(outf) for epoch in range(epoch_num): current_lr = lr / 2**int(epoch / 40) for param_group in optimizer.param_groups: param_group['lr'] = current_lr train_epoch( model, optimizer, train_loader, criterion, epoch, writer=writer, criterion2=criterion2, lamda=lamda) test( model, test_loader, criterion, epoch, writer=writer, criterion2=criterion2) if (epoch + 1) % 10 == 0: torch.save(model.state_dict(), os.path.join(outf, 'model_{}.pth'.format(epoch))) writer.close() torch.save(model.state_dict(), os.path.join(outf, 'model.pth'))
def __init__(self, mdir): # Pretty much copy-paste from Rollout Generator initialization self.action_space = gym.spaces.Box(np.array([-1, 0, 0]), np.array([1, 1, 1])) self.observation_space = gym.spaces.Box(low=0, high=255, shape=(64, 64, 3), dtype=np.uint8) device = settings.device vae_input_size = (settings.reduced_image_channels, settings.reduced_image_width, settings.reduced_image_height) self.vae = VAE(input_size=vae_input_size, latent_dim=settings.vae_latent_dim).to(device) vae_savefile = mdir / 'vae.pt' self.vae.load_state_dict(torch.load(vae_savefile)) self.vae.eval() self.mdrnn = MixtureDensityLSTMCell( settings.vae_latent_dim, settings.action_space_size, settings.mdrnn_hidden_dim, settings.mdrnn_num_gaussians).to(device) mdrnn_savefile = mdir / 'mdrnn.pt' state = torch.load(mdrnn_savefile) new_state = {} for k, v in state.items(): new_k = k.rstrip('_l0') new_k = new_k.replace('lstm', 'lstm_cell') new_state[new_k] = v self.mdrnn.load_state_dict(new_state) self.mdrnn.eval() input_size = (settings.vae_latent_dim + settings.mdrnn_hidden_dim) output_size = 3 self.controller = Controller(input_size, output_size).to(device) controller_savefille = mdir / 'controller.pt' if Path(controller_savefille).exists(): self.controller.load_state_dict(torch.load(controller_savefille)) self.controller.eval() self.env = gym.make('CarRacing-v0') self.device = device self.time_limit = 1000 self.latent = torch.randn(1, self.vae.latent_dim).to(self.device) # Multiply by two because have current, and instinctual next-state # prediction self.hidden = [ torch.zeros(1, settings.mdrnn_hidden_dim).to(self.device) for _ in range(2) ] self.obs = None self.visual_obs = None self.monitor = None self.figure = None
def __init__(self, input_shape, log_dir, filters=32, kernel_size=2, pre_latent_size=64, latent_size=2): # initialise FreyVAE specific variables self.filters = filters self.kernel_size = kernel_size self.pre_latent_size = pre_latent_size self.latent_size = latent_size # call parent constructor VAE.__init__(self, input_shape, log_dir)
def create_vae(conf, vocab): # emb = torchtext.vocab.GloVe(conf.vector, conf.n_embed) # vae = VAE(conf, emb) vae = VAE(conf) vae.embedding.weight.data.copy_(vocab.vectors) vae = on_cuda(vae) trainer_vae = torch.optim.Adam(vae.parameters(), lr=conf.lr) return vae, trainer_vae
def train(self, disp_best=False, save_best=False): env = gym.make('CarRacing-v0', verbose=0) env = wrap_env(env, W, H, gray_scale=False) latent_size = 32 vae_model = VAE(latent_size) checkpoint_dir = './vae_ckpt/' latest = tf.train.latest_checkpoint(checkpoint_dir) vae_model.load_weights(latest) init = './saved/0297result_372.95431161691926' self.elite_from_file(init) #self.evolve(generations=0, pop_size=64, mode='init', n_end_eval=2, n_end_candidate_eval=8) #print(self.elite) #self.evolve(generations=1, pop_size=64, mode='normal') print(self.elite) if disp_best: self.elite.evaluate(env, vae_model, disp=True, n=1) for i in range(1, 1000): st = time() print('Era', i) #masks = self.elite.change_size(rnn_plus=0, controller_plus=16) #self.evolve(generations=20, pop_size=64, n_end_eval=2, n_end_candidate_eval=16) #self.evolve(generations=4, pop_size=64, mode='normal', masks=masks, n_end_eval=3, n_end_candidate_eval=8, n_survivors=8, eps=0.01) #print('rand evolve 1') masks = self.elite.get_random_mask( p=0.08, mask_mask=[False, False, False, True, True, False, False], with_bias=True)[0] self.evolve(generations=0, pop_size=96, mode='normal', masks=masks, n_end_eval=3, n_end_candidate_eval=16, n_survivors=8, eps=0.025) #print('rand evolve 2') #masks = self.elite.get_random_mask(p=1, mask_mask=[False, False, False, False, False, True, True], with_bias=True)[0] #self.evolve(generations=4, pop_size=64, mode='normal', masks=masks, n_end_eval=3, n_end_candidate_eval=8, n_survivors=8, eps=0.01) print('elite:', self.elite) if disp_best: self.elite.evaluate(env, vae_model, disp=True, n=1) if save_best: self.elite.save_all(name='{era:04d}'.format(era=i) + 'result:' + str(self.elite.mean_result)) print('TT:', time() - st) #print(masks) env.close()
def main(_): model_path = 'models/' + FLAGS.name data = load_data(FLAGS.dataset, one_hot=True, validation_size=10000) # Define and instantiate VAE model if FLAGS.vae_type == 'vae': vae = VAE(network_architecture=network_architecture( FLAGS.vae_type, FLAGS.latent_dim), batch_size=FLAGS.batch_size, learn_rate=FLAGS.learn_rate) elif FLAGS.vae_type == 'conv': vae = ConvVAE(network_architecture=network_architecture( FLAGS.vae_type, FLAGS.latent_dim), batch_size=FLAGS.batch_size, learn_rate=FLAGS.learn_rate) else: raise ValueError( "Autoencoder type should be either conv or vae. Received: {}.". format(FLAGS.vae_type)) with tf.Session() as sess: np.random.seed(FLAGS.seed) tf.set_random_seed(FLAGS.seed) saver = tf.train.Saver() saver.restore(sess, model_path) print("Model restored from: %s" % model_path) # Sample a test input and see how well the VAE can reconstruct these samples x_sample = data.test.next_batch(FLAGS.batch_size)[0] x_reconstruct = vae.reconstruct(sess, x_sample) plt.figure(figsize=(8, 12)) for i in range(5): plt.subplot(5, 2, 2 * i + 1) plt.imshow(x_sample[i].reshape(IMAGE_SIZE, IMAGE_SIZE), vmin=0, vmax=1, cmap='gray') plt.title("Test input") plt.colorbar() plt.subplot(5, 2, 2 * i + 2) plt.imshow(x_reconstruct[i].reshape(IMAGE_SIZE, IMAGE_SIZE), vmin=0, vmax=1, cmap='gray') plt.title("Reconstruction") plt.colorbar() plt.tight_layout() plt.show() visualise_latent_space(sess, vae, data.test) if FLAGS.latent_dim == 2: plot_reconstructions(sess, vae, FLAGS.batch_size)
def main(): """ Generate images from a saved model """ train_data = UnlabeledContact( data='/home/ygx/data/fspeptide/fs_peptide.npy') print('Number of samples: {}'.format(len(train_data))) trainloader = DataLoader(train_data, batch_size=args.batch_size) #encoder = Encoder(input_size=args.input_size, latent_size=args.latent_size) #decoder = Decoder(latent_size=args.latent_size, output_size=args.input_size) #vae = VAE(encoder, decoder, use_cuda=args.use_cuda) vae = VAE() # Load saved model vae.load_state_dict(torch.load(args.model_path + args.model_name)) if args.use_cuda: #encoder = encoder.cuda() #decoder = decoder.cuda() vae = vae.cuda() latent_arrys = [] recon_arrys = [] for batch_idx, data in enumerate(trainloader): inputs = data['cont_matrix'] inputs = inputs.resize_(args.batch_size, 1, 21, 21) inputs = inputs.float() if args.use_cuda: inputs = inputs.cuda() inputs = Variable(inputs) #latent_array = encoder(inputs).data.cpu().numpy() #print('latent_array has shape {}'.format(latent_array.shape)) #latent_arrys.append(latent_array) #reconstructed_array = vae(inputs).data.cpu().numpy() reconstructed_array, mu, _ = vae(inputs) reconstructed_array = reconstructed_array.data.cpu().numpy() latent_array = mu.data.cpu().numpy() recon_arrys.append(reconstructed_array) latent_arrys.append(latent_array) if batch_idx % 100 == 0: print('Saving progress: {:.3f}%'.format(batch_idx * 100. / len(trainloader))) print('\nNumber of images prepared: {}'.format(len(latent_arrys))) latent_stacked = np.stack(latent_arrys, axis=0) latent_filename = 'latent_imgs_fc' np.save(args.latent_save_path + latent_filename, latent_stacked) recon_stacked = np.stack(recon_arrys, axis=0) recon_filename = 'recon_imgs_fc' np.save(args.recon_save_path + recon_filename, recon_stacked)
def main(): use_cuda = args.use_cuda train_data = UnlabeledContact(data=args.data_dir) print('Number of samples: {}'.format(len(train_data))) trainloader = DataLoader(train_data, batch_size=args.batch_size) # Contact matrices are 21x21 input_size = 441 encoder = Encoder(input_size=input_size, latent_size=3) decoder = Decoder(latent_size=3, output_size=input_size) vae = VAE(encoder, decoder, use_cuda=use_cuda) criterion = nn.MSELoss() if use_cuda: encoder = nn.DataParallel(encoder) decoder = nn.DataParallel(decoder) encoder = encoder.cuda().half() decoder = decoder.cuda().half() vae = nn.DataParallel(vae) vae = vae.cuda().half() criterion = criterion.cuda().half() optimizer = optim.SGD(vae.parameters(), lr=0.01) clock = AverageMeter(name='clock16', rank=0) epoch_loss = 0 total_loss = 0 end = time.time() for epoch in range(15): for batch_idx, data in enumerate(trainloader): inputs = data['cont_matrix'] # inputs = inputs.resize_(args.batch_size, 1, 21, 21) inputs = inputs.float() if use_cuda: inputs = inputs.cuda().half() inputs = Variable(inputs) optimizer.zero_grad() dec = vae(inputs) ll = latent_loss(vae.z_mean, vae.z_sigma) loss = criterion(dec, inputs) + ll loss.backward() optimizer.step() epoch_loss += loss.data[0] clock.update(time.time() - end) end = time.time() if batch_idx % args.log_interval == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(trainloader.dataset), 100. * batch_idx / len(trainloader), loss.data[0])) clock.save(path='/home/ygx/libraries/mds/molecules/molecules/linear_vae')
def main(): mnist_train = datasets.MNIST('./../data', train=True, transform=transforms.Compose( [transforms.ToTensor()]), download=False) mnist_train = DataLoader(mnist_train, batch_size=128, shuffle=True) mnist_test = datasets.MNIST('../data', train=False, transform=transforms.Compose( [transforms.ToTensor()]), download=False) mnist_test = DataLoader(mnist_test, batch_size=128, shuffle=False) # 无监督学习,不需要label x, _ = iter(mnist_train).next() print('x: ', x.shape) device = torch.device('cuda') model = VAE().to(device) print(model) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) viz = visdom.Visdom() for epoch in range(100): for batchidx, (x, _) in enumerate(mnist_train): # [b,1,28,28] x = x.to(device) x_hat, kld = model(x) loss = criterion(x_hat, x) if kld is not None: elbo = -loss - 1.0 * kld loss = -elbo # backprop optimizer.zero_grad() loss.backward() optimizer.step() print('epoch: ', epoch, ' loss:', loss.item(), 'kld: ', kld.item()) x, _ = iter(mnist_test).next() x = x.to(device) with torch.no_grad(): x_hat, kld = model(x) viz.images(x, nrow=8, win='x', opts=dict(title='x')) viz.images(x_hat, nrow=8, win='x_hat', opts=dict(title='x_hat'))
def train(): mnist_train = datasets.MNIST('../data/mnist', train = True, transform=transforms.Compose([transforms.ToTensor()]), download=True) mnist_train = DataLoader(mnist_train, batch_size = 32, shuffle = True) mnist_test = datasets.MNIST('../data/mnist', train = False, transform=transforms.Compose([transforms.ToTensor()]), download=True) mnist_test = DataLoader(mnist_test, batch_size = 32, shuffle = True) #不需要label,因为是无监督学习 x, _ = iter(mnist_train).next() print('x:', x.shape) device = torch.device('cuda') model = VAE().to(device) criteon = nn.MSELoss() # loss function optimzer = optim.Adam(model.parameters(), lr = 1e-3) print(model) vis = visdom.Visdom() for epoch in range(1000): # 训练过程 for batchIdx, (x, _) in enumerate(mnist_train): #forwardp [b, 1, 28, 28] x = x.to(device) x_hat, kld = model(x) loss = criteon(x_hat, x) if kld is not None: elbo = - loss - 1.0 * kld lossX = - elbo #backward optimzer.zero_grad() lossX.backward() optimzer.step() # 打印loss print('epoch:', epoch, ' lossX:', lossX.item(),' loss:', loss.item(),' kld:', kld.item()) # 测试过程 x, _ = iter(mnist_test).next() x = x.to(device) with torch.no_grad(): #测试不用梯度 x_hat, _ = model(x) vis.images(x, nrow=8, win='x', opts=dict(title='x')) #画输入 vis.images(x_hat, nrow=8, win='x_hat', opts=dict(title='x_hat')) #画输出
def get_model(model_name='pretrained', drop_last=2): # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if model_name == 'vae': model = VAE(latent_dim=200, nf=128) #.to(device) feature_extractor = model.get_feature_extractor() else: model = models.resnet50(pretrained=True) feature_extractor = torch.nn.Sequential( *(list(model.children())[:-drop_last])) feature_extractor.eval() return feature_extractor
def __init__(self, args, obs_raw_shape): self.args = args self.obs_raw_shape = obs_raw_shape self._tasks = None self._latents = None self._trajectory_current = None self._rewards_current = None self.fit_counter = 0 self.history = History(args) self.clusterer = VAE(args, obs_raw_shape[0]) self._fitted = False
def main(): mnist_train = datasets.MNIST('mnist', True, transform=transforms.Compose([ transforms.ToTensor() ]), download=True) mnist_train = DataLoader(mnist_train, batch_size=32, shuffle=True) mnist_test = datasets.MNIST('mnist', False, transform=transforms.Compose([ transforms.ToTensor() ]), download=True) mnist_test = DataLoader(mnist_test, batch_size=32, shuffle=True) x, _ = iter(mnist_train).next() print('x:', x.shape) device = torch.device('cuda') # model = AE().to(device) model = VAE().to(device) criteon = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) print(model) viz = visdom.Visdom() for epoch in range(1000): for batchidx, (x, _) in enumerate(mnist_train): # [b, 1, 28, 28] x = x.to(device) x_hat, kld = model(x)#x_hat重建后的 loss = criteon(x_hat, x) if kld is not None: elbo = - loss - 1.0 * kld loss = - elbo # backprop optimizer.zero_grad() loss.backward() optimizer.step() # kld比loss低才正常 print(epoch, 'loss:', loss.item(), 'kld:', kld.item()) x, _ = iter(mnist_test).next()#_不要忘记占位符了,否则会出现x.size(0)没有size这个功能 x = x.to(device) with torch.no_grad(): x_hat, kld = model(x) viz.images(x, nrow=8, win='x', opts=dict(title='x')) viz.images(x_hat, nrow=8, win='x_hat', opts=dict(title='x_hat'))
print (len(dataset[ii][0])) # single timepoint print (dataset[ii][0][0].shape) #action [1] a_t+1 print (dataset[ii][0][1].shape) #state [2,84,84] s_t state_dataset = [] for i in range(len(dataset)): for t in range(len(dataset[i])): state_dataset.append(dataset[i][t][1]) print (len(state_dataset)) print('Init VAE') vae = VAE() vae.cuda() load_ = 1 train_ = 1 viz_ = 1 if load_: load_epoch = 50 path_to_load_variables = home+'/Documents/tmp/breakout_2frames/vae_params'+str(load_epoch)+'.ckpt' vae.load_params(path_to_load_variables) epochs = 100 if load_: path_to_save_variables = home+'/Documents/tmp/breakout_2frames/vae_params'+str(epochs+load_epoch)+'.ckpt'
''' # Load MNIST data in a format suited for tensorflow from tensorflow.examples.tutorials.mnist import input_data mnist = input_data.read_data_sets('MNIST_data/', one_hot=True) # ------------------------------------------------- nsamples = mnist.train.num_examples data_dim = 784 dataset_name = 'MNIST' batch_size = 100 epochs = 75 network_structure = [data_dim, 500, 500, 20] vae = VAE(network_structure, learning_rate=0.001, batch_size=batch_size) final_cost = 0. # training cycle start = time.time() for epoch in range(epochs): avg_cost = 0. total_batch = int(nsamples/batch_size) for i in range(total_batch): sample, _ = mnist.train.next_batch(batch_size) cost, label = vae.fit(sample) avg_cost += cost / nsamples * batch_size final_cost = avg_cost