def _log_policy(self, z, k, c, log_pi, mu, Sigma):
z, k, c, log_pi, mu, Sigma = [
torch.from_numpy(x) if type(x) is np.ndarray else x
for x in [z, k, c, log_pi, mu, Sigma]
]
# mean, cov = self._get_mean_cov_z(c, k, mu, Sigma)
# log_z_given_k_c = multivariate_normal.MultivariateNormal(mean, precision_matrix=torch.inverse(cov)).log_prob(z)
# log_k_given_c = self._log_responsabilities(c, log_pi, mu, Sigma).T[range(c.shape[0]), k]
#
# return log_z_given_k_c + log_k_given_c
# mu = mu.squeeze(3)
log_pi_k = log_pi[k]
log_z_given_k = multivariate_normal.MultivariateNormal(
mu[k], Sigma[k]).log_prob(z)
m, s = self.get_params_c_given_z_k(z, k)
log_c_given_z_k = multivariate_normal.MultivariateNormal(m,
s).log_prob(c)
log_c_z_k = log_c_given_z_k + log_z_given_k + log_pi_k
stack = []
for i in range(self._n_cluster):
k_i = torch.tensor([i]).expand([k.shape[0]])
m, s = self.get_params_c_given_k(k_i, mu, Sigma)
log_c_given_k = multivariate_normal.MultivariateNormal(
m, s).log_prob(c)
log_pi_k = log_pi[k_i]
stack.append(log_pi_k + log_c_given_k)
log_c = sum_logs(torch.stack(stack, 0), axis=0)
return log_c_z_k - log_c
def langevin_sampling(zs, z_dim, fake_labels, generator, discriminator,
batch_size, langevin_rate, langevin_noise_std,
langevin_decay, langevin_decay_steps, langevin_steps,
device):
scaler = 1.0
apply_decay = langevin_decay > 0 and langevin_decay_steps > 0
mean = torch.zeros(z_dim, device=device)
prior_std = torch.eye(z_dim, device=device)
lgv_std = prior_std * langevin_noise_std
prior = MN.MultivariateNormal(loc=mean, covariance_matrix=prior_std)
lgv_prior = MN.MultivariateNormal(loc=mean, covariance_matrix=lgv_std)
for i in range(langevin_steps):
zs = autograd.Variable(zs, requires_grad=True)
fake_images = generator(zs, fake_labels, eval=True)
fake_dict = discriminator(fake_images, fake_labels, eval=True)
energy = -prior.log_prob(zs) - fake_dict["adv_output"]
z_grads = losses.cal_deriv(inputs=zs, outputs=energy, device=device)
zs = zs - 0.5 * langevin_rate * z_grads + (
langevin_rate**0.5) * lgv_prior.sample([batch_size]) * scaler
if apply_decay and (i + 1) % langevin_decay_steps == 0:
langevin_rate *= langevin_decay
scaler *= langevin_decay
return zs
def generate_data(samp_size):
mean_1 = torch.Tensor([4,4])
covar_1 = torch.Tensor([[2,1],[1,2]])
sampler_1 = mn.MultivariateNormal(mean_1,covar_1)
mean_2 = torch.Tensor([-4,-4])
covar_2 = torch.Tensor([[2,1],[1,2]])
sampler_2 = mn.MultivariateNormal(mean_2,covar_2)
sample_lebel_1 = sampler_1.sample(sample_shape=(int(samp_size/2),1))
sample_lebel_2 = sampler_2.sample(sample_shape=(int(samp_size/2),1))
return sample_lebel_1,sample_lebel_2
def sample_big_data(feature_size,
nb_of_classes,
samples_per_class,
inter_class_distance=2):
full_data = {}
data_size = nb_of_classes * samples_per_class
means_ = init_means(feature_size, nb_of_classes, inter_class_distance)
covariances = init_covariances(feature_size, nb_of_classes, 1)
full_data['data_train'] = torch.zeros(data_size, feature_size)
full_data['data_test'] = torch.zeros(data_size, feature_size)
full_data['labels_train'] = torch.zeros(data_size)
full_data['labels_test'] = torch.zeros(data_size)
bar = Bar('Generating data ', max=nb_of_classes)
for idx_class in range(nb_of_classes):
bar.next()
data_sampler = mv_n.MultivariateNormal(
torch.DoubleTensor(means_[idx_class]),
torch.DoubleTensor(covariances[idx_class]))
full_data['data_train'][idx_class * samples_per_class:(idx_class + 1) *
samples_per_class] = data_sampler.sample(
(samples_per_class, ))
full_data['data_test'][idx_class * samples_per_class:(idx_class + 1) *
samples_per_class] = data_sampler.sample(
(samples_per_class, ))
full_data['labels_train'][idx_class *
samples_per_class:(idx_class + 1) *
samples_per_class] = idx_class
full_data['labels_test'][idx_class *
samples_per_class:(idx_class + 1) *
samples_per_class] = idx_class
bar.finish()
return full_data
def generate_data(
samp_size, mean_1, covar_1, mean_2,
covar_2): # generate synthetic data with a multivariate normal
# distribution
mean_1 = torch.Tensor(mean_1)
covar_1 = torch.Tensor(covar_1)
sampler_1 = mn.MultivariateNormal(mean_1, covar_1)
mean_2 = torch.Tensor(mean_2)
covar_2 = torch.Tensor(covar_2)
sampler_2 = mn.MultivariateNormal(mean_2, covar_2)
sample_lebel_1 = sampler_1.sample(sample_shape=(int(samp_size / 2), 1))
sample_lebel_2 = sampler_2.sample(sample_shape=(int(samp_size / 2), 1))
return sample_lebel_1, sample_lebel_2
def __init__(self,
net,
v_net,
w_momentum,
w_weight_decay,
noise_add=False):
"""
Args:
net
w_momentum: weights momentum
"""
self.net = net
# self.v_net = copy.deepcopy(net)
self.v_net = v_net
self.w_momentum = w_momentum
self.w_weight_decay = w_weight_decay
self.noise_add = noise_add
shape_gaussian = {}
if self.noise_add:
for param in self.net.alphas():
shape_gaussian[param.data.shape] = gaussian.MultivariateNormal(
torch.zeros(param.data.shape),
torch.eye(param.data.shape[-1]))
self.shape_gaussian = shape_gaussian
def train(epoch, model, image_size, latent_dim, prior='Gauss'):
model.train()
train_loss = list(range(len(train_loader)))
train_model_kernel = 0
train_encoder_kernel = 0
train_cross_kernel = 0
recon_loss = list(range(len(train_loader)))
sinkhorn_solver = SinkhornSolver(epsilon=0.01, iterations=20)
start_time = time.time()
for batch_idx, (real_data, _) in enumerate(train_loader):
real_data = real_data.to(device)
real_data = real_data.type(torch.float32)
optimizer.zero_grad()
if prior == 'Gauss':
latent_priors = multivariate_normal.MultivariateNormal(loc=torch.zeros(latent_dim),
covariance_matrix=torch.eye(latent_dim)).\
sample(sample_shape=(real_data.size()[0],)).to(device)
else:
latent_priors = Variable(
-2 * torch.rand(real_data.size()[0], latent_dim) + 1,
requires_grad=False).to(device)
mu, logvar, latent_encoded = model.a_encoder(
real_data.view(-1, image_size))
decoded_data = model.a_decoder(latent_priors)
latent_decoded_data, _, _ = model.a_encoder(decoded_data)
reconstructed_data = model.a_decoder(latent_encoded)
observable_error = ((real_data.view(-1, image_size) -
reconstructed_data).pow(2).mean(-1)).mean()
C1 = compute_cost(decoded_data, real_data.view(-1, image_size))
C2 = compute_cost(decoded_data, reconstructed_data)
C3 = compute_cost(latent_priors - latent_decoded_data,
latent_encoded - mu)
C4 = compute_cost(torch.zeros_like(reconstructed_data),
real_data.view(-1, image_size) - reconstructed_data)
loss, _ = sinkhorn_solver(decoded_data,
real_data.view(-1, image_size),
C=C1 + C2 + C3 + C4)
loss.backward()
train_loss[batch_idx] = loss.item()
recon_loss[batch_idx] = observable_error.item()
optimizer.step()
if batch_idx % args.log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(real_data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
end_time = time.time()
end_time = end_time - start_time
print(
'====> Epoch: {} Average loss: {:.4f}, Model Kernel: {:.6f},Encoder Kernel: {:.6f}, Cross Kernel: {:.6f}, Observable Error: {:.6f} Time: {:.6f}'
.format(epoch,
np.array(train_loss).mean(0), train_model_kernel / batch_idx,
train_encoder_kernel / batch_idx,
train_cross_kernel / batch_idx,
np.array(recon_loss).mean(0), end_time))
return np.array(train_loss).mean(0)
def forward(self, state):
H = F.relu(self.w1(state))
mu = self.w2(H)
value = self.w3(H)
sigma = self.w4(H)[0]
covar = torch.diag(torch.exp(sigma))
m = multivariate_normal.MultivariateNormal(mu, covar)
return m, value
def generate_latent(sample_labels, means, covariances):
# res = np.array([
# np.random.multivariate_normal(means[e], covariances[e])
# for e in sample_labels
# ])
# return torch.tensor(res).float()
# print(len(means[sample_labels[0]]), len(covariances[sample_labels[0]]))
distribution = mn.MultivariateNormal(means[sample_labels[0]],
covariances[sample_labels[0]])
fake_z = distribution.sample((1, ))
for c in range(1, len(sample_labels[0:])):
fake_z = torch.cat(
(fake_z,
mn.MultivariateNormal(means[sample_labels[c]],
covariances[sample_labels[c]]).sample(
(1, ))),
dim=0)
return fake_z.float()
def sample_hypernet(hypernet):
netE, W1, W2, W3 = hypernet
mean, cov = torch.zeros(300), torch.eye(300)
D = N.MultivariateNormal(mean, cov)
z = D.sample((32,)).cuda()
z.requires_grad = True
codes = netE(z)
l1 = W1(codes[0])
l2 = W2(codes[1])
l3 = W3(codes[2])
return l1, l2, l3
def initialize_synthetic_sampler(dim,
nb_classes,
intra_class_distance,
max_axis=1,
epsilon=1e-3):
means_ = init_means(dim, nb_classes, intra_class_distance, epsilon)
covariances = init_covariances(dim, nb_classes, max_axis, False)
data_class_sampler = {}
for idx_class in range(nb_classes):
data_class_sampler[idx_class] = mv_n.MultivariateNormal(
torch.DoubleTensor(means_[idx_class]),
torch.DoubleTensor(covariances[idx_class]))
return data_class_sampler
def generate_full(self, c, noise=True, isomorphic_noise=False):
"""
Sample w, z, k from a vector of contextes.
:param c:
:param noise:
:param isomorphic_noise:
:return:
"""
c = torch.from_numpy(c).view(1, c.shape[1])
p = np.exp(
self._log_responsabilities(c, self._last_log_pi, self._last_mu,
self._last_Sigma).detach().numpy()).T
k = [np.random.choice(range(self._n_cluster), p=p_i) for p_i in p]
O, o = self._Omega[k], self._omega[k]
# mean, prec = self._get_stable_mean_prec_z(c, k, self._last_mu, self._last_Sigma) # TODO: check single param here
mean, prec = self._get_mean_prec_z(
c, k, self._last_mu,
self._last_Sigma) # TODO: check single param here
if noise:
m = multivariate_normal.MultivariateNormal(mean,
precision_matrix=prec)
z = m.sample()
else:
z = mean.unsqueeze(0)
if isomorphic_noise:
m = multivariate_normal.MultivariateNormal(
(O @ z.unsqueeze(2)).squeeze(2) + o,
self.sigma_sq[k].unsqueeze(1).unsqueeze(2) *
torch.eye(self._action_dim).expand(len(k), self._action_dim,
self._action_dim))
w = m.sample()
else:
w = O @ z.unsqueeze(2) + o.unsqueeze(2)
return w.detach().numpy().squeeze(), z.detach().numpy().squeeze(
), np.array(k).squeeze()
def Z(d, beta):
"""
Returns a d-dimensional vector drawn from \mathcal{N}(0,\beta^{2}I_{d})
Arguments:
d : int
Dimensionality
beta : float
Standard deviation for the isotropic Gaussian
Returns:
d-dimensional vector
"""
return mvn.MultivariateNormal(torch.zeros(d), torch.diag(torch.ones(d) * beta**2))
def _test_log_policy(self, z, k, c, log_pi, mu, Sigma):
z, k, c, log_pi, mu, Sigma = [
torch.from_numpy(x) if type(x) is np.ndarray else x
for x in [z, k, c, log_pi, mu, Sigma]
]
mean, prec = self._get_stable_mean_prec_z(c, k, mu, Sigma)
log_z_given_k_c = multivariate_normal.MultivariateNormal(
mean, precision_matrix=prec).log_prob(z)
log_k_given_c = self._log_responsabilities(c, log_pi, mu,
Sigma).T[range(c.shape[0]),
k]
return log_z_given_k_c + log_k_given_c
def _log_p_c_k(self, c, k, mu, Sigma):
"""
compute the log-probability of the context c given the cluster k.
:param c: xontext
:param k: cluster
:return:
"""
n = k.shape[0]
mean = self._C[k] @ mu[k].unsqueeze(2) + self._c[k].unsqueeze(2)
base_cov = self.sigma_sq[k].view(-1, 1, 1) \
* torch.eye(self._context_dim, dtype=torch.float64).unsqueeze(0).repeat(n, 1, 1)
cov = base_cov + self._C[k] @ Sigma[k] @ torch.transpose(
self._C[k], -2, -1)
# TODO before wat squeeze()
m = multivariate_normal.MultivariateNormal(mean.squeeze(-1), cov)
return m.log_prob(c)
def sample_action(self,
pen_down_prob,
o_pi,
o_mu1,
o_mu2,
o_sigma1,
o_sigma2,
o_corr,
no_sigma=False):
# define action sampling as well
# pen_out = torch.where(pen_down_prob<0.5, torch.zeros(pen_down_prob.size()[0]).cuda(),torch.ones(pen_down_prob.size()[0]).cuda())
_, pen_out = torch.sort(pen_down_prob, 1, descending=True)
pen_out = pen_out[:, 0:1].float()
opi_max, opi_sorted = torch.sort(o_pi, 1, descending=True)
actions = []
for i in range(opi_max.size()[0]):
ind = opi_sorted[i, 0]
cur_mu1 = o_mu1[:, ind][i]
cur_mu2 = o_mu2[:, ind][i]
cur_sigma1 = o_sigma1[:, ind][i]
cur_sigma2 = o_sigma2[:, ind][i]
cur_corr = o_corr[:, ind][i]
# print(cur_mu1,cur_mu2,cur_sigma1,cur_sigma2,cur_corr)
sigma_up = [cur_sigma1**2, cur_corr * cur_sigma1 * cur_sigma2]
sigma_down = [cur_corr * cur_sigma1 * cur_sigma2, cur_sigma2**2]
if no_sigma:
sigma = torch.eye(2) #.cuda()
else:
sigma = torch.Tensor([sigma_up, sigma_down])
m = mn.MultivariateNormal(torch.Tensor([cur_mu1, cur_mu2]), sigma)
try:
action = m.sample().cuda()
except:
action = torch.Tensor([cur_mu1, cur_mu2]).cuda()
actions.append(action)
actions = torch.stack(actions, 0)
# print(type(pen_out))
actions = torch.cat((pen_out, actions), 1)
# print('actions',actions.shape)
return actions
def gaussian_neg_log_likelihood(parameters, target):
"""Negative log likelihood loss for the Gaussian distribution
B = batch size, D = dimension of target (num classes), N = ensemble size
Args:
parameters (torch.tensor((B, D)), torch.tensor((B, D))):
mean values and variances of y|x for every x in
batch.
target (torch.tensor((B, N, D))): sample from the normal
distribution, if not an ensemble prediction N=1.
"""
mean, var = parameters
loss = 0
for batch_index, (mean_b, cov_b) in enumerate(zip(mean, var)):
cov_mat_b = torch.diag(cov_b)
distr = torch_mvn.MultivariateNormal(mean_b, cov_mat_b)
log_prob = distr.log_prob(target[batch_index, :, :])
loss -= torch.mean(log_prob) / target.size(0)
return loss
def run_worker():
rpc.init_rpc(name=f"trainer_{config.rank}",
rank=config.rank,
world_size=config.world_size)
logger.info("Logger is set - training start")
# set default gpu device id
torch.cuda.set_device(config.gpus[0])
# set seed
np.random.seed(config.seed)
torch.manual_seed(config.seed)
torch.cuda.manual_seed_all(config.seed)
torch.backends.cudnn.benchmark = True
# get data with meta info
# input_size, input_channels, n_classes, train_data = utils.get_data(
# config.dataset, config.data_path, cutout_length=0, validation=False)
#
net_crit = nn.CrossEntropyLoss().to(device)
# model = SearchCNNController(input_channels, config.init_channels, n_classes, config.layers,
# net_crit, device_ids=config.gpus)
# model = model.to(device)
model = TrainerNet(net_crit)
# weights optimizer
# w_optim = torch.optim.SGD(model.weights(), config.w_lr, momentum=config.w_momentum,
# weight_decay=config.w_weight_decay)
w_optim = DistributedOptimizer(torch.optim.SGD,
model.weights(),
lr=config.w_lr,
momentum=config.w_momentum,
weight_decay=config.w_weight_decay)
# alphas optimizer
# alpha_optim = torch.optim.Adam(model.alphas(), config.alpha_lr, betas=(0.5, 0.999),
# weight_decay=config.alpha_weight_decay)
alpha_optim = DistributedOptimizer(torch.optim.Adam,
model.alphas(),
lr=config.alpha_lr,
betas=(0.5, 0.999),
weight_decay=config.alpha_weight_decay)
# split data to train/validation
n_train = len(train_data)
split = n_train // 2
world = config.world_size
rank = config.rank
indices = list(range(n_train))
train_sampler = torch.utils.data.sampler.SubsetRandomSampler(
indices[int(rank * split / world):int((rank + 1) * split / world)])
valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(
indices[split + int(rank * (n_train - split) / world):split +
int(int((rank + 1) * (n_train - split) / world))])
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=config.batch_size,
sampler=train_sampler,
num_workers=config.workers,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(train_data,
batch_size=config.batch_size,
sampler=valid_sampler,
num_workers=config.workers,
pin_memory=True)
# lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
# w_optim, config.epochs, eta_min=config.w_lr_min)
lrs_rrefs = []
for opt_rref in w_optim.remote_optimizers:
lrs_rrefs.append(
rpc.remote(opt_rref.owner(),
create_lr_scheduler,
args=(opt_rref, )))
v_model = SearchCNNController(input_channels,
config.init_channels,
n_classes,
config.layers,
nn.CrossEntropyLoss().to(device),
device_ids=config.gpus).to(device)
architect = Architect(model, v_model, config.w_momentum,
config.w_weight_decay, noise_add)
if noise_add:
logger.info("Adding noise")
for param in model.parameters():
shape_gaussian[param.data.shape] = gaussian.MultivariateNormal(
torch.zeros(param.data.shape), torch.eye(param.data.shape[-1]))
else:
logger.info("Not adding noise")
# training loop
best_top1 = 0.
for epoch in range(config.epochs):
with dist_autograd.context() as cid:
futs = []
for lrs_rref in lrs_rrefs:
futs.append(
rpc.rpc_async(lrs_rref.owner(),
lrs_step,
args=(lrs_rref, )))
[fut.wait() for fut in futs]
lr = remote_method(get_lrs_value,
lrs_rrefs.owner(),
args=(lrs_rrefs[0], ))
# lr_scheduler.step()
# lr = lr_scheduler.get_lr()[0]
# model.print_alphas(logger)
# training
train(train_loader, valid_loader, model, architect, w_optim,
alpha_optim, lr, epoch)
# validation
cur_step = (epoch + 1) * len(train_loader)
top1 = validate(valid_loader, model, epoch, cur_step)
# log
# genotype
genotype = model.genotype()
logger.info("genotype = {}".format(genotype))
# genotype as a image
plot_path = os.path.join(config.plot_path,
"EP{:02d}".format(epoch + 1))
caption = "Epoch {}".format(epoch + 1)
plot(genotype.normal, plot_path + "-normal", caption)
plot(genotype.reduce, plot_path + "-reduce", caption)
# save
if best_top1 < top1:
best_top1 = top1
best_genotype = genotype
is_best = True
else:
is_best = False
utils.save_checkpoint(model, config.path, is_best)
print("")
logger.info("Final best Prec@1 = {:.4%}".format(best_top1))
logger.info("Best Genotype = {}".format(best_genotype))
rpc.shutdown()
def DoModelInferenceForFrame(agentsObservedOnFrame, forcedHistoryDict=None):
global history_pos
# Update the history if forced param is used
if forcedHistoryDict != None:
for key, value in forcedHistoryDict.items():
assert isinstance(
value, np.ndarray), "value is not instance of numpy array"
assert value.shape is not (NUM_FRAMES_TO_OBSERVE, 2)
history_pos[key] = value
# Update the history of agents seen with the new observed values
for key, value in agentsObservedOnFrame.items():
# If agent was not already in the history pos, init everything with local value
if key not in history_pos:
history_pos[key] = np.tile(value, [NUM_FRAMES_TO_OBSERVE, 1])
else: # Else, just put his new pos in the end of history
values = history_pos[key]
values[0:NUM_FRAMES_TO_OBSERVE -
1] = values[1:NUM_FRAMES_TO_OBSERVE]
values[NUM_FRAMES_TO_OBSERVE - 1] = value
# Do simulation using the model
# ------------------------------------------
# Step 1: fill the input
numAgentsThisFrame = len(agentsObservedOnFrame)
# Absolute observed trajectories
obs_traj = np.zeros(shape=(NUM_FRAMES_TO_OBSERVE, numAgentsThisFrame, 2),
dtype=np.float32)
# Zero index is main, others are following
obs_traj[:, MAIN_AGENT_INDEX, :] = history_pos[MAIN_AGENT_NAME]
index = 1
indexToAgentNameMapping = {}
indexToAgentNameMapping[MAIN_AGENT_INDEX] = MAIN_AGENT_NAME
for key, value in agentsObservedOnFrame.items():
if key != MAIN_AGENT_NAME:
obs_traj[:, index, :] = history_pos[key]
indexToAgentNameMapping[index] = key
index += 1
# Relative observed trajectories
obs_traj_rel = np.zeros(shape=(NUM_FRAMES_TO_OBSERVE, numAgentsThisFrame,
2),
dtype=np.float32)
obs_traj_rel[1:, :, :] = obs_traj[1:, :, :] - obs_traj[:-1, :, :]
seq_start_end = np.array([[0, numAgentsThisFrame]
]) # We have only 1 batch containing all agents
# Transform them to torch tensors
obs_traj = torch.from_numpy(obs_traj).type(torch.float)
obs_traj_rel = torch.from_numpy(obs_traj_rel).type(torch.float)
# Permute to the normal input size of the models [num_agents, 2, num_frames]
obs_traj = obs_traj.permute(1, 2, 0).unsqueeze(0)
obs_traj_rel = obs_traj_rel.permute(1, 2, 0).unsqueeze(0)
seq_start_end = torch.from_numpy(seq_start_end)
V_obs, A_obs = utils.seq_to_graph(obs_traj, obs_traj_rel)
if len(V_obs.shape) < 4:
V_obs = V_obs.unsqueeze(0)
if (len(A_obs.shape) < 4):
A_obs = A_obs.unsqueeze(0)
V_obs_tmp = V_obs.permute(0, 3, 1, 2)
V_obs_tmp = V_obs_tmp.cuda()
A_obs_tmp = A_obs.squeeze(0).cuda()
V_pred, _ = generator(V_obs_tmp, A_obs_tmp)
V_pred = V_pred.permute(0, 2, 3, 1)
V_pred = V_pred.squeeze(0)
num_of_objs = obs_traj_rel.shape[1]
V_pred = V_pred[:, :num_of_objs, :]
sx = torch.exp(V_pred[:, :, 2]) # sx
sy = torch.exp(V_pred[:, :, 3]) # sy
corr = torch.tanh(V_pred[:, :, 4]) # corr
cov = torch.zeros(V_pred.shape[0], V_pred.shape[1], 2, 2).cuda()
cov[:, :, 0, 0] = sx * sx
cov[:, :, 0, 1] = corr * sx * sy
cov[:, :, 1, 0] = corr * sx * sy
cov[:, :, 1, 1] = sy * sy
mean = V_pred[:, :, 0:2]
mvnormal = torchdist.MultivariateNormal(mean, cov)
V_pred = mvnormal.sample()
V_x = seq_to_nodes(obs_traj.data.cpu().numpy().copy())
V_x_rel_to_abs = nodes_rel_to_nodes_abs(
V_obs.data.cpu().numpy().squeeze(0).copy(), V_x[0, :, :].copy())
V_pred_rel_to_abs = nodes_rel_to_nodes_abs(
V_pred.data.cpu().numpy().copy(), V_x[-1, :, :].copy())
#pred_traj_fake_rel = generator(obs_traj, obs_traj_rel, seq_start_end)
#pred_traj_fake = relative_to_abs(pred_traj_fake_rel, obs_traj[-1])
# Take the first predicted position and add it to history
pred_traj_fake = V_pred_rel_to_abs
newMainAgentPos = pred_traj_fake[0][0] # Agent 0 is our main agent
return newMainAgentPos
def wikics(method, device, num_hid, lr, dropout, num_epochs, mask, beta, n_MC,
n_u, n_c, n_jobs):
# Create "data" and "results" folder
data_folder = osp.join(here, "data")
results_folder = osp.join(here, "results")
MLP_folder = osp.join(results_folder, "MLP")
SAGE_folder = osp.join(results_folder, "SAGE")
create_folder(data_folder)
create_folder(results_folder)
create_folder(MLP_folder)
create_folder(SAGE_folder)
# Get WikiCS dataset
dataset = WikiCS(root=data_folder)
data = dataset[0]
X = data.x.to(device)
y = data.y.to(device)
edge_index = data.edge_index.to(device)
# Get splits
train_idx = data.train_mask[:, mask]
val_idx = data.val_mask[:, mask]
test_idx = data.test_mask
# AWGN channel
Z = lambda n, sigma: mvn.MultivariateNormal(
torch.zeros(n), torch.diag(torch.ones(n) * sigma**2))
Z_ = Z(num_hid, beta)
# Instantiate the model, optimizer and criterion
if method == "MLP":
noisy_model = MLP(X.size(-1), num_hid, dataset.num_classes, Z_,
dropout, device).to(device)
elif method == "SAGE":
noisy_model = SAGE(X.size(-1), num_hid, dataset.num_classes, Z_,
dropout, device).to(device)
else:
raise ValueError("Invalid method")
optimizer = Adam(noisy_model.parameters(), lr=lr)
criterion = F.nll_loss
# Train model
losses, train_accs, val_accs, test_accs, I_1, I_2 = [], [], [], [], [], []
for epoch in range(1, 1 + num_epochs):
if method == "MLP":
loss = train(noisy_model, X, y, train_idx, optimizer, criterion)
train_acc, val_acc, test_acc = test(noisy_model, X, y, train_idx,
val_idx, test_idx)
u_S, c_S = get_samples(noisy_model, X[train_idx])
noisy_model.eval()
z, _, _ = noisy_model(X[train_idx])
elif method == "SAGE":
loss = train(noisy_model, X, y, train_idx, optimizer, criterion,
edge_index)
train_acc, val_acc, test_acc = test(noisy_model, X, y, train_idx,
val_idx, test_idx, edge_index)
u_S, c_S = get_samples(noisy_model, X, edge_index, train_idx)
noisy_model.eval()
z, _, _ = noisy_model(X, edge_index)
z = z[train_idx]
kde = KDE(n_jobs=n_jobs)
I_T_1 = compute_I(u_S[0], c_S[0], Z_, n_MC, kde, device, n_u, n_c)
I_T_2 = compute_I(u_S[1], c_S[1], Z_, n_MC, kde, device, n_u, n_c)
print(f"Epoch: {epoch:02d}", f"Loss: {loss:.4f}",
f"Train: {100 * train_acc:.2f}%", f"Valid: {100 * val_acc:.2f}%",
f"Test: {100 * test_acc:.2f}%", f"I(X,T_1]) = {I_T_1:.4f}",
f"I(X,T_2) = {I_T_2:.4f}")
losses.append(loss)
train_accs.append(train_acc)
val_accs.append(val_acc)
test_accs.append(test_acc)
I_1.append(I_T_1)
I_2.append(I_T_2)
# Create a folder for every epoch to store checkpoint and t-SNE embeddings
chkpt = {
"epoch": epoch,
"state_dict": noisy_model.state_dict(),
"loss": loss,
"train accuracy": 100 * train_acc,
"valid accuracy": 100 * val_acc,
"test accuracy": 100 * test_acc,
"I(X:T_1)": I_T_1,
"I(X:T_2)": I_T_2
}
if method == "MLP":
epoch_folder = osp.join(MLP_folder, f"{epoch}")
elif method == "SAGE":
epoch_folder = osp.join(SAGE_folder, f"{epoch}")
create_folder(epoch_folder)
torch.save(chkpt, osp.join(epoch_folder, f"{epoch}.pt"))
x_coord, y_coord = zip(*TSNE().fit_transform(z.cpu().numpy()))
colormap = np.array(Category20_20 + Category20b_20 + Accent8)
labels = np.array([int(l) for l in data.y[train_idx]])
plt.scatter(x_coord, y_coord, c=colormap[labels])
plt.savefig(osp.join(epoch_folder, f"{epoch}.png"))
plt.close()
return losses, train_accs, val_accs, test_accs, I_1, I_2
decoder.module.load_state_dict(checkpointD['decoder_state_dict'])
decoder.to(device)
encoder.eval()
decoder.eval()
Z = list()
with torch.no_grad():
for i in range(class_num):
Z.append(torch.zeros((1, z_dim), dtype=torch.float))
for i, (img, label) in enumerate(train_loader):
img = img.to(device)
mu, log_sigmoid = encoder(img)
z = reparameterization(mu, log_sigmoid)
z = z.view(-1, 1, z_dim)
Z = batch2one(Z, label, z, class_num)
N = []
for i in range(class_num):
label_mean = torch.mean(Z[i][1:], dim=0).double()
label_cov = torch.from_numpy(np.cov(Z[i][1:].numpy(),
rowvar=False)).double()
m = mn.MultivariateNormal(label_mean, label_cov)
sample = m.sample((2048, ))
sample = sample.float()
sample = sample.to(device).view(z_dim, -1)
fake = decoder(sample)
N.append(m)
torch.save({'distribution': N}, model_save_path + "class_distribution.dt")
def forward(self, obs, deterministic=False, with_logprob=True):
"""Perform forward pass through the network.
Args:
obs (torch.Tensor): The tensor of observations.
deterministic (bool, optional): Whether we want to use a deterministic
policy (used at test time). When true the mean action of the stochastic
policy is returned. If false the action is sampled from the stochastic
policy. Defaults to False.
with_logprob (bool, optional): Whether we want to return the log probability
of an action. Defaults to True.
Returns:
torch.Tensor, torch.Tensor: The actions given by the policy, the log
probabilities of each of these actions.
"""
# Create base distribution
base_distribution = mn.MultivariateNormal(torch.zeros(self.a_dim),
torch.eye(self.a_dim))
epsilon = base_distribution.sample((obs.shape[0], ))
# Calculate required variables
net_out = self.net(obs)
mu = self.mu_layer(net_out)
log_sigma = self.log_sigma(net_out)
log_sigma = torch.clamp(log_sigma, self._log_std_min,
self._log_std_max)
sigma = torch.exp(log_sigma)
# Create bijection
squash_bijector = transforms.TanhTransform()
affine_bijector = transforms.ComposeTransform([
transforms.AffineTransform(mu, sigma),
])
# Calculate raw action
raw_action = bijector(epsilon)
# Check summing axis
sum_axis = 0 if obs.shape.__len__() == 1 else 1
# Pre-squash distribution and sample
# TODO: Check if this has the right size. LAC samples from base distribution
# Sample size is memory buffer size!
pi_distribution = Normal(mu, sigma)
raw_action = (
pi_distribution.rsample()
# DEBUG: The tensorflow implmentation samples
) # Sample while using the parameterization trick
# Compute log probability in squashed gaussian
if with_logprob:
# Compute logprob from Gaussian, and then apply correction for Tanh
# squashing. NOTE: The correction formula is a little bit magic. To get an
# understanding of where it comes from, check out the original SAC paper
# (arXiv 1801.01290) and look in appendix C. This is a more
# numerically-stable equivalent to Eq 21. Try deriving it yourself as a
# (very difficult) exercise. :)
logp_pi = pi_distribution.log_prob(raw_action).sum(axis=-1)
logp_pi -= (2 * (np.log(2) - raw_action -
F.softplus(-2 * raw_action))).sum(axis=1)
else:
logp_pi = None
# Calculate scaled action and return the action and its log probability
clipped_a = torch.tanh(
raw_action) # Squash gaussian to be between -1 and 1
# Get clipped mu
# FIXME: Is this okay LAC also squashes this output?!
# clipped_mu = torch.tanh(mu) # LAC version
clipped_mu = mu
# Return action and log likelihood
# Debug: The LAC expects a distribution we already return the log probabilities
return clipped_a, clipped_mu, logp_pi
def particlefilter(G, J, U, V, lam, r, y, P_process, P_obs, Np):
"""
particle filter function for estimating latent dynamics of the TAP brain
B, Ns, Nr, Ny, T = no. of batches, latent varibles, neurons, input variables, time steps
inputs:
G : message passing parameters
J : interaction matrix, tensor of shape Ns x Ns
U : embedding matrix. tensor of shape Nr x Ns
V : input mapping matrix, tensor of shape Ns x Ny
lam: low pass filtering constant for TAP dynamics
r : tensor of shape B x Nr x T
y : tensor of shape B x Ny (Ny = no. of input variables)
P_process : inverse of covariance of process noise, tensor of shape Ns x Ns
P_obs : inverse of covariance of observation noise, tensor of shape Nr x Nr
Np : no. of particles to use
outputs:
LL : observed data log likelihood, tensor of shape B
xhat: estimated latent dynamics, tensor of shape B x Ns x T
ParticlesAll : particle trajectories, tensor of shape B x Ns x Np x T
WVec:particle weights, tensor of shape B x Np
"""
if len(r.shape) < 3:
r.unsqueeze_(0) # this is to ensure shape is B x Nr x T
y.unsqueeze_(0) # this is to ensure shape is B x Ny x T
B, Nr, T = r.shape # no. of batches, no. of neurons, no. of time steps
Ns, Ny = J.shape[0], y.shape[0] # no. of latent variables, no. of inputs
device = r.device
dtype = r.dtype
# Compute the inverse covariance of the proposal distribution q = p(x_t | x_(t-1), r_t)
# Notation: Q for covariance, P for inverse covariance
P_1 = torch.mm(P_obs, U) # intermediate matrix 1
P_2 = torch.mm(U.t(), P_1) # intermediate matrix 2
P_proposal = P_process + P_2
P_proposal = (P_proposal +
P_proposal.t()) / 2 # make it a perfectly symmetric matrix
Q_proposal = P_proposal.inverse()
# Define the noise processes
mvn_process = MVN.MultivariateNormal(loc=torch.zeros(Ns,
device=device,
dtype=dtype),
precision_matrix=P_process)
mvn_proposal = MVN.MultivariateNormal(loc=torch.zeros(Ns,
device=device,
dtype=dtype),
precision_matrix=P_proposal)
# Generate initial particles X_0 ~ N( pinv(U)r_0, Q_process )
# At the time step zero generate K x Np particles and pick Np particles with highest weights
K = 10
r0 = r[..., 0]
x = torch.matmul(torch.pinverse(U), r0.unsqueeze(2)) + mvn_process.rsample(
sample_shape=torch.Size([B, K * Np])).permute(0, 2, 1)
# Compute the initial weights of the particles
P_3 = torch.mm(torch.pinverse(U).t(), P_process) # intermediate matrix 3
logWVec = -0.5 * ((x * torch.matmul(P_2, x)).sum(dim=1) - 2 *
(r0.unsqueeze(2) * torch.matmul(P_1, x)).sum(dim=1))
logWVec += 0.5 * ((x * torch.matmul(P_process, x)).sum(dim=1) - 2 *
(r0.unsqueeze(2) * torch.matmul(P_3, x)).sum(dim=1))
log_e = torch.max(logWVec, dim=1)[0] # find maximum log weight
logWVec -= log_e.unsqueeze(1) # subtract the maximum
# retain only the Np best particles
for b in range(B):
idx = torch.argsort(logWVec[b], descending=True)
logWVec[b] = logWVec[b, idx]
x[b] = x[b, :, idx]
logWVec, x = logWVec[:, 0:Np], x[..., 0:Np]
# normalized initial weights
WVec = torch.exp(logWVec)
WVec = WVec / torch.sum(WVec, dim=1).unsqueeze(1)
ParticlesAll = torch.zeros((B, Ns, Np, T), device=device, dtype=dtype)
ParticlesAll[..., 0] = x
# normalization constant for the weights
# log_nu = 0.5*np.log(np.linalg.det(P_process.data.numpy())) + 0.5*np.log(np.linalg.det(P_obs.data.numpy())) -0.5*np.log(np.linalg.det(P_proposal.data.numpy()))
# log_nu += -0.5*Nr*np.log(2*np.pi)
log_nu = 0
# # Compute log p(r_0) to initialize the observed data log likelihood
# P_4 = torch.mm(P_1, torch.mm(P_2.inverse(),P_1.t()))
# LL = -0.5*(Nr - Ns)*np.log(2*np.pi) + 0.5*np.log(np.linalg.det(P_obs.data.numpy())) - 0.5*np.log(np.linalg.det(P_2.data.numpy()))
# if B == 1:
# LL += -0.5*torch.matmul(r0.unsqueeze(1),torch.matmul(P_obs - P_4,r0.unsqueeze(2))).item()
# else:
# LL += -0.5*torch.matmul(r0.unsqueeze(1),torch.matmul(P_obs - P_4,r0.unsqueeze(2))).squeeze()
LL = 0
#savedparticles = []
#savedparticles.append(x[0].data.numpy())
for tt in range(1, T):
# resample particles based on their weights if sample diversity is low
ESS = 1 / torch.sum(WVec**2, dim=1)
for b in range(B):
if ESS[b] < Np / 2 and tt != T - 1:
idx = resampleSystematic_torch(WVec[b], Np, device, dtype)
ParticlesAll[b] = ParticlesAll[b, :, idx]
x = ParticlesAll[..., tt - 1]
yt = y[..., tt - 1]
rt = r[..., tt]
Minvr = torch.matmul(P_obs, rt.unsqueeze(2)) # size B x Nr x 1
rMinvr = torch.matmul(rt.unsqueeze(1), Minvr) # size B x 1 x 1
UMinvr = torch.matmul(U.t(), Minvr) # size B x Ns x 1
f_tap = TAPnonlinearity(x, yt.unsqueeze(2), G, J, V, lam)
Pinvf_tap = torch.matmul(P_process, f_tap)
v = Pinvf_tap + UMinvr
mu_proposal = torch.matmul(Q_proposal,
v) # mean of the proposal distribution
# sample new particles from proposal distribution
ParticlesNew = mu_proposal + mvn_proposal.rsample(
sample_shape=torch.Size([B, Np])).permute(0, 2, 1)
# log of incremental weights
log_alpha = log_nu - 0.5 * (rMinvr.squeeze(2) + torch.sum(
f_tap * Pinvf_tap - v * mu_proposal, dim=1))
# update log weights
logWVec = torch.log(WVec) + log_alpha
log_e = torch.max(logWVec, dim=1)[0] # find maximum log weight
logWVec -= log_e.unsqueeze(1) # subtract the maximum
# unnormalized weights
WVec = torch.exp(logWVec)
# update log likelihood
LL += torch.log(torch.sum(WVec, dim=1)) + log_e
# normalize the weights
WVec = WVec / torch.sum(WVec, dim=1).unsqueeze(1)
# append particles
ParticlesAll[..., tt] = ParticlesNew
#savedparticles.append(np.copy(ParticlesAll[0,:,:,0:tt+1].data.numpy()))
#savedparticles.append(ParticlesNew[0].data.numpy())
xhat = torch.sum(ParticlesAll * WVec.view(B, 1, Np, 1), dim=2).squeeze(2)
return LL, xhat, ParticlesAll, WVec
def create_full(shape):
mean = torch.zeros(shape)
cov = torch.ones((shape, shape))
D = N.MultivariateNormal(mean, cov)
return D
def test(KSTEPS=20): # K=20
global loader_test, model
model.eval()
ade_bigls = []
fde_bigls = []
ade_bigls_mean = []
fde_bigls_mean = []
raw_data_dict = {}
step = 0
for batch in loader_test:
step += 1
# Get data
batch = [tensor.cuda() for tensor in batch]
obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_gt_rel, non_linear_ped, \
loss_mask, V_obs, A_obs, V_tr, A_tr, obs_classes = batch
num_of_objs = obs_traj_rel.shape[1]
# Forward
V_obs_tmp = V_obs.permute(0, 3, 1, 2)
V_pred, _ = model(V_obs_tmp, A_obs.squeeze(), obs_classes)
V_pred = V_pred.permute(0, 2, 3, 1)
V_tr = V_tr.squeeze()
A_tr = A_tr.squeeze()
V_pred = V_pred.squeeze()
num_of_objs = obs_traj_rel.shape[1]
V_pred, V_tr = V_pred[:, :num_of_objs, :], V_tr[:, :num_of_objs, :]
# For now I have my bi-variate parameters
sx = torch.exp(V_pred[:, :, 2]) # sx
sy = torch.exp(V_pred[:, :, 3]) # sy
corr = torch.tanh(V_pred[:, :, 4]) # corr
cov = torch.zeros(V_pred.shape[0], V_pred.shape[1], 2, 2).cuda()
cov[:, :, 0, 0] = sx * sx
cov[:, :, 0, 1] = corr * sx * sy
cov[:, :, 1, 0] = corr * sx * sy
cov[:, :, 1, 1] = sy * sy
mean = V_pred[:, :, 0:2]
mvnormal = torchdist.MultivariateNormal(mean, cov)
# Now sample 20 samples
ade_ls = {}
fde_ls = {}
V_x = seq_to_nodes(obs_traj.data.cpu().numpy().copy())
V_x_rel_to_abs = nodes_rel_to_nodes_abs(
V_obs.data.cpu().numpy().squeeze().copy(), V_x[0, :, :].copy())
V_y = seq_to_nodes(pred_traj_gt.data.cpu().numpy().copy())
V_y_rel_to_abs = nodes_rel_to_nodes_abs(
V_tr.data.cpu().numpy().squeeze().copy(), V_x[-1, :, :].copy())
raw_data_dict[step] = {}
raw_data_dict[step]['obs'] = copy.deepcopy(V_x_rel_to_abs)
raw_data_dict[step]['trgt'] = copy.deepcopy(V_y_rel_to_abs)
raw_data_dict[step]['pred'] = []
for n in range(num_of_objs):
ade_ls[n] = []
fde_ls[n] = []
for k in range(KSTEPS):
V_pred = mvnormal.sample()
V_pred_rel_to_abs = nodes_rel_to_nodes_abs(
V_pred.data.cpu().numpy().squeeze().copy(),
V_x[-1, :, :].copy())
raw_data_dict[step]['pred'].append(
copy.deepcopy(V_pred_rel_to_abs))
for n in range(num_of_objs):
pred = []
target = []
obsrvs = []
number_of = []
pred.append(V_pred_rel_to_abs[:, n:n + 1, :] *
10) # scaling factor
target.append(V_y_rel_to_abs[:, n:n + 1, :] * 10)
obsrvs.append(V_x_rel_to_abs[:, n:n + 1, :] * 10)
number_of.append(1)
ade_ls[n].append(ade(pred, target, number_of))
fde_ls[n].append(fde(pred, target, number_of))
for n in range(num_of_objs):
ade_bigls.append(min(ade_ls[n]))
fde_bigls.append(min(fde_ls[n]))
ade_bigls_mean.append(sum(ade_ls[n]) / len(ade_ls[n]))
fde_bigls_mean.append(sum(fde_ls[n]) / len(fde_ls[n]))
ade_ = sum(ade_bigls) / len(ade_bigls)
fde_ = sum(fde_bigls) / len(fde_bigls)
aade_ = sum(ade_bigls_mean) / len(ade_bigls_mean)
afde_ = sum(fde_bigls_mean) / len(fde_bigls_mean)
return ade_, fde_, aade_, afde_, raw_data_dict
#
torch.manual_seed(0)
# First generate R D-dimensional data points from p(x|alpha) when alpha = 1
alpha = 1
D = 2
R = 10
gamma = 1.2
x_mean = torch.zeros(D)
x_mean[0] = alpha
x_samples = multivariate_normal.MultivariateNormal(x_mean,
torch.eye(D)).sample([R])
# exact target posterior p(alpha|x_samples)
mu = torch.mean(x_samples[:, 0])
sigma = 1 / R**0.5
target = normal.Normal(mu, sigma)
# score function of p(alpha|x_samples)
score_p = lambda alpha_: jacobian(target.log_prob, alpha_)[0]
sequence_len = 50
# score function for each individual likelihood contribution
def create_d(shape):
mean = torch.zeros(shape)
cov = torch.eye(shape)
D = N.MultivariateNormal(mean, cov)
return D
def kl_ind_standard(indG):
dimension = indG.event_shape[0]
delta_2_norm = torch.pow(torch.norm(indG.mean, dim=-1), 2)
log_det_ind = torch.sum(torch.log(indG.variance), dim=-1)
trace_term = torch.sum(indG.variance, dim=-1)
kl_score = 0.5 * (delta_2_norm + trace_term - log_det_ind - dimension)
return kl_score
#
if __name__ == '__main__':
from torch.distributions import independent
from torch.distributions import multivariate_normal, normal
loc = torch.rand(3)
scale = 2 + torch.rand(3)
xx = multivariate_normal.MultivariateNormal(loc, torch.eye(3) * scale)
xy = multivariate_normal.MultivariateNormal(torch.zeros(3), torch.eye(3))
print(loc)
print(scale)
d1 = independent.Independent(normal.Normal(loc, scale), 1)
xz = multivariate_normal.MultivariateNormal(
loc,
torch.eye(3) * torch.pow(scale, 2))
loc = torch.rand(3)
scale = 2 + torch.rand(3)
print(loc)
print(scale)
d2 = independent.Independent(normal.Normal(loc, scale), 1)
d3 = independent.Independent(normal.Normal(torch.zeros(3), torch.ones(3)),
1)
def sample_z_like(shape, scale=1., grad=True):
mean = torch.zeros(shape[1])
cov = torch.eye(shape[1])
D = N.MultivariateNormal(mean, cov)
z = D.sample((shape[0], )).cuda()
return scale * z
def test(KSTEPS=20):
global loader_test, model
model.eval()
ade_bigls = []
fde_bigls = []
raw_data_dict = {}
step = 0
for batch in loader_test:
step += 1
#Get data
batch = [tensor.cuda() for tensor in batch]
obs_traj, pred_traj_gt, obs_traj_rel, pred_traj_gt_rel, non_linear_ped,\
loss_mask,V_obs,A_obs,V_tr,A_tr = batch
# A_tr is actually never used.
num_of_objs = obs_traj_rel.shape[1]
#Forward
#V_obs = batch,seq,node,feat
#V_obs_tmp = batch,feat,seq,node
V_obs_tmp = V_obs.permute(0, 3, 1, 2)
V_pred, _ = model(V_obs_tmp, A_obs.squeeze())
# print(V_pred.shape)
# torch.Size([1, 5, 12, 2])
# torch.Size([12, 2, 5])
V_pred = V_pred.permute(0, 2, 3, 1)
# torch.Size([1, 12, 2, 5])>>seq,node,feat
# V_pred= torch.rand_like(V_tr).cuda()
V_tr = V_tr.squeeze()
A_tr = A_tr.squeeze()
V_pred = V_pred.squeeze()
num_of_objs = obs_traj_rel.shape[1]
V_pred, V_tr = V_pred[:, :num_of_objs, :], V_tr[:, :num_of_objs, :]
#print(V_pred.shape)
#For now I have my bi-variate parameters
#normx = V_pred[:,:,0:1]
#normy = V_pred[:,:,1:2]
sx = torch.exp(V_pred[:, :, 2]) #sx
sy = torch.exp(V_pred[:, :, 3]) #sy
corr = torch.tanh(V_pred[:, :, 4]) #corr
cov = torch.zeros(V_pred.shape[0], V_pred.shape[1], 2, 2).cuda()
cov[:, :, 0, 0] = sx * sx
cov[:, :, 0, 1] = corr * sx * sy
cov[:, :, 1, 0] = corr * sx * sy
cov[:, :, 1, 1] = sy * sy
mean = V_pred[:, :, 0:2]
mvnormal = torchdist.MultivariateNormal(mean, cov)
### Rel to abs
##obs_traj.shape = torch.Size([1, 6, 2, 8]) Batch, Ped ID, x|y, Seq Len
#Now sample 20 samples
ade_ls = {}
fde_ls = {}
V_x = seq_to_nodes(obs_traj.data.cpu().numpy().copy())
V_x_rel_to_abs = nodes_rel_to_nodes_abs(
V_obs.data.cpu().numpy().squeeze().copy(), V_x[0, :, :].copy())
V_y = seq_to_nodes(pred_traj_gt.data.cpu().numpy().copy())
V_y_rel_to_abs = nodes_rel_to_nodes_abs(
V_tr.data.cpu().numpy().squeeze().copy(), V_x[-1, :, :].copy())
raw_data_dict[step] = {}
raw_data_dict[step]['obs'] = copy.deepcopy(V_x_rel_to_abs)
raw_data_dict[step]['trgt'] = copy.deepcopy(V_y_rel_to_abs)
raw_data_dict[step]['pred'] = []
for n in range(num_of_objs):
ade_ls[n] = []
fde_ls[n] = []
for k in range(KSTEPS):
V_pred = mvnormal.sample()
#V_pred = seq_to_nodes(pred_traj_gt.data.numpy().copy())
V_pred_rel_to_abs = nodes_rel_to_nodes_abs(
V_pred.data.cpu().numpy().squeeze().copy(),
V_x[-1, :, :].copy())
raw_data_dict[step]['pred'].append(
copy.deepcopy(V_pred_rel_to_abs))
# print(V_pred_rel_to_abs.shape) #(12, 3, 2) = seq, ped, location
for n in range(num_of_objs):
pred = []
target = []
obsrvs = []
number_of = []
pred.append(V_pred_rel_to_abs[:, n:n + 1, :])
target.append(V_y_rel_to_abs[:, n:n + 1, :])
obsrvs.append(V_x_rel_to_abs[:, n:n + 1, :])
number_of.append(1)
ade_ls[n].append(ade(pred, target, number_of))
fde_ls[n].append(fde(pred, target, number_of))
for n in range(num_of_objs):
ade_bigls.append(min(ade_ls[n]))
fde_bigls.append(min(fde_ls[n]))
ade_ = sum(ade_bigls) / len(ade_bigls)
fde_ = sum(fde_bigls) / len(fde_bigls)
return ade_, fde_, raw_data_dict
