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):
