def __init__(self, im_shape, auroc=None, lr=None, batch_size=None, epochs=None, progress_bar=None, early_stopping=None, layer_idx=None, z_dim=None, dataset=None, num_workers=None, train_digit=None, test_digit=None, mvtec_object=None, inference_mode=None, sample_during_training=None, sample_every_n_epoch=None, eval=None, model_version=None, init_seed=None): """ args: lr - Learning rate used for training inference_mode - Which kind of inferencing to use to calculate the score to be backpropagated on options are 'mean_sum' and 'normal_diff' layer_idx - Index for which layer of the network to use the activations and gradients for for attention map generation z_dim - The latent dimension size to use for encoding/decoding auroc - Boolean switch, indicating whether we have access to target masks and if we want to compute the AUROC scores for them. MNIST doesn't have target masks, UCSD and MVTEC do Other arguments are listed are not functional, but for logging, as save_hyperparameters() allows us to track all arguments with tensorboard """ super().__init__() # Set defaults to None when arguments are irrelevant if 'mvtec' not in dataset: mvtec_object = None elif 'mnist' not in dataset: train_digit, test_digit = None, None self.save_hyperparameters() self.vae = VAE(self.hparams.layer_idx, self.hparams.z_dim, self.hparams.im_shape) if self.hparams.auroc: self.roc = metrics.ROC(pos_label=1)
def train_vae(): batch_size = 64 epochs = 1000 latent_dimension = 100 patience = 10 device = torch.device( 'cuda:0') if torch.cuda.is_available() else torch.device('cpu') # load data train_loader, valid_loader, _ = get_data_loader('data', batch_size) model = VAE(latent_dimension).to(device) optim = Adam(model.parameters(), lr=1e-3) val_greater_count = 0 last_val_loss = 0 for e in range(epochs): running_loss = 0 model.train() for i, (images, _) in enumerate(train_loader): images = images.to(device) model.zero_grad() outputs, mu, logvar = model(images) loss = compute_loss(images, outputs, mu, logvar) running_loss += loss loss.backward() optim.step() running_loss = running_loss / len(train_loader) model.eval() with torch.no_grad(): val_loss = 0 for images, _ in valid_loader: images = images.to(device) outputs, mu, logvar = model(images) loss = compute_loss(images, outputs, mu, logvar) val_loss += loss val_loss /= len(valid_loader) if val_loss > last_val_loss: val_greater_count += 1 else: val_greater_count = 0 last_val_loss = val_loss torch.save( { 'epoch': e, 'model': model.state_dict(), 'running_loss': running_loss, 'optim': optim.state_dict(), }, "vae/upsample_checkpoint_{}.pth".format(e)) print("Epoch: {} Train Loss: {}".format(e + 1, running_loss.item())) print("Epoch: {} Val Loss: {}".format(e + 1, val_loss.item())) if val_greater_count >= patience: break
def load_robot_VAE(self): if self.target_robot == 'BAXTER': load_name = '[Baxter]Motion_vae_Superset' elif self.target_robot == 'NAO': load_name = '[NAO]Motion_vae_Superset' elif self.target_robot == 'C3PO': load_name = '[C3PO]Motion_vae_Superset' else: print('Unknown robot..') raise ValueError load_path = os.path.join('./trained_models', 'vae') state_dict = torch.load(os.path.join(load_path, load_name + '.pt'), map_location='cpu') rModel = VAE(14, latent_dim=np.prod(self.args.action_shape), use_batch_norm=False, activation='Tanh') rModel.load_state_dict(state_dict) rModel.eval() print('Motion VAE Load Success!') if not self.args.symm_policy: self.myrobot.rModel = rModel return rModel
def train(config): # Print all configs to confirm parameter settings print_flags() # Initialize the model that we are going to use # model = LSTMLM(vocabulary_size=vocab_size, model = VAE(vocabulary_size=vocab_size, dropout=1 - config.dropout_keep_prob, lstm_num_hidden=config.lstm_num_hidden, lstm_num_layers=config.lstm_num_layers, lstm_num_direction=config.lstm_num_direction, num_latent=config.num_latent, device=device) model.to(device) # Setup the loss and optimizer criterion = nn.CrossEntropyLoss(ignore_index=1, reduction='sum') optimizer = optim.Adam(model.parameters(), lr=config.learning_rate) # Store some measures iteration = list() tmp_loss = list() train_loss = list() val_nll = list() val_perp = list() val_acc = list() val_elbo = list() train_perp = list() train_acc = list() train_elbo = list() train_nll = list() iter_i = 0 best_perp = 1e6 while True: # when we run out of examples, shuffle and continue for train_batch in get_minibatch(train_data, batch_size=config.batch_size): # Only for time measurement of step through network t1 = time.time() iter_i += 1 model.train() optimizer.zero_grad() inputs, targets, lengths_in_batch = prepare_minibatch( train_batch, vocab) # zeros in dim = (num_layer*num_direction * batch * lstm_hidden_size) # we have bidrectional single layer LSTM h_0 = torch.zeros( config.lstm_num_layers * config.lstm_num_direction, inputs.shape[0], config.lstm_num_hidden).to(device) c_0 = torch.zeros( config.lstm_num_layers * config.lstm_num_direction, inputs.shape[0], config.lstm_num_hidden).to(device) # pred, _, _ = model(inputs, h_0, c_0) decoder_output, KL_loss = model(inputs, h_0, c_0, lengths_in_batch, config.importance_sampling_size) reconstruction_loss = 0.0 for k in range(config.importance_sampling_size): # the first argument for criterion, ie, crossEntrooy must be (batch, classes(ie vocab size), sent_length), so we need to permute the last two dimension of decoder_output (batch, sent_length, vocab_classes) # decoder_output[k] =decoder_output[k].permute(0, 2, 1) doesnt work reconstruction_loss += criterion( decoder_output[k].permute(0, 2, 1), targets) # get the mean of the k samples of z reconstruction_loss = reconstruction_loss / config.importance_sampling_size KL_loss = KL_loss / config.importance_sampling_size print('At iter', iter_i, ', rc_loss=', reconstruction_loss.item(), ' KL_loss = ', KL_loss.item()) total_loss = (reconstruction_loss + KL_loss) / config.batch_size tmp_loss.append(total_loss.item()) total_loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.max_norm) optimizer.step() if iter_i % config.eval_every == 0: eval_data = val_data eval_data_flag = 'val' print('Evaluating with validation at iteration ', iter_i, '...') if iter_i % config.eval_every_train == 0: eval_data = train_data eval_data_flag = 'train' print('Evaluating with training instead, at iteration ', iter_i, '...') model.eval() ppl_total = 0.0 validation_elbo_loss = 0.0 validation_lengths = list() nll_per_eval = list() match = list() with torch.no_grad(): # computing ppl, match, and accuracy for validation_th, val_sen in enumerate(eval_data): val_input, val_target = prepare_example(val_sen, vocab) # zeros in dim = (num_layer*num_direction, # batch=config.importance_sampling_size, lstm_hidden_size) h_0 = torch.zeros( config.lstm_num_layers * config.lstm_num_direction, config.importance_sampling_size, config.lstm_num_hidden).to(device) c_0 = torch.zeros( config.lstm_num_layers * config.lstm_num_direction, config.importance_sampling_size, config.lstm_num_hidden).to(device) # append the sent length of this particular validation example validation_lengths.append(val_input.size(1)) # feed into models decoder_output, KL_loss_validation = model( val_input, h_0, c_0, [val_input.size(1)], config.importance_sampling_size) # decoder_output.size() = (k, batchsize=1, val_input.size(1)(ie sent_length), vocabsize) # prediction.size() = (k, sent_len, vocabsize) # prediction_mean.size() = (sent_len, vocabsize), ie averaged over # k samples (and squeezed) prediction = nn.functional.softmax(torch.squeeze( decoder_output, dim=1), dim=2) prediction_mean = torch.mean(prediction, 0) # averaged over k ppl_per_example = 0.0 # sentence length, ie 1 word/1 timestamp for each loop for j in range(prediction.shape[1]): # 0 as the target is the same for the k samples ppl_per_example -= torch.log( prediction_mean[j][int(val_target[0][j])]) ppl_total += ppl_per_example if validation_th % 300 == 0: print(' ppl_per_example at the ', validation_th, eval_data_flag, 'case = ', ppl_per_example) tmp_match = compute_match_vae(prediction_mean, val_target) match.append(tmp_match) # calculate validation elbo # decoder_output.size() = (k, batchsize=1, val_input.size(1)(ie sent_length), vocabsize) # the first argument for criterion, ie, crossEntrooy must be (batch, classes(ie vocab size), sent_length), so we need to permute the last two dimension of decoder_output to get (k, batchsize=1, vocab_classes, sent_length) # then we loop over k to get (1, vocab_classes, sent_len) decoder_output_validation = decoder_output.permute( 0, 1, 3, 2) reconstruction_loss = 0 for k in range(config.importance_sampling_size): reconstruction_loss += criterion( decoder_output_validation[k], val_target) validation_elbo_loss += (reconstruction_loss + \ KL_loss_validation) / config.importance_sampling_size nll_per_eval.append(ppl_per_example) ppl_total = torch.exp(ppl_total / sum(validation_lengths)) print('ppl_total for iteration ', iter_i, ' = ', ppl_total) accuracy = sum(match) / sum(validation_lengths) print('accuracy for iteration ', iter_i, ' = ', accuracy) # loss of the previous iterations (up the after last eval) avg_loss = sum(tmp_loss) / len(tmp_loss) tmp_loss = list() # reinitialize to zero validation_elbo_loss = validation_elbo_loss / len(val_data) if ppl_total < best_perp: best_perp = ppl_total torch.save(model.state_dict(), "./models/vae_best.pt") # Instead of rewriting the same file, we can have new ones: # model_saved_name = datetime.now().strftime("%Y-%m-%d_%H%M") + './models/vae_best.pt' # torch.save(model.state_dict(), model_saved_name) nll = sum(nll_per_eval) print( "[{}] Train Step {:04d}/{:04d}, " "Validation Perplexity = {:.4f}, Validation loss ={:.4f}, Training Loss = {:.4f}, NLL = {:.4f}" "Validation Accuracy = {:.4f}".format( datetime.now().strftime("%Y-%m-%d %H:%M"), iter_i, config.train_steps, ppl_total, validation_elbo_loss, avg_loss, nll, accuracy)) # update/save eval results everytime iteration.append(iter_i) train_loss.append(avg_loss) np.save('./np_saved_results/train_loss.npy', train_loss + ['till_iter_' + str(iter_i)]) if eval_data_flag == 'val': val_perp.append(ppl_total.item()) val_acc.append(accuracy) val_elbo.append(validation_elbo_loss.item()) val_nll.append(nll) np.save('./np_saved_results/val_perp.npy', val_perp + ['till_iter_' + str(iter_i)]) np.save('./np_saved_results/val_acc.npy', val_acc + ['till_iter_' + str(iter_i)]) np.save('./np_saved_results/val_elbo.npy', val_elbo + ['till_iter_' + str(iter_i)]) np.save('./np_saved_results/val_nll.npy', val_elbo + ['till_iter_' + str(iter_i)]) if eval_data_flag == 'train': train_perp.append(ppl_total.item()) train_acc.append(accuracy) train_elbo.append(validation_elbo_loss.item()) train_nll.append(nll) np.save('./np_saved_results/train_perp.npy', train_perp + ['till_iter_' + str(iter_i)]) np.save('./np_saved_results/train_acc.npy', train_acc + ['till_iter_' + str(iter_i)]) np.save('./np_saved_results/train_elbo.npy', train_elbo + ['till_iter_' + str(iter_i)]) np.save('./np_saved_results/train_nll.npy', train_elbo + ['till_iter_' + str(iter_i)]) if iter_i == config.train_steps: break if iter_i == config.train_steps: break print('Done training!') print('+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-') print('Testing...') print('+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-') model.load_state_dict(torch.load('./models/vae_best.pt')) model.eval() ppl_total = 0.0 validation_elbo_loss = 0.0 validation_lengths = list() nll_per_eval = list() match = list() with torch.no_grad(): # computing ppl, match, and accuracy # too large too slow lets stick with first 1000/1700 first for validation_th, val_sen in enumerate(test_data): val_input, val_target = prepare_example(val_sen, vocab) # zeros in dim = (num_layer*num_direction, # batch=config.importance_sampling_size, lstm_hidden_size) h_0 = torch.zeros( config.lstm_num_layers * config.lstm_num_direction, config.importance_sampling_size, config.lstm_num_hidden).to(device) c_0 = torch.zeros( config.lstm_num_layers * config.lstm_num_direction, config.importance_sampling_size, config.lstm_num_hidden).to(device) # append the sent length of this particular validation example validation_lengths.append(val_input.size(1)) # feed into models decoder_output, KL_loss_validation = model( val_input, h_0, c_0, [val_input.size(1)], config.importance_sampling_size) # decoder_output.size() = (k, batchsize=1, val_input.size(1)(ie sent_length), vocabsize) # prediction.size() = (k, sent_len, vocabsize) # prediction_mean.size() = (sent_len, vocabsize), ie averaged over k # samples (and squeezed) prediction = nn.functional.softmax(torch.squeeze(decoder_output, dim=1), dim=2) prediction_mean = torch.mean(prediction, 0) # averaged over k ppl_per_example = 0.0 # sentence length, ie 1 word/1 timestamp for each loop for j in range(prediction.shape[1]): # 0 as the target is the same for the k samples ppl_per_example -= torch.log(prediction_mean[j][int( val_target[0][j])]) ppl_total += ppl_per_example tmp_match = compute_match_vae(prediction_mean, val_target) match.append(tmp_match) # calculate validation elbo # decoder_output.size() = (k, batchsize=1, val_input.size(1)(ie sent_length), vocabsize) # the first argument for criterion, ie, crossEntrooy must be (batch, classes(ie vocab size), sent_length), so we need to permute the last two dimension of decoder_output to get (k, batchsize=1, vocab_classes, sent_length) # then we loop over k to get (1, vocab_classes, sent_len) decoder_output_validation = decoder_output.permute(0, 1, 3, 2) reconstruction_loss = 0 for k in range(config.importance_sampling_size): reconstruction_loss += criterion(decoder_output_validation[k], val_target) validation_elbo_loss += (reconstruction_loss + \ KL_loss_validation) / config.importance_sampling_size nll_per_eval.append(ppl_total) ppl_total = torch.exp(ppl_total / sum(validation_lengths)) accuracy = sum(match) / sum(validation_lengths) validation_elbo_loss = validation_elbo_loss / len(test_data) nll = sum(nll_per_eval) print('Test Perplexity on the best model is: {:.3f}'.format(ppl_total)) print( 'Test ELBO on the best model is: {:.3f}'.format(validation_elbo_loss)) print('Test accuracy on the best model is: {:.3f}'.format(accuracy)) print('Test NLL on the best model is: {:.3f}'.format(nll)) print('+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-') with open('./result/vae_test.txt', 'a') as file: file.write( 'Learning Rate = {}, Train Step = {}, ' 'Dropout = {}, LSTM Layers = {}, ' 'Hidden Size = {}, Test Perplexity = {:.3f}, Test ELBO = {:.3f}, Test NLL = {:.3f}' 'Test Accuracy = {}\n'.format(config.learning_rate, config.train_steps, 1 - config.dropout_keep_prob, config.lstm_num_layers, config.lstm_num_hidden, ppl_total, validation_elbo_loss, nll, accuracy)) file.close() print('Sampling...') print('+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-') # model.load_state_dict(torch.load('./models/vae_best_lisa.pt')) model.load_state_dict( torch.load('./models/vae_best_lisa.pt', map_location=lambda storage, loc: storage)) with torch.no_grad(): sentences = model.sample(config.sample_size, vocab) sentences_pruned_EOS = [[] for x in range(config.sample_size)] for i in range(len(sentences)): for j in range(len(sentences[i])): if sentences[i][j] != 'EOS': sentences_pruned_EOS[i].append(sentences[i][j]) else: break with open('./result/vae_test_greedy_new.txt', 'a') as file: for idx, sen in enumerate(sentences_pruned_EOS): if idx == 0: file.write('\n Greedy: \n') file.write('Sampling \n{}: {}\n'.format(idx, ' '.join(sen))) else: file.write('Sampling \n{}: {}\n'.format(idx, ' '.join(sen))) print('Interpolating...') print('+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-') #interpolation with torch.no_grad(): sentences = model.interpolation(vocab) sentences_pruned_EOS = [[], [], [], [], []] for i in range(len(sentences)): for j in range(len(sentences[i])): if sentences[i][j] != 'EOS': sentences_pruned_EOS[i].append(sentences[i][j]) else: break with open('./result/vae_test_interpolate.txt', 'a') as file: file.write('\n Interpolation: \n') file.write('Sampling z1:\n {}\n'.format(' '.join( sentences_pruned_EOS[0]))) file.write('Sampling z2:\n {}\n'.format(' '.join( sentences_pruned_EOS[1]))) file.write('Sampling z1+z2/2:\n {}\n'.format(' '.join( sentences_pruned_EOS[2]))) file.write('Sampling z1*0.8+z2*0.2:\n {}\n'.format(' '.join( sentences_pruned_EOS[3]))) file.write('Sampling z1*0.2+z2*0.8:\n {}\n'.format(' '.join( sentences_pruned_EOS[4]))) print('Test case reconstruction...') print('+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-') test_sen = test_data[101] # print('test_sen', test_sen) test_input, _ = prepare_example(test_sen, vocab) # print('test_input',test_input) # zeros in dim = (num_layer*num_direction, # batch=config.importance_sampling_size, lstm_hidden_size) h_0 = torch.zeros(config.lstm_num_layers * config.lstm_num_direction, config.importance_sampling_size, config.lstm_num_hidden).to(device) c_0 = torch.zeros(config.lstm_num_layers * config.lstm_num_direction, config.importance_sampling_size, config.lstm_num_hidden).to(device) # feed into models reconstructed_sentences = model.test_reconstruction(test_input, vocab) sentences_pruned_EOS = [[] for x in range(10)] for i in range(len(reconstructed_sentences)): for j in range(len(reconstructed_sentences[i])): if reconstructed_sentences[i][j] != 'EOS': sentences_pruned_EOS[i].append(reconstructed_sentences[i][j]) else: break with open('./result/vae_test_reconstruct.txt', 'a') as file: file.write('\n The sentence to reconstruct:\n {}\n'.format(' '.join( test_sen[1:]))) for x in range(10): file.write('Sample: {} \n {}\n'.format( x, ' '.join(sentences_pruned_EOS[x]))) ''' t_loss = plt.figure(figsize = (6, 4)) plt.plot(iteration, train_loss) plt.xlabel('Iteration') plt.ylabel('Training Loss') t_loss.tight_layout() t_loss.savefig('./result/vae_training_loss.eps', format='eps') v_perp = plt.figure(figsize = (6, 4)) plt.plot(iteration, val_perp) plt.xlabel('Iteration') plt.ylabel('Validation Perplexity') v_perp.tight_layout() v_perp.savefig('./result/vae_validation_perplexity.eps', format='eps') v_acc = plt.figure(figsize = (6, 4)) plt.plot(iteration, val_acc) plt.xlabel('Iteration') plt.ylabel('Validation Accuracy') v_acc.tight_layout() v_acc.savefig('./result/vae_validation_accuracy.eps', format='eps') v_elbo = plt.figure(figsize = (6, 4)) plt.plot(iteration, val_elbo) plt.xlabel('Iteration') plt.ylabel('Validation ELBO') v_elbo.tight_layout() v_elbo.savefig('./result/vae_validation_elbo.eps', format='eps') print('Figures are saved.') print('+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-') ''' return 0
from torchsummary import summary from torch.utils.data import DataLoader import numpy as np torch.manual_seed(2) place = '119' if place == 'azure': data_path = '/home/hung/DLCV2018SPRING/hw4/hw4_data/' else: data_path = '/home/lilioo826/hw4_data/' train_faceDataset = FaceDataset(data_path + 'train', data_path + 'train.csv', transforms.ToTensor()) train_dataloader = DataLoader(train_faceDataset, batch_size=20, num_workers=1) cuda = True model = VAE(64, 1e-6) # print(model) if cuda: model.cuda() # summary(model, (3,64,64)) # exit() epoch_num = 100 model.train() optimizer = optim.Adam(model.parameters(), lr=1e-4) klds = [] mses = [] for epoch in range(epoch_num): print('epoch {}'.format(epoch + 1)) epoch_kld = 0
@author: Caroline Pacheco @email: [email protected] """ import pandas as pd from vae_model import VAE from trainer import * from samples import * from data import * from LBP.eq_validity import equation_validity import collections import random # train a vae instance model model = VAE(vocab, vector).to(device) fit(model, train_data) # generate sample equations sample = sample.take_samples(model, n_batch, 3000) # validate generated equations valid, not_valid, max_equations, unseenval = equation_validity(sample) # save output file df = pd.DataFrame(unseenval) df.to_csv(r'output/my_equation.txt', header=None, index=None, sep='\t', mode='a')
data_set, [math.floor(TRAIN_VAL_SPLIT_CONSTANT * data_set.shape[0])]) print('x_train' 's shape is: ' + str(x_train.shape)) print('shuffling...') train_ds = tf.data.Dataset.from_tensor_slices(x_train).shuffle(500000).batch( batch_size) print('shuffle of train ds completed.') test_ds = tf.data.Dataset.from_tensor_slices(x_val).batch(batch_size) print('shuffle completed.') print('model declaring...') if args.model: print('model loading from file...') vae = tf.keras.models.load_model(args.model) else: vae = VAE() print('model declared.') # vae.summary() # plot_model(vae, to_file='vae_cnn.png', show_shapes=True) # ret1 = None # ret2 = None best_test_loss = math.inf last_best_epoch = 0 best_model_route = None train_branch = False automatic_switch = True print('hyper-parameter for weight: img:{}, reward:{}, kl:{}, done:{}'.format( loss_weight['image_loss'], loss_weight['reward_loss'], loss_weight['latent_loss'], loss_weight['done_loss'])) print('start training...') for epoch in range(epochs):
import torch import numpy as np import matplotlib.pyplot as plt from matplotlib.widgets import Slider, Button from arguments import Loader, get_args from vae_model import VAE args = get_args() loader = Loader('./dataset', 10000) vae = VAE() vae.load_state_dict(torch.load('./trained_models/{}'.format(args.run_id))) vae.eval() # Assuming code dimension = 32 nb_axes = 32 ax_margin = 0.02 ax_width = 0.1 ax_updown_margin = 0.1 ax_left = 0.85 inc = (1. - 2. * ax_updown_margin) / (float(nb_axes)) f, ax = plt.subplots() plt.subplots_adjust(left=0.1, right=0.8) slider_min = -2. slider_max = 2. all_axes = [
from trainer import * from vae_model import VAE from data import * from samples import * model = VAE(vocab, vector).to(device) fit(model, train_data) model.eval() sample = sample.take_samples(model, n_batch) print(sample)
sModel = VAE_SK(75, 75, use_batch_norm=False, activation='ReLU') sModel.load_state_dict(state_dict) sModel.eval() # Load Motion vae Model if myrobot.robot_name == 'BAXTER_MIMIC': load_name = '[Baxter]Motion_vae_Superset' elif myrobot.robot_name == 'NAO_MIMIC': load_name = '[NAO]Motion_vae_Superset' elif myrobot.robot_name == 'C3PO_MIMIC': load_name = '[C3PO]Motion_vae_Superset' print('Motion Vae Name: ', load_name) state_dict = torch.load(os.path.join(load_path, load_name + '.pt')) rModel = VAE(14, latent_dim=np.prod(action_shape), use_batch_norm=False, activation='Tanh') # action: 10*20 rModel.load_state_dict(state_dict) rModel.eval() if not args.symm_policy: myrobot.rModel = rModel # -----[ For socket Communication, Client ] ----- if ntu_socket_comm: ntu_s = SocketCom('localhost', ntu_tcp_port) ntu_s.socket_connect() num_eval_frames = 5000 reward_list = []
class ExpVAE(pl.LightningModule): def __init__(self, im_shape, auroc=None, lr=None, batch_size=None, epochs=None, progress_bar=None, early_stopping=None, layer_idx=None, z_dim=None, dataset=None, num_workers=None, train_digit=None, test_digit=None, mvtec_object=None, inference_mode=None, sample_during_training=None, sample_every_n_epoch=None, eval=None, model_version=None, init_seed=None): """ args: lr - Learning rate used for training inference_mode - Which kind of inferencing to use to calculate the score to be backpropagated on options are 'mean_sum' and 'normal_diff' layer_idx - Index for which layer of the network to use the activations and gradients for for attention map generation z_dim - The latent dimension size to use for encoding/decoding auroc - Boolean switch, indicating whether we have access to target masks and if we want to compute the AUROC scores for them. MNIST doesn't have target masks, UCSD and MVTEC do Other arguments are listed are not functional, but for logging, as save_hyperparameters() allows us to track all arguments with tensorboard """ super().__init__() # Set defaults to None when arguments are irrelevant if 'mvtec' not in dataset: mvtec_object = None elif 'mnist' not in dataset: train_digit, test_digit = None, None self.save_hyperparameters() self.vae = VAE(self.hparams.layer_idx, self.hparams.z_dim, self.hparams.im_shape) if self.hparams.auroc: self.roc = metrics.ROC(pos_label=1) def loss_f(self, x_rec, x, stats): """ Function which calculates the VAE loss using BCE and KL divergence x - that original "target" images x_rec - reconstructed images from the model stats - contains distribution p and q, the sampled latent z, and the output of the encoder, mu and _logvar """ # Unpack and define some variables for calculating the loss p, q, z, mu, log_var = stats im_shape = x.shape[2:] # For MVTEC images, we use mean MSE for reconstruction loss if im_shape == (256, 256): L_rec = F.mse_loss(x_rec, x, reduction='mean') # For others, we use summed BCE for reconstruction loss) else: L_rec = F.binary_cross_entropy(x_rec, x, reduction='sum') # Compute KL divergence between encoder and unit Gaussian log_qz = q.log_prob(z) log_pz = p.log_prob(z) kl = log_qz - log_pz kl = kl.mean() L_reg = kl*0.1 # Compute final loss loss = L_rec + L_reg return loss def training_step(self, batch, batch_idx): """ Defines a single training iteration with the goal of reconstructing the input """ x, _ = batch x_rec, stats = self.vae(x) loss = self.loss_f(x_rec, x, stats) self.log('train_loss', loss) return loss def validation_step(self, batch, batch_idx, loader_idx): """ Defines a single validation iteration with the goal of reconstructing the input """ x, _ = batch # First dataloader contains images of the class we're training on, so we compute # regular VAE loss if loader_idx == 0: x_rec, stats = self.vae(x) loss = self.loss_f(x_rec, x, stats) self.log('val_loss', loss) return loss elif loader_idx == 1: _, ground_truth = batch # For MNIST, there are no ground truth masks, so we don't compute them if self.hparams.auroc: x, ground_truth = batch ground_truth = ground_truth.to(self.device) # Compute attention maps _, attmaps, _ = self.forward(x) # Update the ROC with new predictions and ground truths self.roc.update(attmaps, ground_truth) def validation_epoch_end(self, outputs): """ After going through the entire validation set, we compute the final ROC curve accumulated overall predictions and target masks, then compute the AUROC """ if self.hparams.auroc: fpr, tpr, thresholds = self.roc.compute() fpr, idx = torch.sort(fpr, descending=False) tpr, thresholds = tpr[idx], thresholds[idx] auroc = auc(fpr, tpr) self.log('auroc', auroc) def test_step(self, batch, batch_idx, loader_idx): """ Defines a single testing iteration with the goal of reconstructing the input For MNIST, it measures VAE loss with goal of reconstructing the input and attention maps on outlier class images For UCSD and MVTEC, it also saves ground truth masks and input images """ x, _ = batch if loader_idx == 0: x_rec, stats = self.vae(x) loss = self.loss_f(x_rec, x, stats) self.log('test_loss', loss) return loss elif loader_idx == 1: _, ground_truth = batch x_rec, M, colormaps = self.forward(x) x_rec, M, colormaps = x_rec.detach().cpu(), M.detach().cpu(), colormaps.detach().cpu() colormaps = make_grid(colormaps) save_image(colormaps.float(), f'{self.trainer.logger.log_dir}/batch{batch_idx}-attmaps.png') # For MNIST, there are no ground truth masks, so we don't compute them if self.hparams.auroc: x, ground_truth = batch ground_truth = ground_truth.to(self.device) # Compute attention maps _, attmaps, _ = self.forward(x) # Update the ROC with new predictions and ground truths self.roc.update(attmaps, ground_truth) def test_epoch_end(self, outputs): """ After going through the entire test set, we compute the final ROC curve accumulated overall predictions and target masks, then compute the AUROC """ if self.hparams.auroc: # Compute ROC, then compute AUROC and log the value for the whole test set fpr, tpr, thresholds = self.roc.compute() fpr, idx = torch.sort(fpr, descending=False) tpr, thresholds = tpr[idx], thresholds[idx] auroc = auc(fpr, tpr) self.log('auroc_test', auroc) # Divide thresholds from ROC into 100 equally separated thresholds step_size = int(len(thresholds)/100) thresholds = thresholds[::step_size] # Find best best threshold based off of best IOU best_iou = 0 best_threshold = -1 # For each threshold, compute IOU for whole test set for i, threshold in enumerate(thresholds): test_dataloader = self.trainer.datamodule.test_dataloader()[1] ious = [] for batch_idx, (x, y) in enumerate(test_dataloader): x, y = x.to(self.device), y.to(self.device) x_rec, M, colormaps = self.forward(x) bloc_map = self.gen_bloc_map(M, threshold) iou_score = iou(bloc_map, y) ious.append(iou_score.detach().cpu().item()) avg_iou = np.mean(ious) if avg_iou > best_iou: best_iou = avg_iou best_threshold = threshold self.trainer.logger.experiment.add_scalar('avg_iou', avg_iou, i) self.trainer.logger.experiment.add_scalar('threshold', threshold, i) # Log best iou and threshold self.log('best_iou', best_iou) self.log('best_threshold', best_threshold) # Now, using best threshold, generate the binary localization maps for # all images in the test set and log/save them for batch_idx, (x, y) in enumerate(test_dataloader): x, y = x.to(self.device), y.to(self.device) x_rec, M, colormaps = self.forward(x) bloc_map = self.gen_bloc_map(M, best_threshold) # Save the binary localization maps bloc_map = bloc_map.detach().cpu() bloc_map_grid = make_grid(bloc_map).float() save_image(bloc_map_grid, f'{self.trainer.logger.log_dir}/batch{batch_idx}-blocmaps.png') self.trainer.logger.experiment.add_image('blocmaps', bloc_map_grid.numpy(), batch_idx) # Save the input images x = x.detach().cpu() x = self.trainer.datamodule.unnormalize_batch(x) x_grid = make_grid(x).float() save_image(x_grid, f'{self.trainer.logger.log_dir}/batch{batch_idx}-input.png') self.trainer.logger.experiment.add_image('input', x_grid.numpy(), batch_idx) # Save teh target masks y = y.detach().cpu() y_grid = make_grid(y).float() save_image(y_grid, f'{self.trainer.logger.log_dir}/batch{batch_idx}-targets.png') self.trainer.logger.experiment.add_image('targets', y_grid.numpy(), batch_idx) def forward(self, x): """ Forward function which reconstructs input, and also returns attention, color and binary localization maps Note that this function is not used during training, only for evaluation """ # Set model to eval, and zero out gradients self.eval() self.zero_grad() # Make sure gradients are enabled (PyTorch Lightning disables gradients for validation loop, which calls this function) with torch.set_grad_enabled(True): x = x.to(self.device) # Push images through the network to get reconstruction, and mu for computing the score to backprop on x_rec, stats = self.vae(x) p, q, z, mu, log_var = stats n_channels = x.shape[1] # For mean sum inference, we simply sum the mu vector to compute the score if self.hparams.inference_mode == 'mean_sum': score = torch.sum(mu) elif self.hparams.inference_mode == 'normal_diff': z = self.vae.norm_diff_reparametrize(mu, log_var) score = torch.sum(z) # Make sure the score if score.requires_grad == False: score.requires_grad_() score.backward(retain_graph=True) # Retrieve the activations and gradients from the specific layer dz_da, A = self.vae.get_layer_data() # We can now compute the attention maps M and create the color maps dz_da = dz_da / (torch.sqrt(torch.mean(torch.square(dz_da))) + 1e-5) alpha = F.avg_pool2d(dz_da, kernel_size=dz_da.shape[2:]) A, alpha = A, alpha A, alpha = A.unsqueeze(0), alpha.unsqueeze(1) M = F.conv3d(A, (alpha), padding=0, groups=len(alpha)).squeeze(0).squeeze(1) M = F.interpolate(M.unsqueeze(1), size=self.hparams.im_shape[1:], mode='bilinear', align_corners=False) M = torch.abs(M) colormaps = self.create_colormap(x, M) # Zero out the gradients again, and put model back into train mode self.zero_grad() self.train() return x_rec, M, colormaps def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr) return optimizer def set_normal_stats(self, mu, log_var): """ Sets the VAE's inference method use the difference between the distribution of the trained embeddings vs the distribution of the outlier class """ self.vae.configure_normal(mu=mu, log_var=log_var) self.hparams.inference_mode = 'normal_diff' def create_colormap(self, x, attmaps, unnormalize=True): """ Creates and returns a colormap from the attention map and original input image x - original input images attmaps - attention maps from the model inferred from the input images """ if unnormalize: x = self.trainer.datamodule.unnormalize_batch(x) attmaps = attmaps.detach() n_channels = x.shape[1] if n_channels == 1: x = x.repeat(1, 3, 1, 1) colormaps = torch.zeros(x.shape) for i in range(x.size(0)): raw_image = x[i] * 255.0 ndarr = raw_image.permute(1, 2, 0).cpu().byte().numpy() im = Image.fromarray(ndarr.astype(np.uint8)) r_im = np.asarray(im) gcam = get_cam(r_im, attmaps[i].squeeze().cpu().data.numpy()) colormaps[i] = torch.from_numpy(gcam).permute(2,0,1)/255 permute = [2,1,0] colormaps = colormaps[:, permute] return colormaps def gen_bloc_map(self, M, threshold): # Generates a binary localizatino map given a threshold M = (M > threshold).to(torch.int) return M
if epoch%args.save_interval == 0: save_path = '/home/mehdi/Codes/ML3/WorldModels/trained_models/' try: os.makedirs(save_path) except OSError: pass if args.use_cuda: model.cpu() torch.save(model.state_dict(),save_path + args.run_id) if args.use_cuda: model.cuda() if(epoch%args.im_interval == 0): current_sample, current_recon, current_im = run_sample(model, loader, args) image_to_writer(writer, epoch, current_sample, current_recon, current_im) loader = Loader('/home/mehdi/Codes/ML3/WorldModels/dataset', 10000) writer = initialize_writer() args = get_args() vae = VAE() small_init(vae) train(vae, loader, args, writer)
def main(): if (FLAGS.save_path != ''): print(">>>>>>>>>>>>>>>FLAGS.save_path: ", FLAGS.save_path) from vae_model import VAE import EmoData as ED import cv2 import pickle batch_size = 10 vae_model = VAE((160, 240, 1), batch_size, FLAGS.num_au) pp = ED.image_pipeline.FACE_pipeline( histogram_normalization=True, grayscale=True, resize=True, rotation_range=3, width_shift_range=0.03, height_shift_range=0.03, zoom_range=0.03, random_flip=True, ) batch_size = 10 def get_test_latent(test_file_names, N_batch): file_names_batch = np.reshape( test_file_names[:N_batch * batch_size], [N_batch, batch_size]) z_arr = [] for file_path in file_names_batch: imgs = [] for filename in file_path: img = cv2.imread(filename) imgs.append(img) img_arr, pts, pts_raw = pp.batch_transform(imgs, preprocessing=True, augmentation=False) weights, z = vae_model.computeLatentVal( img_arr, FLAGS.vae_model, FLAGS.au_idx) z_arr.append(z) return np.concatenate(z_arr) test_subjects = os.listdir(FLAGS.testset_dir) test_subjects.sort() test_subjects = test_subjects[FLAGS.test_start_idx - 14:FLAGS.test_start_idx - 14 + FLAGS.test_num] print("test_subjects: ", test_subjects) test_z_arr = [] out = open(FLAGS.save_path + "testset_z_arr.pkl", 'wb') for test_subject in test_subjects: data = pickle.load(open(FLAGS.testset_dir + test_subject, "rb"), encoding='latin1') N_batch = int(len(data['test_file_names']) / batch_size) test_file_names = data['test_file_names'][:N_batch * batch_size] print( test_subject.split(".")[0], " original total len:", len(data['test_file_names'])) print( test_subject.split(".")[0], " rounded down total len:", len(test_file_names)) pickle.dump( {test_subject: get_test_latent(test_file_names, N_batch)}, out, protocol=2) # test_z_arr.append() # print("test_z_arr : ", test_z_arr) # print("test_z_arr size: ", np.array(test_z_arr).shape) out.close() else: data_generator = DataGenerator(FLAGS.update_batch_size * 2, FLAGS.meta_batch_size) data_generator.make_data_tensor()
["truedist" if is_truedist else "embdist"] + ["2K" if is_smalldata else "38K"] + ["binaryreward" if is_binaryreward else ""]) kwargs = {} if not is_truedist: assert is_image # n_trajs = args.n_trajs with open(test_data, 'rb') as f: test_tasks = pkl.load(f) raw_transitions = np.load(transition_file) if not os.path.exists(save_path): os.makedirs(save_path) if is_image: model = VAE(image_channels=3, z_dim=10).cuda() model.load_state_dict(torch.load(embedding_params)) kwargs['model'] = model transform = transforms.Compose([ transforms.Resize(64), transforms.CenterCrop(64), transforms.ToTensor(), # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) """ Set up the environment """ env = BlockEnv() env.reset() env.viewer_setup()
def make_data_tensor(self, train=True): if train: folders = self.metatrain_character_folders print("train folders: ", folders) else: folders = self.metatest_character_folders print("test folders: ", folders) # make list of files print('Generating filenames') inputa_files = [] inputb_files = [] labelas = [] labelbs = [] # To have totally different inputa and inputb, they should be sampled at the same time and then splitted. for sub_folder in folders: # 쓰일 task수만큼만 경로 만든다. 이 task들이 iteration동안 어차피 반복될거니까 # random.shuffle(sampled_character_folders) labels_and_images = get_images(sub_folder, range(self.num_classes), FLAGS.kshot_seed, nb_samples=self.num_samples_per_class, shuffle=True) # make sure the above isn't randomized order labels = [li[0] for li in labels_and_images] # 0 0 1 1 = off off on on filenames = [li[1] for li in labels_and_images] # Split data into a/b k = int(self.num_samples_per_class / 2) # = FLAGS.update_batch_size filenames = np.array(filenames).reshape(self.num_classes, self.num_samples_per_class) for files_per_class in filenames: for i in range(k): inputa_files.append(files_per_class[2*i]) inputb_files.append(files_per_class[2*i+1]) labels = np.array(labels).reshape(self.num_classes, self.num_samples_per_class) for labels_per_class in labels: for i in range(k): labelas.append(labels_per_class[2*i]) labelbs.append(labels_per_class[2*i+1]) print(">>>> inputa_files: ", inputa_files) print('----------------------------------------------------------------------------') print(">>>> inputb_files: ", inputb_files) print(">>>>>>>>>>>>>>>>> vae_model: ", FLAGS.vae_model) print(">>>>>>>>>>>>>>>>>> random seed for kshot: ", FLAGS.kshot_seed) print(">>>>>>>>>>>>>>>>>> random seed for weight: ", FLAGS.weight_seed) ################################################################################# # inputa_files has (n*k * num_of_task) files. # make it to batch of which size is (n*k) : thus, the total number of batch = num_of_task batch_size = int(self.num_classes * FLAGS.update_batch_size) N_batch = num_of_task = int(len(inputa_files) / batch_size) # len(inputa_files)/nk = num of task vae_model = VAE((self.img_size[0], self.img_size[1], 1), batch_size, FLAGS.num_au) def latent_feature(file_names): file_names_batch = np.reshape(file_names, [N_batch, batch_size]) z_arr = [] for file_bath in file_names_batch: imgs = [] for filename in file_bath: img = cv2.imread(filename) imgs.append(img) pp = ED.image_pipeline2.FACE_pipeline( histogram_normalization=True, grayscale=True, resize=True, rotation_range=3, width_shift_range=0.03, height_shift_range=0.03, zoom_range=0.03, random_flip=True, ) img_arr, pts, pts_raw = pp.batch_transform(imgs, preprocessing=True, augmentation=False) weights, z = vae_model.computeLatentVal(img_arr, FLAGS.vae_model, FLAGS.au_idx) z_arr.append(z) return np.concatenate(z_arr), weights inputa_latent_feat, self.pred_weights = latent_feature(inputa_files) inputb_latent_feat, self.pred_weights = latent_feature(inputb_files) ################################################################################# import pickle pp = ED.image_pipeline.FACE_pipeline( histogram_normalization=True, grayscale=True, resize=True, rotation_range=3, width_shift_range=0.03, height_shift_range=0.03, zoom_range=0.03, random_flip=True, ) batch_size = 10 def get_test_latent(test_file_names, N_batch): file_names_batch = np.reshape(test_file_names[:N_batch * batch_size], [N_batch, batch_size]) z_arr = [] for file_path in file_names_batch: imgs = [] for filename in file_path: img = cv2.imread(filename) imgs.append(img) img_arr, pts, pts_raw = pp.batch_transform(imgs, preprocessing=True, augmentation=False) weights, z = vae_model.computeLatentVal(img_arr, FLAGS.vae_model, FLAGS.au_idx) z_arr.append(z) return np.concatenate(z_arr) test_subjects = os.listdir(FLAGS.testset_dir) test_subjects.sort() test_subjects = test_subjects[FLAGS.test_start_idx - 14:FLAGS.test_start_idx - 14 + FLAGS.test_num] print("test_subjects: ", test_subjects) test_z_arr = [] for test_subject in test_subjects: data = pickle.load(open(FLAGS.testset_dir + test_subject, "rb"), encoding='latin1') N_batch = int(len(data['test_file_names']) / batch_size) test_file_names = data['test_file_names'][:N_batch * batch_size] print(test_subject.split(".")[0], " original total len:", len(data['test_file_names'])) print(test_subject.split(".")[0], " rounded down total len:", len(test_file_names)) test_z_arr.append(get_test_latent(test_file_names, N_batch)) print("test_z_arr : ", test_z_arr) print("test_z_arr size: ", np.array(test_z_arr).shape) ################################################################################# def get_distance(z_arr): size = z_arr.shape[0] distance_mat = np.zeros((size,size)) for i in range(size): for j in range(size): distance_mat[i,j]= np.linalg.norm(z_arr[i]-z_arr[j]) return distance_mat def get_distance_from_test(z_arr, z_arr2): distance_mat = np.zeros((z_arr.shape[0], len(z_arr2))) print("z_arr len in get_distance_from_test: ", z_arr.shape[0]) for i in range(z_arr.shape[0]): print(i) for j in range(len(z_arr2)): distance = 0 for k in range(z_arr2[j].shape[0]): distance += np.linalg.norm(z_arr[i] - z_arr2[j][k]) distance_mat[i, j] = distance / z_arr2[j].shape[0] return distance_mat btw_a_a = get_distance(inputa_latent_feat) print("btw_a_a.shape: ", btw_a_a.shape) btw_b_b = get_distance(inputb_latent_feat) print("btw_b_b.shape: ", btw_b_b.shape) # alldata=np.array(np.array(inputa_latent_feat.tolist().extend(inputa_latent_feat.tolist()))) # btw_all = get_distance(alldata) # test_file_names, y_lab btw_a_te = get_distance_from_test(inputa_latent_feat, test_z_arr) print("btw_a_te.shape: ", btw_a_te.shape) btw_b_te = get_distance_from_test(inputb_latent_feat, test_z_arr) print("btw_b_te.shape: ", btw_b_te.shape) save_path = "./logs/" + FLAGS.au + "/kshot/seed" + str(FLAGS.kshot_seed) + "/distance/" if not os.path.exists(save_path): os.makedirs(save_path) out = open(save_path + str(FLAGS.update_batch_size) + 'shot.pkl', 'wb') # pickle.dump({'btw_tr_tr': btw_tr_tr, 'btw_te_te': btw_te_te, 'btw_all': btw_all, 'a':inputa_files, 'b':inputb_files}, out, protocol=2) pickle.dump( {'btw_a_a': btw_a_a, 'btw_a_te': btw_a_te, 'a': inputa_files, 'btw_b_b': btw_b_b, 'btw_b_te': btw_b_te, 'a': inputb_files, }, out, protocol=2) out.close()
if args.cuda: torch.cuda.manual_seed(args.seed) kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} train_loader = torch.utils.data.DataLoader(TransientObjectLoader( args.data, train=True, transform=transforms.ToTensor()), batch_size=args.batch_size, shuffle=True, **kwargs) test_loader = torch.utils.data.DataLoader(TransientObjectLoader( args.data, train=False, transform=transforms.ToTensor()), batch_size=args.batch_size, shuffle=True, **kwargs) model = VAE() load_ext = False if osp.exists(args.save): with open(args.save, 'rb') as f: state_dict = torch.load(f) discard = [x for x in state_dict if x.startswith('fc1')] state = model.state_dict() state.update(state_dict) try: model.load_state_dict(state) except Exception: for key in discard: state_dict.pop(key) state = model.state_dict() state.update(state_dict) model.load_state_dict(state)
kld_history = np.load('VAE/kld_loss.npy') mse_history = np.load('VAE/mse_loss.npy') # print(kld_history) # print(mse_history) plt.figure(figsize=(40, 10)) plt.subplot(121) plt.title('KLD loss') plt.plot(kld_history) plt.subplot(122) plt.title('MSE loss') plt.plot(mse_history) plt.savefig(output_path+'/fig1_2.jpg') model = VAE(64, 1e-5) model.load_state_dict(torch.load('VAE/vae_state_model.pth')) if cuda: model = model.cuda() # print(model) data_for_tsne = [] label_for_tsne = [] mse = 0 for (data, label) in test_dataloader: # print(data.size()) # if cuda: # data = data.cuda() # data = Variable(data.cuda()) # recon_img, mu, logvar = model(data) # loss = model.loss_function(data, recon_img, mu, logvar)
help='how many training processes to use (default: 4)') parser.add_argument('--num-steps', type=int, default=20, help='number of forward steps in A3C (default: 20)') # parser.add_argument( '--env-name', default='Pong-ram-v0', help='environment to train on (default: PongDeterministic-v4)') return parser.parse_args() action_map = {0: 2, 1: 3} vae = VAE() vae.load_state_dict( torch.load(VAE_MODEL_PATH, map_location=lambda storage, loc: storage)) def train(rank, args, shared_model, optimizer, counter, lock): env = gym.make(args.env_name) env._max_episode_steps = 100000 env.seed(args.seed + rank) torch.manual_seed(args.seed + rank) model = RNN_vae(2, action_map) model.train() state = env.reset() done = True episode_length = 0
import pickle import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim import numpy as np import gym from vae_model import VAE vae = VAE() vae.load_state_dict(torch.load('trained_models/VAE_1')) vae.eval() def save_data(code, ac, done): class Env: def __init__(self): self.env = gym.make('CarRacing-v0') def reset(self):