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