model.load_state_dict(state_dict)
else:
model.embedding.weight.data.copy_(vocab.vectors)
model = to_device(model)
print(model)
train(model, datasets, args)
if args.save_model is not None:
torch.save(model.state_dict(), args.save_model)
if args.n_generated > 0:
model.eval()
samples, z, y_onehot = model.inference(n=args.n_generated)
intent = y_onehot.data.max(1)[1].cpu().numpy()
delexicalised = idx2word(samples, i2w=i2w, pad_idx=pad_idx)
if args.input_type == 'delexicalised':
labelling, utterance = surface_realisation(samples,
i2w=i2w,
pad_idx=pad_idx)
print('----------GENERATED----------')
for i in range(args.n_generated):
print('Intent : ', i2int[intent[i]])
print('Delexicalised : ', delexicalised[i])
if args.input_type == 'delexicalised':
for i in range(args.n_generated):
print('Lexicalised : ', utterance[i] + '\n')
run['generated'] = utterance
def main(args):
# Image preprocessing
transform = transforms.Compose([
transforms.Resize((224, 224), Image.LANCZOS),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
# Load vocabulary wrapper
with open(args.vocab_path, 'rb') as f:
vocab = pickle.load(f)
pad_idx = vocab.word2idx['<pad>']
sos_idx = vocab.word2idx['<start>']
eos_idx = vocab.word2idx['<end>']
unk_idx = vocab.word2idx['<unk>']
# Build the models
model = CVAE(
vocab_size=len(vocab),
embedding_size=args.embedding_size,
rnn_type=args.rnn_type,
hidden_size=args.hidden_size,
word_dropout=args.word_dropout,
embedding_dropout=args.embedding_dropout,
latent_size=args.latent_size,
max_sequence_length=args.max_sequence_length,
num_layers=args.num_layers,
bidirectional=args.bidirectional,
pad_idx=pad_idx,
sos_idx=sos_idx,
eos_idx=eos_idx,
unk_idx=unk_idx
)
if not os.path.exists(args.load_checkpoint):
raise FileNotFoundError(args.load_checkpoint)
model.load_state_dict(torch.load(args.load_checkpoint))
print("Model loaded from {}".format(args.load_checkpoint))
model.to(device)
model.eval()
# Build data loader
train_data_loader, valid_data_loader = get_loader(args.train_image_dir, args.val_image_dir,
args.train_caption_path, args.val_caption_path, vocab,
args.batch_size,
shuffle=True, num_workers=args.num_workers)
f1 = open('{}/results/generated_captions.txt'.format(dataset_root_dir), 'w')
f2 = open('{}/results/ground_truth_captions.txt'.format(dataset_root_dir), 'w')
for i, (images, captions, lengths) in enumerate(valid_data_loader):
images = images.to(device)
sampled_ids, z = model.inference(n=args.batch_size, c=images)
sampled_ids_batches = sampled_ids.cpu().numpy() # (batch_size, max_seq_length)
captions = captions.cpu().numpy()
# Convert word_ids to words
for j, sampled_ids in enumerate(sampled_ids_batches):
sampled_caption = []
for word_id in sampled_ids:
word = vocab.idx2word[word_id]
sampled_caption.append(word)
if word == '<end>':
break
generated_sentence = ' '.join(sampled_caption)
generated_sentence = generated_sentence.rstrip()
generated_sentence = generated_sentence.replace("\n", "")
generated_sentence = "{0}\n".format(generated_sentence)
if j == 0:
print("RE: {}".format(generated_sentence))
f1.write(generated_sentence)
for g, ground_truth_ids in enumerate(captions):
ground_truth_caption = []
for word_id in ground_truth_ids:
word = vocab.idx2word[word_id]
ground_truth_caption.append(word)
if word == '<end>':
break
ground_truth_sentence = ' '.join(ground_truth_caption)
ground_truth_sentence = ground_truth_sentence.rstrip()
ground_truth_sentence = ground_truth_sentence.replace("\n", "")
ground_truth_sentence = "{0}\n".format(ground_truth_sentence)
if g == 0:
print("GT: {}".format(ground_truth_sentence))
f2.write(ground_truth_sentence)
if i % 10 == 0:
print("This is the {0}th batch".format(i))
f1.close()
f2.close()
class CVAEInterface():
def __init__(self, run_id=1, output_path="", env_path_root=""):
super().__init__()
self.cvae = CVAE(run_id=run_id)
self.device = torch.device('cuda' if CUDA_AVAILABLE else 'cpu')
self.output_path = output_path
self.env_path_root = env_path_root
if self.output_path is not None:
if os.path.exists(self.output_path):
shutil.rmtree(self.output_path)
os.mkdir(self.output_path)
def load_dataset(self, dataset_root, data_type="arm", mode="train"):
assert (data_type == "both" or data_type == "arm"
or data_type == "base")
assert (mode == "train" or mode == "test")
# Should show different count and path for different modes
print("Loading {} dataset for mode : {}, path : {}".format(
data_type, mode, dataset_root))
self.data_type = data_type
paths_dataset = PathsDataset(type="FULL_STATE")
c_test_dataset = PathsDataset(type="CONDITION_ONLY")
env_dir_paths = os.listdir(dataset_root)
# Get all C vars to test sample generation on each
all_condition_vars = []
for env_dir_index in filter(lambda f: f[0].isdigit(), env_dir_paths):
env_paths_file = os.path.join(dataset_root, env_dir_index,
"data_{}.txt".format(data_type))
env_paths = np.loadtxt(env_paths_file)
# 4 to 16
if IGNORE_START:
start = env_paths[:, X_DIM:2 * X_DIM]
samples = env_paths[:, :X_DIM]
euc_dist = np.linalg.norm(start - samples, axis=1)
far_from_start = np.where(euc_dist > 5.0)
print(far_from_start)
env_paths = env_paths[far_from_start[0], :]
condition_vars = env_paths[:, 2 * X_DIM:2 * X_DIM + C_DIM]
else:
if mode == "train":
# Testing, less points near start to reduce them in sampled output
start = env_paths[:, X_DIM:X_DIM + POINT_DIM]
samples = env_paths[:, :X_DIM]
euc_dist = np.linalg.norm(start - samples, axis=1)
far_from_start = np.where(euc_dist > 2.0)
# print(far_from_start)
env_paths = env_paths[far_from_start[0], :]
condition_vars = env_paths[:, X_DIM:X_DIM + C_DIM]
# print(env_paths.shape)
# Stuff for train dataloader
# Take only required elements
# env_paths = env_paths[:, :X_DIM + C_DIM]
env_paths = np.hstack((env_paths[:, :X_DIM], condition_vars))
# Uniquify to remove duplicates
env_paths = np.unique(env_paths, axis=0)
env_index = np.empty((env_paths.shape[0], 1))
env_index.fill(env_dir_index)
data = np.hstack((env_index, env_paths))
paths_dataset.add_env_paths(data.tolist())
# Stuff for test dataloader
env_index = np.empty((condition_vars.shape[0], 1))
env_index.fill(env_dir_index)
data = np.hstack((env_index, condition_vars))
all_condition_vars += data.tolist()
print("Added {} states from {} environment".format(
env_paths.shape[0], env_dir_index))
dataloader = DataLoader(paths_dataset,
batch_size=TRAIN_BATCH_SIZE,
shuffle=True)
if data_type != "both":
# Depending on which dataset is being loaded, set the right variables
if mode == "train":
self.train_dataloader = dataloader
self.train_paths_dataset = paths_dataset
elif mode == "test":
self.test_condition_vars = np.unique(all_condition_vars,
axis=0)
print("Unique test conditions count : {}".format(
self.test_condition_vars.shape[0]))
# Tile condition variables to predict given number of samples for x
all_condition_vars_tile = np.repeat(self.test_condition_vars,
TEST_SAMPLES, 0)
c_test_dataset.add_env_paths(all_condition_vars_tile.tolist())
c_test_dataloader = DataLoader(c_test_dataset,
batch_size=TEST_BATCH_SIZE,
shuffle=False)
self.test_dataloader = c_test_dataloader
else:
arm_test_dataset = PathsDataset(type="CONDITION_ONLY")
base_test_dataset = PathsDataset(type="CONDITION_ONLY")
all_condition_vars = np.array(all_condition_vars)
self.test_condition_vars = np.delete(all_condition_vars, [4, 5],
axis=1)
self.test_condition_vars = np.unique(self.test_condition_vars,
axis=0)
print("Unique test conditions count : {}".format(
self.test_condition_vars.shape[0]))
# print(self.test_condition_vars)
arm_condition_vars = np.insert(self.test_condition_vars,
2 * POINT_DIM,
1,
axis=1)
arm_condition_vars = np.insert(arm_condition_vars,
2 * POINT_DIM,
0,
axis=1)
arm_condition_vars = np.repeat(arm_condition_vars, TEST_SAMPLES, 0)
arm_test_dataset.add_env_paths(arm_condition_vars.tolist())
arm_test_dataloader = DataLoader(arm_test_dataset,
batch_size=TEST_BATCH_SIZE,
shuffle=False)
base_condition_vars = np.insert(self.test_condition_vars,
2 * POINT_DIM,
0,
axis=1)
base_condition_vars = np.insert(base_condition_vars,
2 * POINT_DIM,
1,
axis=1)
base_condition_vars = np.repeat(base_condition_vars, TEST_SAMPLES,
0)
base_test_dataset.add_env_paths(base_condition_vars.tolist())
base_test_dataloader = DataLoader(base_test_dataset,
batch_size=TEST_BATCH_SIZE,
shuffle=False)
if mode == "train":
self.train_dataloader = dataloader
elif mode == "test":
self.arm_test_dataloader = arm_test_dataloader
self.base_test_dataloader = base_test_dataloader
def visualize_train_data(self, num_conditions=1):
# Pick a random condition
# Find all states for that condition
# Plot them
print("Plotting input data for {} random conditions".format(
num_conditions))
all_input_paths = np.array(self.train_paths_dataset.paths)[:, 1:]
env_ids = np.array(self.train_paths_dataset.paths)[:, :1]
# print(all_input_paths[0,:])
for c_i in range(num_conditions):
rand_index = np.random.randint(0, all_input_paths.shape[0])
condition = all_input_paths[rand_index, 2:]
env_id = env_ids[rand_index, 0]
# print(condition)
# condition_samples = np.argwhere(all_input_paths[:,2:] == condition)
# indices = np.where(all_input_paths[:,2:] == condition)
# Find all samples corresponding to this condition
indices = np.where(
np.isin(all_input_paths[:, 2:], condition).all(axis=1))[0]
# print(indices)
x = all_input_paths[indices, :2]
fig = self.plot(x, condition, env_id=env_id)
self.cvae.tboard.add_figure('train_data/condition_{}'.format(c_i),
fig, 0)
# print(all_input_paths[indices,:])
self.cvae.tboard.flush()
def visualize_map(self, env_id):
path = "{}/{}.txt".format(self.env_path_root, int(env_id))
plt.title('Environment - {}'.format(env_id))
with open(path, "r") as f:
line = f.readline()
while line:
line = line.split(" ")
# print(line)
if "wall" in line[0] or "table" in line[0]:
x = float(line[1])
y = float(line[2])
l = float(line[4])
b = float(line[5])
rect = Rectangle((x - l / 2, y - b / 2), l, b)
plt.gca().add_patch(rect)
line = f.readline()
plt.draw()
def plot(self, x, c, env_id=None, suffix=0, write_file=False, show=False):
'''
Plot samples and environment - from train input or predicted output
'''
# print(c)
if IGNORE_START:
goal = c[0:2]
else:
start = c[0:2]
goal = c[2:4]
# For given conditional, plot the samples
fig1 = plt.figure(figsize=(10, 6), dpi=80)
# ax1 = fig1.add_subplot(111, aspect='equal')
plt.scatter(x[:, 0], x[:, 1], color="green", s=70, alpha=0.1)
if IGNORE_START == False:
plt.scatter(start[0], start[1], color="blue", s=70, alpha=0.6)
plt.scatter(goal[0], goal[1], color="red", s=70, alpha=0.6)
if env_id is not None:
self.visualize_map(env_id)
# wall_locs = c[4:]
# i = 0
# while i < wall_locs.shape[0]:
# plt.scatter(wall_locs[i], wall_locs[i+1], color="green", s=70, alpha=0.6)
# i = i + 2
plt.xlabel('x')
plt.ylabel('y')
plt.xlim(0, X_MAX)
plt.ylim(0, Y_MAX)
if write_file:
plt.savefig('{}/gen_points_fig_{}.png'.format(
self.output_path, suffix))
np.savetxt('{}/gen_points_{}.txt'.format(self.output_path, suffix),
x,
fmt="%.2f",
delimiter=',')
np.savetxt('{}/start_goal_{}.txt'.format(self.output_path, suffix),
np.vstack((start, goal)),
fmt="%.2f",
delimiter=',')
if show:
plt.show()
# plt.close(fig1)
return fig1
def load_saved_cvae(self, decoder_path):
print("Loading saved CVAE")
self.cvae.load_decoder(decoder_path)
# base_cvae = CVAE(run_id=run_id)
# base_decoder_path = 'experiments/cvae/base/decoder-final.pkl'
# base_cvae.load_decoder(base_decoder_path)
# for iteration, batch in enumerate(dataloader):
def test_single(self,
env_id,
sample_size=1000,
c_test=None,
visualize=True):
self.cvae.eval()
c_test_gpu = torch.from_numpy(c_test).float().to(self.device)
c_test_gpu = torch.unsqueeze(c_test_gpu, dim=0)
x_test = self.cvae.inference(sample_size=sample_size, c=c_test_gpu)
x_test = x_test.detach().cpu().numpy()
if visualize:
self.plot(x_test,
c_test,
env_id=env_id,
show=False,
write_file=True,
suffix=0)
return x_test
def test(self, epoch, dataloader, write_file=False, suffix=""):
x_test_predicted = []
self.cvae.eval()
for iteration, batch in enumerate(dataloader):
# print(batch)
c_test_data = batch['condition'].float().to(self.device)
# print(c_test_data[0,:])
x_test = self.cvae.batch_inference(c=c_test_data)
x_test_predicted += x_test.detach().cpu().numpy().tolist()
# print(x_test.shape)
if iteration % LOG_INTERVAL == 0 or iteration == len(
dataloader) - 1:
print(
"Test Epoch {:02d}/{:02d} Batch {:04d}/{:d}, Iteration {}".
format(epoch, num_epochs, iteration,
len(dataloader) - 1, iteration))
x_test_predicted = np.array(x_test_predicted)
# print(x_test_predicted.shape)
# Draw plot for each unique condition
for c_i in range(self.test_condition_vars.shape[0]):
x_test = x_test_predicted[c_i * TEST_SAMPLES:(c_i + 1) *
TEST_SAMPLES]
# Fine because c_test is used only for plotting, we dont need arm/base label here
c_test = self.test_condition_vars[c_i, 1:]
env_id = self.test_condition_vars[c_i, 0]
# print(self.test_condition_vars[c_i,:])
fig = self.plot(x_test,
c_test,
env_id=env_id,
suffix=c_i,
write_file=write_file)
self.cvae.tboard.add_figure(
'test_epoch_{}/condition_{}_{}'.format(epoch, c_i, suffix),
fig, 0)
if c_i % LOG_INTERVAL == 0:
print("Plotting condition : {}".format(c_i))
self.cvae.tboard.flush()
# for c_i in range(c_test_data.shape[0]):
# c_test = c_test_data[c_i,:]
# c_test_gpu = torch.from_numpy(c_test).float().to(device)
# x_test = cvae_model.inference(n=TEST_SAMPLES, c=c_test_gpu)
# x_test = x_test.detach().cpu().numpy()
# fig = plot(x_test, c_test)
# cvae_model.tboard.add_figure('test_epoch_{}/condition_{}'.format(epoch, c_i), fig, 0)
# if c_i % 50 == 0:
# print("Epoch : {}, Testing condition count : {} ".format(epoch, c_i))
def train(self,
run_id=1,
num_epochs=1,
initial_learning_rate=0.001,
weight_decay=0.0001):
optimizer = torch.optim.Adam(self.cvae.parameters(),
lr=initial_learning_rate,
weight_decay=weight_decay)
for epoch in range(num_epochs):
for iteration, batch in enumerate(self.train_dataloader):
# print(batch['condition'][0,:])
self.cvae.train()
x = batch['state'].float().to(self.device)
c = batch['condition'].float().to(self.device)
recon_x, mean, log_var, z = self.cvae(x, c)
# print(recon_x.shape)
loss = self.cvae.loss_fn(recon_x, x, mean, log_var)
optimizer.zero_grad()
loss.backward()
optimizer.step()
counter = epoch * len(self.train_dataloader) + iteration
if iteration % LOG_INTERVAL == 0 or iteration == len(
self.train_dataloader) - 1:
print(
"Train Epoch {:02d}/{:02d} Batch {:04d}/{:d}, Iteration {}, Loss {:9.4f}"
.format(epoch, num_epochs, iteration,
len(self.train_dataloader) - 1, counter,
loss.item()))
self.cvae.tboard.add_scalar('train/loss', loss.item(),
counter)
# cvae.eval()
# c_test = c[0,:]
# x_test = cvae.inference(n=TEST_SAMPLES, c=c_test)
# x_test = x_test.detach().cpu().numpy()
# fig = plot(x_test, c_test)
# cvae.tboard.add_figure('test/samples', fig, counter)
if epoch % TEST_INTERVAL == 0 or epoch == num_epochs - 1:
# Test CVAE for all c by drawing samples
if self.data_type != "both":
self.test(epoch, self.test_dataloader)
else:
self.test(epoch, self.arm_test_dataloader, suffix="arm")
self.test(epoch, self.base_test_dataloader, suffix="base")
if epoch % SAVE_INTERVAL == 0 and epoch > 0:
self.cvae.save_model_weights(counter)
self.cvae.save_model_weights('final')
