coords = np.zeros((ssp_pred.shape[0]*ssp_pred.shape[1], 2))
print("Computing predicted locations and true locations")
# Using all data, one chunk at a time
for ri in range(args.rollout_length):
if args.spatial_encoding == '2d':
# copying 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_pred.detach().numpy()[ri, :, :]
# copying 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_outputs.detach().numpy()[:, ri, :]
else:
# computing 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[ri, :, :],
heatmap_vectors, xs, ys
)
# computing 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, ri, :],
heatmap_vectors, xs, ys
)
rmse = np.sqrt((np.linalg.norm(predictions - coords, axis=1)**2).mean())
print(args.spatial_encoding)
print(rmse)
print("")
if args.save_file:
np.savez(args.save_file, rmse=rmse)
# print("lstm_outputs.shape", lstm_outputs.shape)
# predictions = np.zeros((ssp_pred.shape[1]*ssp_pred.shape[2], 2))
# coords = np.zeros((ssp_pred.shape[1]*ssp_pred.shape[2], 2))
predictions = np.zeros((ssp_pred.shape[0] * ssp_pred.shape[1], 2))
coords = np.zeros((ssp_pred.shape[0] * ssp_pred.shape[1], 2))
print("Computing predicted locations and true locations")
# Using all data, one chunk at a time
for ri in range(rollout_length):
if args.encoding == 'ssp':
# computing 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) *
ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[ri, :, :],
heatmap_vectors, xs, ys)
# computing 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) *
ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, ri, :], heatmap_vectors,
xs, ys)
elif args.encoding == '2d':
# copying 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) *
ssp_pred.shape[1], :] = ssp_pred.detach().numpy()[
ri, :, :]
# copying 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) *
cosine_loss.data.item() * c_f, epoch)
# print("ssp_pred.shape", ssp_pred.shape)
# print("ssp_outputs.shape", ssp_outputs.shape)
# Just use start and end location to save on memory and computation
predictions_start = np.zeros((ssp_pred.shape[1], 2))
coords_start = np.zeros((ssp_pred.shape[1], 2))
predictions_end = np.zeros((ssp_pred.shape[1], 2))
coords_end = np.zeros((ssp_pred.shape[1], 2))
if args.encoding == 'ssp':
print("computing prediction locations")
predictions_start[:, :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[0, :, :], heatmap_vectors, xs,
ys)
predictions_end[:, :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[-1, :, :], heatmap_vectors, xs,
ys)
print("computing ground truth locations")
coords_start[:, :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, 0, :], heatmap_vectors, xs,
ys)
coords_end[:, :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, -1, :], heatmap_vectors,
xs, ys)
elif args.encoding == '2d':
print("copying prediction locations")
predictions_start[:, :] = ssp_pred.detach().numpy()[0, :, :]
predictions_end[:, :] = ssp_pred.detach().numpy()[-1, :, :]
encodings = np.zeros((n_samples, args.dim))
for i in range(n_samples):
encodings[i, :] = encoding_func(np.array([rand_pos[i, 0], rand_pos[i, 1]]))
# encodings[i, :] /= np.linalg.norm(encodings[i, :])
encodings += np.random.normal(loc=0,
scale=args.noise,
size=(n_samples, args.dim))
# encodings /= encodings.sum(axis=1)[:, np.newaxis]
predictions = ssp_to_loc_v(
# flat_heatmap_vectors,
encodings,
heatmap_vectors,
xs,
ys)
print(predictions)
coords = predictions.copy()
fig_pred, ax_pred = plt.subplots()
print("plotting predicted locations")
plot_predictions_v(predictions / args.ssp_scaling,
coords / args.ssp_scaling,
ax_pred,
min_val=args.limit_low,
max_val=args.limit_high)
# batch_data = viz_eval[out_p_filt][:, -1, :]
batch_data = viz_eval[out_p_filt][:, 10:, :].mean(axis=1)
true_ssps = vis_output[bi, :, :]
print('pred.shape', batch_data.shape)
print('true_ssps.shape', true_ssps.shape)
wall_overlay = np.sum(true_ssps, axis=1) == 0
print('wall_overlay.shape', wall_overlay.shape)
hmv = get_encoding_heatmap_vectors(xs, ys, args.dim, encoding_func, normalize=False)
predictions = np.zeros((res * res, 2))
# computing 'predicted' coordinates, where the agent thinks it is
predictions[:, :] = ssp_to_loc_v(
batch_data,
hmv, xs, ys
)
plot_predictions_v(
predictions=predictions, coords=coords,
ax=ax[bi],
min_val=limit_low,
max_val=limit_high,
fixed_axes=True,
)
plt.show()
# # For each neuron, contains the average activity at each spatial bin
# # Computing for both ground truth and predicted location
# rate_maps_pred = np.zeros((model.lstm_hidden_size, len(xs), len(ys)))
# rate_maps_truth = np.zeros((model.lstm_hidden_size, len(xs), len(ys)))
print("Computing predicted locations and true locations")
# Using all data, one chunk at a time
for ri in range(rollout_length):
# trim out head direction info if that was included by only looking up to args.encoding_dim
# computing 'predicted' coordinates, where the agent thinks it is
pred = ssp_pred.detach().numpy()[ri, :, :args.dim]
# pred = pred / pred.sum(axis=1)[:, np.newaxis]
predictions[ri * ssp_pred.shape[1]:(ri + 1) *
ssp_pred.shape[1], :] = ssp_to_loc_v(
pred, heatmap_vectors, xs, ys)
# computing 'ground truth' coordinates, where the agent should be
coord = ssp_outputs.detach().numpy()[:, ri, :args.dim]
# coord = coord / coord.sum(axis=1)[:, np.newaxis]
coords[ri * ssp_pred.shape[1]:(ri + 1) *
ssp_pred.shape[1], :] = ssp_to_loc_v(coord, heatmap_vectors, xs,
ys)
# reshaping activations and converting to numpy array
lstm_activations[ri * ssp_pred.shape[1]:(ri + 1) *
ssp_pred.shape[1], :] = lstm_outputs.detach().numpy()[
ri, :, :]
dense_activations[
ri * ssp_pred.shape[1]:(ri + 1) *
ssp_pred.shape[1], :] = dense_outputs.detach().numpy()[ri, :, :]
ssp_pred, ssp_outputs.to(device),
torch.ones(ssp_pred.shape[0]).to(device))
mse_loss = mse_criterion(ssp_pred, ssp_outputs.to(device))
print("test mse loss", mse_loss.data.item())
print("test cosine loss", mse_loss.data.item())
# One prediction and ground truth coord for every element in the batch
# NOTE: this is assuming the eval set only has one giant batch
predictions = np.zeros((ssp_pred.shape[0], 2))
coords = np.zeros((ssp_pred.shape[0], 2))
if spatial_decoding_type == 'ssp':
print("computing prediction locations")
predictions[:, :] = ssp_to_loc_v(ssp_pred.detach().cpu().numpy()[:, :],
eval_heatmap_vectors, eval_xs,
eval_ys)
print("computing ground truth locations")
coords[:, :] = ssp_to_loc_v(ssp_outputs.detach().cpu().numpy()[:, :],
eval_heatmap_vectors, eval_xs, eval_ys)
elif spatial_decoding_type == '2d':
print("copying prediction locations")
predictions[:, :] = ssp_pred.detach().cpu().numpy()[:, :]
print("copying ground truth locations")
coords[:, :] = ssp_outputs.detach().cpu().numpy()[:, :]
elif spatial_decoding_type == '2d-normalized':
print("copying prediction locations")
predictions[:, :] = ssp_pred.detach().cpu().numpy()[:, :]
time_steps = clean_pred_over_time.shape[0]
print('time_steps', time_steps)
steps_per_item = 100
n_samples = int(time_steps / steps_per_item)
print('n_samples', n_samples)
clean_pred = np.zeros((n_samples, args.dim))
# clean_pred[:, :] = clean_pred_over_time[50::steps_per_item, :]
# take the mean from the middle of the estimate
for i in range(n_samples):
clean_pred[i, :] = clean_pred_over_time[i * steps_per_item +
10:i * steps_per_item +
100, :].mean(axis=0)
truth_coord_pred = np.zeros((n_samples, 2))
truth_coord_pred[:, :] = ssp_to_loc_v(clean_true[:n_samples, :],
heatmap_vectors, xs, ys)
error = np.zeros((n_samples, ))
for i in range(n_samples):
error[i] = np.linalg.norm(truth_coord_pred[i, :] - coords[i, :])
n_close = len(np.where(error < .5)[0])
print("truth fraction within 0.5 is {}".format(1.0 * n_close / n_samples))
print("true ssp cosine dist is {}".format(
np.diagonal(
cdist(clean_true[:n_samples, :],
clean_true[:n_samples, :],
metric='cosine')).mean()))
noisy_coord_pred = np.zeros((n_samples, 2))
noisy_coord_pred[:, :] = ssp_to_loc_v(noisy[:n_samples, :], heatmap_vectors,
xs, ys)
def run_eval(self, model, writer, epoch):
with torch.no_grad():
# Everything is in one batch, so this loop will only happen once
for i, data in enumerate(self.dataloader):
# sensor_inputs, map_ids, ssp_outputs = data
# sensors and map ID combined
combined_inputs, ssp_outputs = data
# ssp_pred = model(sensor_inputs, map_ids)
ssp_pred = model(combined_inputs)
cosine_loss = self.cosine_criterion(
ssp_pred, ssp_outputs, torch.ones(ssp_pred.shape[0]))
mse_loss = self.mse_criterion(ssp_pred, ssp_outputs)
print("test mse loss", mse_loss.data.item())
print("test cosine loss", mse_loss.data.item())
writer.add_scalar('test_mse_loss', mse_loss.data.item(), epoch)
writer.add_scalar('test_cosine_loss', cosine_loss.data.item(),
epoch)
# One prediction and ground truth coord for every element in the batch
# NOTE: this is assuming the eval set only has one giant batch
predictions = np.zeros((ssp_pred.shape[0], 2))
coords = np.zeros((ssp_pred.shape[0], 2))
if self.spatial_encoding == 'ssp':
print("computing prediction locations")
predictions[:, :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[:, :], self.heatmap_vectors,
self.xs, self.ys)
print("computing ground truth locations")
coords[:, :] = ssp_to_loc_v(ssp_outputs.detach().numpy()[:, :],
self.heatmap_vectors, self.xs,
self.ys)
elif self.spatial_encoding == '2d':
print("copying prediction locations")
predictions[:, :] = ssp_pred.detach().numpy()[:, :]
print("copying ground truth locations")
coords[:, :] = ssp_outputs.detach().numpy()[:, :]
fig_pred, ax_pred = plt.subplots()
fig_truth, ax_truth = plt.subplots()
print("plotting predicted locations")
plot_predictions_v(
# predictions / self.ssp_scaling,
# coords / self.ssp_scaling,
predictions,
coords,
ax_pred,
# min_val=0,
# max_val=2.2
min_val=self.xs[0],
max_val=self.xs[-1],
)
writer.add_figure("predictions", fig_pred, epoch)
# Only plot ground truth if epoch is 0
if epoch == 0:
print("plotting ground truth locations")
plot_predictions_v(
# coords / self.ssp_scaling,
# coords / self.ssp_scaling,
coords,
coords,
ax_truth,
# min_val=0,
# max_val=2.2
min_val=self.xs[0],
max_val=self.xs[-1],
)
writer.add_figure("ground truth", fig_truth, epoch)
def run_eval(self, model, writer, epoch):
with torch.no_grad():
# Everything is in one batch, so this loop will only happen once
for i, data in enumerate(self.dataloader):
combined_inputs, ssp_inputs, ssp_outputs = data
ssp_pred = model(combined_inputs, ssp_inputs)
# NOTE: need to permute axes of the targets here because the output is
# (sequence length, batch, units) instead of (batch, sequence_length, units)
# could also permute the outputs instead
# NOTE: for cosine loss the input needs to be flattened first
cosine_loss = self.cosine_criterion(
ssp_pred.reshape(ssp_pred.shape[0] * ssp_pred.shape[1],
ssp_pred.shape[2]),
ssp_outputs.permute(1, 0, 2).reshape(
ssp_pred.shape[0] * ssp_pred.shape[1],
ssp_pred.shape[2]),
torch.ones(ssp_pred.shape[0] * ssp_pred.shape[1]))
mse_loss = self.mse_criterion(ssp_pred,
ssp_outputs.permute(1, 0, 2))
print("test mse loss", mse_loss.data.item())
print("test cosine loss", mse_loss.data.item())
writer.add_scalar('test_mse_loss', mse_loss.data.item(), epoch)
writer.add_scalar('test_cosine_loss', cosine_loss.data.item(),
epoch)
# Just use start and end location to save on memory and computation
predictions_start = np.zeros((ssp_pred.shape[1], 2))
coords_start = np.zeros((ssp_pred.shape[1], 2))
predictions_end = np.zeros((ssp_pred.shape[1], 2))
coords_end = np.zeros((ssp_pred.shape[1], 2))
if self.spatial_encoding == 'ssp':
print("computing prediction locations")
predictions_start[:, :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[0, :, :], self.heatmap_vectors,
self.xs, self.ys)
predictions_end[:, :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[-1, :, :], self.heatmap_vectors,
self.xs, self.ys)
print("computing ground truth locations")
coords_start[:, :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, 0, :],
self.heatmap_vectors, self.xs, self.ys)
coords_end[:, :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, -1, :],
self.heatmap_vectors, self.xs, self.ys)
elif self.spatial_encoding == '2d':
print("copying prediction locations")
predictions_start[:, :] = ssp_pred.detach().numpy()[0, :, :]
predictions_end[:, :] = ssp_pred.detach().numpy()[-1, :, :]
print("copying ground truth locations")
coords_start[:, :] = ssp_outputs.detach().numpy()[:, 0, :]
coords_end[:, :] = ssp_outputs.detach().numpy()[:, -1, :]
fig_pred_start, ax_pred_start = plt.subplots()
fig_truth_start, ax_truth_start = plt.subplots()
fig_pred_end, ax_pred_end = plt.subplots()
fig_truth_end, ax_truth_end = plt.subplots()
print("plotting predicted locations")
plot_predictions_v(
predictions_start / self.ssp_scaling,
coords_start / self.ssp_scaling,
ax_pred_start,
min_val=self.xs[0],
max_val=self.xs[-1],
)
plot_predictions_v(
predictions_end / self.ssp_scaling,
coords_end / self.ssp_scaling,
ax_pred_end,
min_val=self.xs[0],
max_val=self.xs[-1],
)
writer.add_figure("predictions start", fig_pred_start, epoch)
writer.add_figure("predictions end", fig_pred_end, epoch)
# Only plotting ground truth if the epoch is 0
if epoch == 0:
print("plotting ground truth locations")
plot_predictions_v(
coords_start / self.ssp_scaling,
coords_start / self.ssp_scaling,
ax_truth_start,
min_val=self.xs[0],
max_val=self.xs[-1],
)
plot_predictions_v(
coords_end / self.ssp_scaling,
coords_end / self.ssp_scaling,
ax_truth_end,
min_val=self.xs[0],
max_val=self.xs[-1],
)
writer.add_figure("ground truth start", fig_truth_start, epoch)
writer.add_figure("ground truth end", fig_truth_end, epoch)
def run_and_gather_activations(
seed=13,
n_samples=1000,
dataset='../../lab/reproducing/data/path_integration_trajectories_logits_200t_15s_seed13.npz',
model_path='../output/ssp_path_integration/clipped/Mar22_15-24-10/ssp_path_integration_model.pt',
encoding='ssp',
rollout_length=100,
batch_size=10,
n_place_cells=256,
encoding_func=None, # added for frozen-learned encoding option
):
torch.manual_seed(seed)
np.random.seed(seed)
data = np.load(dataset)
x_axis_vec = data['x_axis_vec']
y_axis_vec = data['y_axis_vec']
pc_centers = data['pc_centers']
#pc_activations = data['pc_activations']
if encoding == 'ssp':
encoding_dim = 512
ssp_scaling = data['ssp_scaling']
elif encoding == '2d':
encoding_dim = 2
ssp_scaling = 1
elif encoding == 'pc':
dim = n_place_cells
ssp_scaling = 1
elif encoding == 'frozen-learned':
encoding_dim = 512
ssp_scaling = 1
elif encoding == 'pc-gauss' or encoding == 'pc-gauss-softmax':
encoding_dim = 512
ssp_scaling = 1
else:
raise NotImplementedError
limit_low = 0 * ssp_scaling
limit_high = 2.2 * ssp_scaling
res = 128 #256
xs = np.linspace(limit_low, limit_high, res)
ys = np.linspace(limit_low, limit_high, res)
if encoding == 'frozen-learned' or encoding == 'pc-gauss' or encoding == 'pc-gauss-softmax':
# encoding for every point in a 2D linspace, for approximating a readout
# FIXME: inefficient but will work for now
heatmap_vectors = np.zeros((len(xs), len(ys), 512))
for i, x in enumerate(xs):
for j, y in enumerate(ys):
heatmap_vectors[i, j, :] = encoding_func(
# batch dim
# np.array(
# [[x, y]]
# )
# no batch dim
np.array(
[x, y]
)
)
heatmap_vectors[i, j, :] /= np.linalg.norm(heatmap_vectors[i, j, :])
else:
# Used for visualization of test set performance using pos = ssp_to_loc(sp, heatmap_vectors, xs, ys)
heatmap_vectors = get_heatmap_vectors(xs, ys, x_axis_vec, y_axis_vec)
model = SSPPathIntegrationModel(unroll_length=rollout_length, sp_dim=encoding_dim)
model.load_state_dict(torch.load(model_path), strict=False)
trainloader, testloader = train_test_loaders(
data,
n_train_samples=n_samples,
n_test_samples=n_samples,
rollout_length=rollout_length,
batch_size=batch_size,
encoding=encoding,
encoding_func=encoding_func,
)
print("Testing")
with torch.no_grad():
# Everything is in one batch, so this loop will only happen once
for i, data in enumerate(testloader):
velocity_inputs, ssp_inputs, ssp_outputs = data
ssp_pred, lstm_outputs = model.forward_activations(velocity_inputs, ssp_inputs)
predictions = np.zeros((ssp_pred.shape[0]*ssp_pred.shape[1], 2))
coords = np.zeros((ssp_pred.shape[0]*ssp_pred.shape[1], 2))
activations = np.zeros((ssp_pred.shape[0]*ssp_pred.shape[1], model.lstm_hidden_size))
assert rollout_length == ssp_pred.shape[0]
# # For each neuron, contains the average activity at each spatial bin
# # Computing for both ground truth and predicted location
# rate_maps_pred = np.zeros((model.lstm_hidden_size, len(xs), len(ys)))
# rate_maps_truth = np.zeros((model.lstm_hidden_size, len(xs), len(ys)))
print("Computing predicted locations and true locations")
# Using all data, one chunk at a time
for ri in range(rollout_length):
if encoding == 'ssp':
# computing 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[ri, :, :],
heatmap_vectors, xs, ys
)
# computing 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, ri, :],
heatmap_vectors, xs, ys
)
elif encoding == '2d':
# copying 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_pred.detach().numpy()[ri, :, :]
# copying 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_outputs.detach().numpy()[:, ri, :]
elif encoding == 'pc':
# (quick hack is to just use the most activated place cell center)
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = pc_to_loc_v(
pc_activations=ssp_outputs.detach().numpy()[:, ri, :],
centers=pc_centers,
jitter=0.01,
)
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = pc_to_loc_v(
pc_activations=ssp_outputs.detach().numpy()[:, ri, :],
centers=pc_centers,
jitter=0.01,
)
elif encoding == 'frozen-learned' or encoding == 'pc-gauss' or encoding == 'pc-gauss-softmax':
# computing 'predicted' coordinates, where the agent thinks it is
pred = ssp_pred.detach().numpy()[ri, :, :]
pred = pred / pred.sum(axis=1)[:, np.newaxis]
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
pred,
heatmap_vectors, xs, ys
)
# computing 'ground truth' coordinates, where the agent should be
coord = ssp_outputs.detach().numpy()[:, ri, :]
coord = coord / coord.sum(axis=1)[:, np.newaxis]
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
coord,
heatmap_vectors, xs, ys
)
# reshaping activations and converting to numpy array
activations[ri*ssp_pred.shape[1]:(ri+1)*ssp_pred.shape[1], :] = lstm_outputs.detach().numpy()[ri, :, :]
return activations, predictions, coords
def run_and_gather_localization_activations(
seed=13,
n_samples=1000,
dataset='../../localization/data/localization_trajectories_5m_200t_250s_seed13.npz',
model_path='../../localization/output/ssp_trajectory_localization/May13_16-00-27/ssp_trajectory_localization_model.pt',
encoding='ssp',
rollout_length=100,
batch_size=10,
):
torch.manual_seed(seed)
np.random.seed(seed)
data = np.load(dataset)
x_axis_vec = data['x_axis_vec']
y_axis_vec = data['y_axis_vec']
ssp_scaling = data['ssp_scaling']
ssp_offset = data['ssp_offset']
# shape of coarse maps is (n_maps, env_size, env_size)
# some npz files had different naming, try both
try:
coarse_maps = data['coarse_maps']
except KeyError:
coarse_maps = data['coarse_mazes']
n_maps = coarse_maps.shape[0]
env_size = coarse_maps.shape[1]
# shape of dist_sensors is (n_maps, n_trajectories, n_steps, n_sensors)
n_sensors = data['dist_sensors'].shape[3]
# shape of ssps is (n_maps, n_trajectories, n_steps, dim)
dim = data['ssps'].shape[3]
limit_low = -ssp_offset * ssp_scaling
limit_high = (env_size - ssp_offset) * ssp_scaling
res = 256
xs = np.linspace(limit_low, limit_high, res)
ys = np.linspace(limit_low, limit_high, res)
# Used for visualization of test set performance using pos = ssp_to_loc(sp, heatmap_vectors, xs, ys)
heatmap_vectors = get_heatmap_vectors(xs, ys, x_axis_vec, y_axis_vec)
model = LocalizationModel(
input_size=2 + n_sensors + n_maps,
unroll_length=rollout_length,
sp_dim=dim
)
model.load_state_dict(torch.load(model_path), strict=False)
trainloader, testloader = localization_train_test_loaders(
data,
n_train_samples=n_samples,
n_test_samples=n_samples,
rollout_length=rollout_length,
batch_size=batch_size,
encoding=encoding,
)
print("Testing")
with torch.no_grad():
# Everything is in one batch, so this loop will only happen once
for i, data in enumerate(testloader):
combined_inputs, ssp_inputs, ssp_outputs = data
ssp_pred, lstm_outputs = model.forward_activations(combined_inputs, ssp_inputs)
predictions = np.zeros((ssp_pred.shape[0]*ssp_pred.shape[1], 2))
coords = np.zeros((ssp_pred.shape[0]*ssp_pred.shape[1], 2))
activations = np.zeros((ssp_pred.shape[0]*ssp_pred.shape[1], model.lstm_hidden_size))
assert rollout_length == ssp_pred.shape[0]
# # For each neuron, contains the average activity at each spatial bin
# # Computing for both ground truth and predicted location
# rate_maps_pred = np.zeros((model.lstm_hidden_size, len(xs), len(ys)))
# rate_maps_truth = np.zeros((model.lstm_hidden_size, len(xs), len(ys)))
print("Computing predicted locations and true locations")
# Using all data, one chunk at a time
for ri in range(rollout_length):
if encoding == 'ssp':
# computing 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_pred.detach().numpy()[ri, :, :],
heatmap_vectors, xs, ys
)
# computing 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_to_loc_v(
ssp_outputs.detach().numpy()[:, ri, :],
heatmap_vectors, xs, ys
)
elif encoding == '2d':
# copying 'predicted' coordinates, where the agent thinks it is
predictions[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_pred.detach().numpy()[ri, :, :]
# copying 'ground truth' coordinates, where the agent should be
coords[ri * ssp_pred.shape[1]:(ri + 1) * ssp_pred.shape[1], :] = ssp_outputs.detach().numpy()[:, ri, :]
# reshaping activations and converting to numpy array
activations[ri*ssp_pred.shape[1]:(ri+1)*ssp_pred.shape[1], :] = lstm_outputs.detach().numpy()[ri, :, :]
return activations, predictions, coords
def main():
parser = argparse.ArgumentParser(
'Train a network to clean up a noisy spatial semantic pointer')
parser.add_argument('--loss',
type=str,
default='cosine',
choices=['cosine', 'mse'])
parser.add_argument('--noise-type',
type=str,
default='memory',
choices=['memory', 'gaussian', 'both'])
parser.add_argument(
'--sigma',
type=float,
default=1.0,
help='sigma on the gaussian noise if noise-type==gaussian')
parser.add_argument('--train-fraction',
type=float,
default=.8,
help='proportion of the dataset to use for training')
parser.add_argument(
'--n-samples',
type=int,
default=10000,
help=
'Number of memories to generate. Total samples will be n-samples * n-items'
)
parser.add_argument('--n-items',
type=int,
default=12,
help='number of items in memory. Proxy for noisiness')
parser.add_argument('--dim',
type=int,
default=512,
help='Dimensionality of the semantic pointers')
parser.add_argument('--hidden-size',
type=int,
default=512,
help='Hidden size of the cleanup network')
parser.add_argument('--limits',
type=str,
default="-5,5,-5,5",
help='The limits of the space')
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--batch-size', type=int, default=32)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--momentum', type=float, default=0.9)
parser.add_argument('--seed', type=int, default=13)
parser.add_argument('--logdir',
type=str,
default='ssp_cleanup',
help='Directory for saved model and tensorboard log')
parser.add_argument('--load-model',
type=str,
default='',
help='Optional model to continue training from')
parser.add_argument(
'--name',
type=str,
default='',
help=
'Name of output folder within logdir. Will use current date and time if blank'
)
parser.add_argument('--weight-histogram',
action='store_true',
help='Save histograms of the weights if set')
parser.add_argument('--use-hex-ssp', action='store_true')
parser.add_argument('--optimizer',
type=str,
default='adam',
choices=['sgd', 'adam', 'rmsprop'])
args = parser.parse_args()
args.limits = tuple(float(v) for v in args.limits.split(','))
np.random.seed(args.seed)
torch.manual_seed(args.seed)
dataset_name = 'data/ssp_cleanup_dataset_dim{}_seed{}_items{}_limit{}_samples{}.npz'.format(
args.dim, args.seed, args.n_items, args.limits[1], args.n_samples)
final_test_samples = 100
final_test_items = 15
final_test_dataset_name = 'data/ssp_cleanup_test_dataset_dim{}_seed{}_items{}_limit{}_samples{}.npz'.format(
args.dim, args.seed, final_test_items, args.limits[1],
final_test_samples)
if not os.path.exists('data'):
os.makedirs('data')
rng = np.random.RandomState(seed=args.seed)
if args.use_hex_ssp:
x_axis_sp, y_axis_sp = get_axes(dim=args.dim, n=3, seed=args.seed)
else:
x_axis_sp = make_good_unitary(args.dim, rng=rng)
y_axis_sp = make_good_unitary(args.dim, rng=rng)
if args.noise_type == 'gaussian':
# Simple generation
clean_ssps = np.zeros((args.n_samples, args.dim))
coords = np.zeros((args.n_samples, 2))
for i in range(args.n_samples):
x = np.random.uniform(low=args.limits[0], high=args.limits[1])
y = np.random.uniform(low=args.limits[2], high=args.limits[3])
clean_ssps[i, :] = encode_point(x,
y,
x_axis_sp=x_axis_sp,
y_axis_sp=y_axis_sp).v
coords[i, 0] = x
coords[i, 1] = y
# Gaussian noise will be added later
noisy_ssps = clean_ssps.copy()
else:
if os.path.exists(dataset_name):
print("Loading dataset")
data = np.load(dataset_name)
clean_ssps = data['clean_ssps']
noisy_ssps = data['noisy_ssps']
else:
print("Generating SSP cleanup dataset")
clean_ssps, noisy_ssps, coords = generate_cleanup_dataset(
x_axis_sp=x_axis_sp,
y_axis_sp=y_axis_sp,
n_samples=args.n_samples,
dim=args.dim,
n_items=args.n_items,
limits=args.limits,
seed=args.seed,
)
print("Dataset generation complete. Saving dataset")
np.savez(
dataset_name,
clean_ssps=clean_ssps,
noisy_ssps=noisy_ssps,
coords=coords,
x_axis_vec=x_axis_sp.v,
y_axis_vec=x_axis_sp.v,
)
# check if the final test set has been generated yet
if os.path.exists(final_test_dataset_name):
print("Loading final test dataset")
final_test_data = np.load(final_test_dataset_name)
final_test_clean_ssps = final_test_data['clean_ssps']
final_test_noisy_ssps = final_test_data['noisy_ssps']
else:
print("Generating final test dataset")
final_test_clean_ssps, final_test_noisy_ssps, final_test_coords = generate_cleanup_dataset(
x_axis_sp=x_axis_sp,
y_axis_sp=y_axis_sp,
n_samples=final_test_samples,
dim=args.dim,
n_items=final_test_items,
limits=args.limits,
seed=args.seed,
)
print("Final test generation complete. Saving dataset")
np.savez(
final_test_dataset_name,
clean_ssps=final_test_clean_ssps,
noisy_ssps=final_test_noisy_ssps,
coords=final_test_coords,
x_axis_vec=x_axis_sp.v,
y_axis_vec=x_axis_sp.v,
)
# Add gaussian noise if required
if args.noise_type == 'gaussian' or args.noise_type == 'both':
noisy_ssps += np.random.normal(loc=0,
scale=args.sigma,
size=noisy_ssps.shape)
n_samples = clean_ssps.shape[0]
n_train = int(args.train_fraction * n_samples)
n_test = n_samples - n_train
assert (n_train > 0 and n_test > 0)
train_clean = clean_ssps[:n_train, :]
train_noisy = noisy_ssps[:n_train, :]
test_clean = clean_ssps[n_train:, :]
test_noisy = noisy_ssps[n_train:, :]
# NOTE: this dataset is actually generic and can take any input/output mapping
dataset_train = CoordDecodeDataset(vectors=train_noisy, coords=train_clean)
dataset_test = CoordDecodeDataset(vectors=test_noisy, coords=test_clean)
dataset_final_test = CoordDecodeDataset(vectors=final_test_noisy_ssps,
coords=final_test_clean_ssps)
trainloader = torch.utils.data.DataLoader(
dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=0,
)
# For testing just do everything in one giant batch
testloader = torch.utils.data.DataLoader(
dataset_test,
batch_size=len(dataset_test),
shuffle=False,
num_workers=0,
)
final_testloader = torch.utils.data.DataLoader(
dataset_final_test,
batch_size=len(dataset_final_test),
shuffle=False,
num_workers=0,
)
model = FeedForward(dim=dataset_train.dim,
hidden_size=args.hidden_size,
output_size=dataset_train.dim)
# Open a tensorboard writer if a logging directory is given
if args.logdir != '':
current_time = datetime.now().strftime('%b%d_%H-%M-%S')
save_dir = osp.join(args.logdir, current_time)
writer = SummaryWriter(log_dir=save_dir)
if args.weight_histogram:
# Log the initial parameters
for name, param in model.named_parameters():
writer.add_histogram('parameters/' + name,
param.clone().cpu().data.numpy(), 0)
mse_criterion = nn.MSELoss()
cosine_criterion = nn.CosineEmbeddingLoss()
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum)
elif args.optimizer == 'rmsprop':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
else:
raise NotImplementedError
for e in range(args.epochs):
print('Epoch: {0}'.format(e + 1))
avg_mse_loss = 0
avg_cosine_loss = 0
n_batches = 0
for i, data in enumerate(trainloader):
noisy, clean = data
if noisy.size()[0] != args.batch_size:
continue # Drop data, not enough for a batch
optimizer.zero_grad()
outputs = model(noisy)
mse_loss = mse_criterion(outputs, clean)
# Modified to use CosineEmbeddingLoss
cosine_loss = cosine_criterion(outputs, clean,
torch.ones(args.batch_size))
avg_cosine_loss += cosine_loss.data.item()
avg_mse_loss += mse_loss.data.item()
n_batches += 1
if args.loss == 'cosine':
cosine_loss.backward()
else:
mse_loss.backward()
# print(loss.data.item())
optimizer.step()
print(avg_cosine_loss / n_batches)
if args.logdir != '':
if n_batches > 0:
avg_cosine_loss /= n_batches
writer.add_scalar('avg_cosine_loss', avg_cosine_loss, e + 1)
writer.add_scalar('avg_mse_loss', avg_mse_loss, e + 1)
if args.weight_histogram and (e + 1) % 10 == 0:
for name, param in model.named_parameters():
writer.add_histogram('parameters/' + name,
param.clone().cpu().data.numpy(),
e + 1)
print("Testing")
with torch.no_grad():
for label, loader in zip(['test', 'final_test'],
[testloader, final_testloader]):
# Everything is in one batch, so this loop will only happen once
for i, data in enumerate(loader):
noisy, clean = data
outputs = model(noisy)
mse_loss = mse_criterion(outputs, clean)
# Modified to use CosineEmbeddingLoss
cosine_loss = cosine_criterion(outputs, clean,
torch.ones(len(loader)))
print(cosine_loss.data.item())
if args.logdir != '':
# TODO: get a visualization of the performance
# show plots of the noisy, clean, and cleaned up with the network
# note that the plotting mechanism itself uses nearest neighbors, so has a form of cleanup built in
xs = np.linspace(args.limits[0], args.limits[1], 256)
ys = np.linspace(args.limits[0], args.limits[1], 256)
heatmap_vectors = get_heatmap_vectors(xs, ys, x_axis_sp,
y_axis_sp)
noisy_coord = ssp_to_loc_v(noisy, heatmap_vectors, xs, ys)
pred_coord = ssp_to_loc_v(outputs, heatmap_vectors, xs, ys)
clean_coord = ssp_to_loc_v(clean, heatmap_vectors, xs, ys)
fig_noisy_coord, ax_noisy_coord = plt.subplots()
fig_pred_coord, ax_pred_coord = plt.subplots()
fig_clean_coord, ax_clean_coord = plt.subplots()
plot_predictions_v(noisy_coord,
clean_coord,
ax_noisy_coord,
min_val=args.limits[0],
max_val=args.limits[1],
fixed_axes=True)
plot_predictions_v(pred_coord,
clean_coord,
ax_pred_coord,
min_val=args.limits[0],
max_val=args.limits[1],
fixed_axes=True)
plot_predictions_v(clean_coord,
clean_coord,
ax_clean_coord,
min_val=args.limits[0],
max_val=args.limits[1],
fixed_axes=True)
writer.add_figure('{}/original_noise'.format(label),
fig_noisy_coord)
writer.add_figure('{}/test_set_cleanup'.format(label),
fig_pred_coord)
writer.add_figure('{}/ground_truth'.format(label),
fig_clean_coord)
# fig_hist = plot_histogram(predictions=outputs, coords=coord)
# writer.add_figure('test set histogram', fig_hist)
writer.add_scalar('{}/test_cosine_loss'.format(label),
cosine_loss.data.item())
writer.add_scalar('{}/test_mse_loss'.format(label),
mse_loss.data.item())
# Close tensorboard writer
if args.logdir != '':
writer.close()
torch.save(model.state_dict(), osp.join(save_dir, 'model.pt'))
params = vars(args)
# # Additionally save the axis vectors used
# params['x_axis_vec'] = list(x_axis_sp.v)
# params['y_axis_vec'] = list(y_axis_sp.v)
with open(osp.join(save_dir, "params.json"), "w") as f:
json.dump(params, f)
for i in range(args.n_samples):
encodings[i, :] = encoding_functions[ei][di*n_seeds + si](
np.array(
[true_pos[i, 0], true_pos[i, 1]]
)
)
for ni, n in enumerate(noise_levels):
print("Noise level: {}".format(n))
noisy_encodings = encodings + np.random.normal(loc=0, scale=n, size=(args.n_samples, d))
predictions = ssp_to_loc_v(
noisy_encodings,
heatmap_vectors, xs, ys
)
# Root mean squared error
# results[ei, di, ni] = np.sqrt(((predictions - true_pos)**2).mean())
rmse = np.sqrt(((predictions - true_pos)**2).mean())
df = df.append(
{
'Dimensions': d,
'Noise Level': n,
'Encoding': e,
'RMSE': rmse,
'Seed': s
},
ignore_index=True,
print(cosine_loss.data.item())
if args.logdir != '':
# TODO: get a visualization of the performance
# show plots of the noisy, clean, and cleaned up with the network
# note that the plotting mechanism itself uses nearest neighbors, so has a form of cleanup built in
xs = np.linspace(-args.limit*2, args.limit*2, 256)
ys = np.linspace(-args.limit*2, args.limit*2, 256)
heatmap_vectors = get_projected_heatmap_vectors(xs, ys, x_axis_sp, y_axis_sp, z_axis_sp)
noisy_coord = ssp_to_loc_v(
noisy,
heatmap_vectors, xs, ys
)
pred_coord = ssp_to_loc_v(
outputs,
heatmap_vectors, xs, ys
)
clean_coord = ssp_to_loc_v(
clean,
heatmap_vectors, xs, ys
)
fig_noisy_coord, ax_noisy_coord = plt.subplots()
fig_pred_coord, ax_pred_coord = plt.subplots()
fig_clean_coord, ax_clean_coord = plt.subplots()
def experiment(dim=512,
n_hierarchy=3,
n_items=16,
seed=0,
limit=5,
res=128,
thresh=0.5,
neural=False,
neurons_per_dim=25,
time_per_item=1.0,
max_items=100):
rng = np.random.RandomState(seed=seed)
X, Y = get_fixed_dim_sub_toriod_axes(
dim=dim,
n_proj=3,
scale_ratio=0,
scale_start_index=0,
rng=rng,
eps=0.001,
)
xs = np.linspace(-limit, limit, res)
ys = np.linspace(-limit, limit, res)
hmv = get_heatmap_vectors(xs, ys, X, Y)
item_vecs = rng.normal(size=(n_items, dim))
for i in range(n_items):
item_vecs[i, :] = item_vecs[i, :] / np.linalg.norm(item_vecs[i, :])
locations = rng.uniform(low=-limit, high=limit, size=(n_items, 2))
if n_hierarchy == 1: # no hierarchy case
# Encode items into memory
mem = np.zeros((dim, ))
for i in range(n_items):
mem += (spa.SemanticPointer(data=item_vecs[i, :]) *
encode_point(locations[i, 0], locations[i, 1], X, Y)).v
mem /= np.linalg.norm(mem)
mem_sp = spa.SemanticPointer(data=mem)
estims = np.zeros((
n_items,
dim,
))
sims = np.zeros((n_items, ))
if neural:
# save time for very large numbers of items
n_exp_items = min(n_items, max_items)
estims = np.zeros((
n_exp_items,
dim,
))
sims = np.zeros((n_exp_items, ))
model = nengo.Network(seed=seed)
with model:
input_node = nengo.Node(
lambda t: item_vecs[int(np.floor(t)) % n_items, :],
size_in=0,
size_out=dim)
mem_node = nengo.Node(mem, size_in=0, size_out=dim)
cconv = nengo.networks.CircularConvolution(
n_neurons=neurons_per_dim, dimensions=dim, invert_b=True)
nengo.Connection(mem_node, cconv.input_a)
nengo.Connection(input_node, cconv.input_b)
out_node = nengo.Node(size_in=dim, size_out=0)
nengo.Connection(cconv.output, out_node)
p_out = nengo.Probe(out_node, synapse=0.01)
sim = nengo.Simulator(model)
sim.run(n_exp_items * time_per_item)
output_data = sim.data[p_out]
timesteps_per_item = int(time_per_item / 0.001)
# timestep offset to cancel transients
offset = 100
for i in range(n_exp_items):
estims[i, :] = output_data[i * timesteps_per_item +
offset:(i + 1) *
timesteps_per_item, :].mean(axis=0)
sims[i] = np.dot(
estims[i, :],
encode_point(locations[i, 0], locations[i, 1], X, Y).v)
pred_locs = ssp_to_loc_v(estims, hmv, xs, ys)
errors = np.linalg.norm(pred_locs - locations[:n_exp_items, :],
axis=1)
accuracy = len(np.where(errors < thresh)[0]) / n_items
rmse = np.sqrt(np.mean(errors**2))
sim = np.mean(sims)
else:
# retrieve items
for i in range(n_items):
estims[i, :] = (mem_sp *
~spa.SemanticPointer(data=item_vecs[i, :])).v
sims[i] = np.dot(
estims[i, :],
encode_point(locations[i, 0], locations[i, 1], X, Y).v)
pred_locs = ssp_to_loc_v(estims, hmv, xs, ys)
errors = np.linalg.norm(pred_locs - locations, axis=1)
accuracy = len(np.where(errors < thresh)[0]) / n_items
rmse = np.sqrt(np.mean(errors**2))
sim = np.mean(sims)
elif n_hierarchy == 2:
# TODO: generate vocab and input sequences
n_ids = int(np.sqrt(n_items))
f_n_ids = np.sqrt(n_items)
id_vecs = rng.normal(size=(n_ids, dim))
for i in range(n_ids):
id_vecs[i, :] = id_vecs[i, :] / np.linalg.norm(id_vecs[i, :])
# items to be included in each ID vec
item_sums = np.zeros((n_ids, dim))
item_loc_sums = np.zeros((n_ids, dim))
for i in range(n_items):
id_ind = min(i // n_ids, n_ids - 1)
# id_ind = min(int(i / f_n_ids), n_ids - 1)
item_sums[id_ind, :] += item_vecs[i, :]
item_loc_sums[id_ind, :] += (
spa.SemanticPointer(data=item_vecs[i, :]) *
encode_point(locations[i, 0], locations[i, 1], X, Y)).v
# Encode id_vecs into memory, each id is bound to something that has similarity to all items in the ID's map
mem = np.zeros((dim, ))
for i in range(n_ids):
# normalize previous memories
item_sums[i, :] = item_sums[i, :] / np.linalg.norm(item_sums[i, :])
item_loc_sums[i, :] = item_loc_sums[i, :] / np.linalg.norm(
item_loc_sums[i, :])
mem += (spa.SemanticPointer(data=id_vecs[i, :]) *
spa.SemanticPointer(data=item_sums[i, :])).v
mem /= np.linalg.norm(mem)
mem_sp = spa.SemanticPointer(data=mem)
estims = np.zeros((
n_items,
dim,
))
sims = np.zeros((n_items, ))
# retrieve items
for i in range(n_items):
# noisy ID for the map with this item
estim_id = (mem_sp * ~spa.SemanticPointer(data=item_vecs[i, :])).v
# get closest clean match
id_sims = np.zeros((n_ids, ))
for j in range(n_ids):
id_sims[j] = np.dot(estim_id, id_vecs[j, :])
best_ind = np.argmax(id_sims)
# clean_id = id_vecs[best_ind, :]
# item_loc_sums comes from the associative mapping from clean_id
estims[i, :] = (
spa.SemanticPointer(data=item_loc_sums[best_ind, :]) *
~spa.SemanticPointer(data=item_vecs[i, :])).v
sims[i] = np.dot(
estims[i, :],
encode_point(locations[i, 0], locations[i, 1], X, Y).v)
pred_locs = ssp_to_loc_v(estims, hmv, xs, ys)
errors = np.linalg.norm(pred_locs - locations, axis=1)
accuracy = len(np.where(errors < thresh)[0]) / n_items
rmse = np.sqrt(np.mean(errors**2))
sim = np.mean(sims)
elif n_hierarchy == 3:
# n_ids = int(np.cbrt(n_items))
f_n_ids = np.cbrt(n_items)
n_ids = int(np.ceil(np.cbrt(n_items)))
n_ids_inner = int(np.ceil(np.sqrt(n_items / n_ids)))
# f_n_ids = np.cbrt(n_items)
id_outer_vecs = rng.normal(size=(n_ids, dim))
id_inner_vecs = rng.normal(size=(n_ids_inner, dim))
for i in range(n_ids):
id_outer_vecs[i, :] = id_outer_vecs[i, :] / np.linalg.norm(
id_outer_vecs[i, :])
# for j in range(n_ids):
# id_inner_vecs[i*n_ids+j, :] = id_inner_vecs[i*n_ids+j, :] / np.linalg.norm(id_inner_vecs[i*n_ids+j, :])
for i in range(n_ids_inner):
id_inner_vecs[i, :] = id_inner_vecs[i, :] / np.linalg.norm(
id_inner_vecs[i, :])
# items to be included in each ID vec
item_outer_sums = np.zeros((n_ids, dim))
# item_inner_sums = np.zeros((n_ids*n_ids, dim))
item_inner_sums = np.zeros((n_ids_inner, dim))
item_loc_outer_sums = np.zeros((n_ids, dim))
# item_loc_inner_sums = np.zeros((n_ids*n_ids, dim))
item_loc_inner_sums = np.zeros((n_ids_inner, dim))
for i in range(n_items):
id_outer_ind = min(int(i / (f_n_ids * f_n_ids)), n_ids - 1)
id_inner_ind = min(int(i / f_n_ids), n_ids_inner - 1)
item_outer_sums[id_outer_ind, :] += item_vecs[i, :]
item_inner_sums[id_inner_ind, :] += item_vecs[i, :]
item_loc_outer_sums[id_outer_ind, :] += (
spa.SemanticPointer(data=item_vecs[i, :]) *
encode_point(locations[i, 0], locations[i, 1], X, Y)).v
item_loc_inner_sums[id_inner_ind, :] += (
spa.SemanticPointer(data=item_vecs[i, :]) *
encode_point(locations[i, 0], locations[i, 1], X, Y)).v
# Encode id_vecs into memory, each id is bound to something that has similarity to all items in the ID's map
mem_outer = np.zeros((dim, ))
mem_inner = np.zeros((
n_ids,
dim,
))
for i in range(n_ids):
# normalize previous memories
item_outer_sums[i, :] = item_outer_sums[i, :] / np.linalg.norm(
item_outer_sums[i, :])
item_loc_outer_sums[i, :] = item_loc_outer_sums[
i, :] / np.linalg.norm(item_loc_outer_sums[i, :])
mem_outer += (spa.SemanticPointer(data=id_outer_vecs[i, :]) *
spa.SemanticPointer(data=item_outer_sums[i, :])).v
for j in range(n_ids_inner):
# normalize previous memories
item_inner_sums[j, :] = item_inner_sums[j, :] / np.linalg.norm(
item_inner_sums[j, :])
item_loc_inner_sums[j, :] = item_loc_inner_sums[
j, :] / np.linalg.norm(item_loc_inner_sums[j, :])
i = min(int(j / n_ids), n_ids - 1)
mem_inner[i, :] += (
spa.SemanticPointer(data=id_inner_vecs[j, :]) *
spa.SemanticPointer(data=item_inner_sums[j, :])).v
mem_inner[i, :] /= np.linalg.norm(mem_inner[i, :])
mem_outer /= np.linalg.norm(mem_outer)
mem_outer_sp = spa.SemanticPointer(data=mem_outer)
estims = np.zeros((
n_items,
dim,
))
sims = np.zeros((n_items, ))
if neural:
# time for each item, in seconds
time_per_item = 1.0
model = nengo.Network(seed=seed)
with model:
inp_node = nengo.Node('?', size_in=0, size_out=dim)
estim_outer_id = nengo.Ensemble(dimension=dim,
n_neurons=dim *
neurons_per_dim)
out_node = nengo.Node(size_in=dim, size_out=0)
p_out = nengo.Probe(out_node, synapse=0.01)
sim = nengo.Simulator(model)
sim.run(n_items * time_per_item)
else:
# non-neural version
# retrieve items
for i in range(n_items):
# noisy outer ID for the map with this item
estim_outer_id = (mem_outer_sp *
~spa.SemanticPointer(data=item_vecs[i, :])).v
# get closest clean match
id_sims = np.zeros((n_ids))
for j in range(n_ids):
id_sims[j] = np.dot(estim_outer_id, id_outer_vecs[j, :])
best_ind = np.argmax(id_sims)
# noisy inner ID for the map with this item
estim_inner_id = (
spa.SemanticPointer(data=mem_inner[best_ind, :]) *
~spa.SemanticPointer(data=item_vecs[i, :])).v
# get closest clean match
id_sims = np.zeros((n_ids_inner))
for j in range(n_ids_inner):
id_sims[j] = np.dot(estim_inner_id, id_inner_vecs[j, :])
best_ind = np.argmax(id_sims)
# item_loc_sums comes from the associative mapping from clean_id
estims[i, :] = (spa.SemanticPointer(
data=item_loc_inner_sums[best_ind, :]) *
~spa.SemanticPointer(data=item_vecs[i, :])).v
sims[i] = np.dot(
estims[i, :],
encode_point(locations[i, 0], locations[i, 1], X, Y).v)
pred_locs = ssp_to_loc_v(estims, hmv, xs, ys)
errors = np.linalg.norm(pred_locs - locations, axis=1)
accuracy = len(np.where(errors < thresh)[0]) / n_items
rmse = np.sqrt(np.mean(errors**2))
sim = np.mean(sims)
else:
# 4 split hierarchy
vocab = spa.Vocabulary(dimensions=dim,
pointer_gen=np.random.RandomState(seed=seed))
filler_id_keys = []
filler_keys = []
mapping = {}
items_left = n_items
n_levels = 0
while items_left > 1:
n_levels += 1
items_left /= 4
print(n_levels)
# Location Values, labelled SSP
for i in range(n_items):
# vocab.populate('Item{}'.format(i))
vocab.add('Loc{}'.format(i),
encode_point(locations[i, 0], locations[i, 1], X, Y).v)
# level IDs, e.g. CITY, PROVINCE, COUNTRY
for i in range(n_levels):
vocab.populate('LevelSlot{}.unitary()'.format(i))
# sp = spa.SemanticPointer()
# Item IDs, e.g. Waterloo_ID
for i in range(n_items):
vocab.populate('ItemID{}.unitary()'.format(i))
# level labels (fillers for level ID slots), e.g. Waterloo_ID, Ontario_ID, Canada_ID
for i in range(n_levels):
for j in range(int(n_items / (4**(n_levels - i - 1)))):
vocab.populate('LevelFillerID{}_{}.unitary()'.format(i, j))
# filler_id_keys.append('LevelFillerID{}_{}'.format(i, j))
# filler_keys.append('LevelFiller{}_{}'.format(i, j))
# mapping['LevelFillerID{}_{}'.format(i, j)] = 'LevelFiller{}_{}'.format(i, j)
# Second last level with item*location pairs
for i in range(int(n_items / 4)):
id_str = []
for k in range(n_levels - 1):
id_str.append('LevelSlot{} * LevelFillerID{}_{}'.format(
k, k, int(i * 4 / (4**(n_levels - k - 1)))))
data_str = []
for j in range(4):
ind = i * 4 + j
data_str.append('ItemID{}*Loc{}'.format(ind, ind))
vocab.populate('Item{} = ({}).normalized()'.format(
# i, ' + '.join(id_str + ['LevelSlot{} * LevelFillerID{}_{}'.format(n_levels - 2, n_levels - 2, j)])
ind,
' + '.join(id_str + [
'LevelSlot{} * LevelFillerID{}_{}'.format(
n_levels - 1, n_levels - 1, j)
])))
# vocab.populate('LevelFiller{}_{} = {}'.format(n_levels - 1, i, ' + '.join(data_str)))
vocab.populate('LevelFiller{}_{} = ({}).normalized()'.format(
n_levels - 2, i, ' + '.join(data_str)))
# only appending the ones used
filler_id_keys.append('LevelFillerID{}_{}'.format(n_levels - 2, i))
filler_keys.append('LevelFiller{}_{}'.format(n_levels - 2, i))
mapping['LevelFillerID{}_{}'.format(
n_levels - 2, i)] = 'LevelFiller{}_{}'.format(n_levels - 2, i)
print(sorted(list(vocab.keys())))
# Given each ItemID, calculate the corresponding Loc
# Can map from ItemID{X} -> Item{X}
# Query based on second last levelID to get the appropriate LevelFillerID
# map from LevelFillerID -> LevelFiller
# do the query LevelFiller *~ ItemID{X} to get Loc{X}
possible_level_filler_id_vecs = np.zeros((int(n_items / 4), dim))
for i in range(int(n_items / 4)):
possible_level_filler_id_vecs[i] = vocab[
'LevelFillerID{}_{}'.format(n_levels - 2, i)].v
estims = np.zeros((
n_items,
dim,
))
sims = np.zeros((n_items, ))
if neural:
# save time for very large numbers of items
n_exp_items = min(n_items, max_items)
estims = np.zeros((
n_exp_items,
dim,
))
sims = np.zeros((n_exp_items, ))
filler_id_vocab = vocab.create_subset(keys=filler_id_keys)
filler_vocab = vocab.create_subset(keys=filler_keys)
filler_all_vocab = vocab.create_subset(keys=filler_keys +
filler_id_keys)
model = nengo.Network(seed=seed)
with model:
# The changing item query. Full expanded item, not just ID
item_input_node = nengo.Node(lambda t: vocab['Item{}'.format(
int(np.floor(t)) % n_items)].v,
size_in=0,
size_out=dim)
# item_input_node = spa.Transcode(lambda t: 'Item{}'.format(int(np.floor(t))), output_vocab=vocab)
# The ID for the changing item query
item_id_input_node = nengo.Node(lambda t: vocab[
'ItemID{}'.format(int(np.floor(t)) % n_items)].v,
size_in=0,
size_out=dim)
# item_id_input_node = spa.Transcode(lambda t: 'ItemID{}'.format(int(np.floor(t))), output_vocab=vocab)
# Fixed memory based on the level slot to access
level_slot_input_node = nengo.Node(
lambda t: vocab['LevelSlot{}'.format(n_levels - 2)].v,
size_in=0,
size_out=dim)
model.cconv_noisy_level_filler = nengo.networks.CircularConvolution(
n_neurons=neurons_per_dim * 2,
dimensions=dim,
invert_b=True)
nengo.Connection(item_input_node,
model.cconv_noisy_level_filler.input_a)
nengo.Connection(level_slot_input_node,
model.cconv_noisy_level_filler.input_b)
# Note: this is set up as heteroassociative between ID and the content (should clean up as well)
model.noisy_level_filler_id_cleanup = spa.ThresholdingAssocMem(
threshold=0.4,
input_vocab=filler_id_vocab,
output_vocab=filler_vocab,
# mapping=vocab.keys(),
mapping=mapping,
function=lambda x: x > 0.)
nengo.Connection(model.cconv_noisy_level_filler.output,
model.noisy_level_filler_id_cleanup.input)
model.cconv_location = nengo.networks.CircularConvolution(
n_neurons=neurons_per_dim * 2,
dimensions=dim,
invert_b=True)
nengo.Connection(model.noisy_level_filler_id_cleanup.output,
model.cconv_location.input_a)
nengo.Connection(item_id_input_node,
model.cconv_location.input_b)
out_node = nengo.Node(size_in=dim, size_out=0)
nengo.Connection(model.cconv_location.output, out_node)
p_out = nengo.Probe(out_node, synapse=0.01)
sim = nengo.Simulator(model)
sim.run(n_exp_items * time_per_item)
output_data = sim.data[p_out]
timesteps_per_item = int(time_per_item / 0.001)
# timestep offset to cancel transients
offset = 100
for i in range(n_exp_items):
estims[i, :] = output_data[i * timesteps_per_item +
offset:(i + 1) *
timesteps_per_item, :].mean(axis=0)
sims[i] = np.dot(
estims[i, :],
encode_point(locations[i, 0], locations[i, 1], X, Y).v)
pred_locs = ssp_to_loc_v(estims, hmv, xs, ys)
errors = np.linalg.norm(pred_locs - locations[:n_exp_items, :],
axis=1)
accuracy = len(np.where(errors < thresh)[0]) / n_items
rmse = np.sqrt(np.mean(errors**2))
sim = np.mean(sims)
else:
# non-neural version
# retrieve items
for i in range(n_items):
noisy_level_filler_id = vocab['Item{}'.format(
i)] * ~vocab['LevelSlot{}'.format(n_levels - 2)]
# cleanup filler id
n_fillers = int(n_items / 4)
sim = np.zeros((n_fillers, ))
for j in range(n_fillers):
sim[j] = np.dot(noisy_level_filler_id.v,
possible_level_filler_id_vecs[j, :])
filler_id_ind = np.argmax(sim)
# query the appropriate filler
loc_estim = vocab['LevelFiller{}_{}'.format(
n_levels - 2,
filler_id_ind)] * ~vocab['ItemID{}'.format(i)]
estims[i, :] = loc_estim.v
sims[i] = np.dot(
estims[i, :],
encode_point(locations[i, 0], locations[i, 1], X, Y).v)
pred_locs = ssp_to_loc_v(estims, hmv, xs, ys)
errors = np.linalg.norm(pred_locs - locations, axis=1)
accuracy = len(np.where(errors < thresh)[0]) / n_items
rmse = np.sqrt(np.mean(errors**2))
sim = np.mean(sims)
return rmse, accuracy, sim
n_samples += len(np.where(wall_overlay == False)[0])
print('wall_overlay.shape', wall_overlay.shape)
hmv = get_encoding_heatmap_vectors(xs,
ys,
args.dim,
encoding_func,
normalize=False)
predictions = np.zeros((res * res, 2))
truth = np.zeros((res * res, 2))
# computing 'predicted' coordinates, where the agent thinks it is
predictions[:, :] = ssp_to_loc_v(batch_data, hmv, xs, ys)
truth[:, :] = ssp_to_loc_v(true_ssps, hmv, xs, ys)
squared_error = np.sum(
np.linalg.norm(predictions[wall_overlay == False, :] -
coords[wall_overlay == False, :],
axis=1)**2)
plot_predictions_v(
predictions=predictions[wall_overlay == False, :],
coords=coords[wall_overlay == False, :],
ax=ax[1, bi],
min_val=limit_low,
max_val=limit_high,
fixed_axes=True,
