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
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
while True:
# Sync with the shared model
# 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 = []
# Low Pass Filter
alpha = 0.9
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)
load_ext = True
if args.cuda:
model.cuda()
reconstruction_function = nn.MSELoss()
reconstruction_function.size_average = False
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
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 = [
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)
# mse += torch.sum(model.latest_loss()[0])
["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()
"""
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):