def loss(self,
inputs,
outputs,
extra_action=True,
first_image=True,
log_error_arr=False):
loss_gt = inputs.traj_seq
loss_pad_mask = inputs.pad_mask
if not first_image:
loss_gt = loss_gt[:, 1:]
loss_pad_mask = loss_pad_mask[:, 1:]
weights = broadcast_final(loss_pad_mask, inputs.traj_seq)
# Skip first frame
losses = self.nll(outputs.distr, loss_gt[:, 1:], weights[:, 1:],
log_error_arr)
if self._hp.regress_actions:
actions_pad_mask = inputs.pad_mask[:, :-1]
loss_actions = outputs.actions
if extra_action:
loss_actions = loss_actions[:, :-1]
weights = broadcast_final(actions_pad_mask, inputs.actions)
losses.dense_action_rec = NLL(self._hp.dense_action_rec_weight) \
(loss_actions, inputs.actions, weights=weights, reduction=[-1, -2], log_error_arr=log_error_arr)
return losses
def attention(self,
q,
k,
v,
nz_k,
start_ind,
end_ind,
dropout=None,
forced_attention_step=None,
attention_weights=None):
def tensor_product(key, sequence):
dims = list(range(len(list(sequence.shape)))[3:])
return (key[:, None] * sequence).sum(dim=dims)
scores = tensor_product(q, k) / math.sqrt(nz_k) * self.temperature
scores = MultiheadAttention.mask_out(scores, start_ind, end_ind)
scores = F.softmax(scores, dim=1)
if forced_attention_step is not None:
scores = torch.zeros_like(scores)
batchwise_assign(scores, forced_attention_step[:, 0].long(), 1.0)
if attention_weights is not None:
scores = attention_weights[..., None].repeat_interleave(
scores.shape[2], 2)
if dropout is not None and dropout.p > 0.0:
scores = dropout(scores)
return (broadcast_final(scores, v) * v).sum(dim=1), scores
def forward(self, input, actions):
net_outputs = self.net(input)
padded_actions = torch.nn.functional.pad(
actions, (0, 0, 0, net_outputs.shape[1] - actions.shape[1], 0, 0))
# TODO quite sure the concatenation is automatic
net_outputs = batch_apply(
self.ac_net,
torch.cat([net_outputs,
broadcast_final(padded_actions, input)],
dim=2))
return net_outputs
def loss(self, inputs, outputs, log_error_arr=False):
losses = self.decoder.loss(inputs,
outputs,
extra_action=False,
log_error_arr=log_error_arr)
# TODO don't place loss on the final image
weights = broadcast_final(inputs.pad_mask[:, 1:], outputs.p_z.mu)
losses.kl = KLDivLoss2(self._hp.kl_weight, breakdown=1, free_nats_per_dim=self._hp.free_nats)\
(outputs.q_z, outputs.p_z, weights=weights, log_error_arr=log_error_arr)
return losses
def loss(self, inputs, model_output):
losses = AttrDict()
# action prediction loss
n_actions = model_output.actions.shape[1]
losses.action_reconst = L2Loss(1.0)(model_output.actions, inputs.actions[:, :n_actions],
weights=broadcast_final(inputs.pad_mask[:, :n_actions], inputs.actions))
# compute total loss
#total_loss = torch.stack([loss[1].value * loss[1].weight for loss in losses.items()]).sum()
#losses.total = AttrDict(value=total_loss)
# losses.total = total_loss*torch.tensor(np.nan) # for checking if backprop works
return losses
def _get_lstm_inputs(self, root, inputs):
"""
:param root:
:return:
"""
device = inputs.reference_tensor.device
batch_size, time = self._hp.batch_size, self._hp.max_seq_len
fullseq_shape = [batch_size, time] + list(inputs.enc_e_0.shape[1:])
lstm_inputs = AttrDict()
# collect start and end indexes and values of all segments
e_0s = torch.zeros(fullseq_shape, dtype=torch.float32, device=device)
e_gs = torch.zeros(fullseq_shape, dtype=torch.float32, device=device)
start_inds, end_inds = torch.zeros((batch_size, time), dtype=torch.float32, device=device), \
torch.zeros((batch_size, time), dtype=torch.float32, device=device)
reset_indicator = torch.zeros((batch_size, time),
dtype=torch.uint8,
device=device)
for segment in root.full_tree(
): # traversing the tree in breadth-first order.
if segment.depth == 0: # if leaf-node
start_ind = torch.ceil(segment.start_ind).type(
torch.LongTensor)
end_ind = torch.floor(segment.end_ind).type(torch.LongTensor)
batchwise_assign(reset_indicator, start_ind, 1)
# TODO iterating over batch must be gone
for ex in range(self._hp.batch_size):
if start_ind[ex] > end_ind[ex]:
continue # happens if start and end floats have no int in between
e_0s[ex, start_ind[ex]:end_ind[ex] +
1] = segment.e_0[ex] # +1 for including end_ind frame
e_gs[ex, start_ind[ex]:end_ind[ex] + 1] = segment.e_g[ex]
start_inds[ex, start_ind[ex]:end_ind[ex] +
1] = segment.start_ind[ex]
end_inds[ex, start_ind[ex]:end_ind[ex] +
1] = segment.end_ind[ex]
# perform linear interpolation
time_steps = torch.arange(time, dtype=torch.float, device=device)
inter = (time_steps - start_inds) / (end_inds - start_inds + 1e-7)
lstm_inputs.reset_indicator = reset_indicator
lstm_inputs.cell_input = (e_gs - e_0s) * broadcast_final(inter,
e_gs) + e_0s
lstm_inputs.reset_input = torch.cat([e_gs, e_0s], dim=2)
return lstm_inputs
def loss(self, inputs, outputs, add_total=True):
losses = AttrDict()
# subgoal reconstruction loss
n_action_output = outputs.actions.shape[1]
loss_weights = broadcast_final(
outputs.pad_mask[:, :n_action_output],
inputs.actions) if 'pad_mask' in outputs else 1
losses.action_reconst = L2Loss(self._hp.action_rec_weight)(
outputs.actions,
outputs.action_targets[:, :n_action_output],
weights=loss_weights)
if self._hp.pred_states:
losses.state_reconst = L2Loss(self._hp.state_rec_weight)(
outputs.states, outputs.state_targets)
return losses
def mask_extra(seq):
return seq * broadcast_final(get_pad_mask(model_output.end_ind, N), seq)
def mean(self):
template = torch.tensor([128, 64, 32, 16, 8, 4, 2, 1])
p = self.p
value = broadcast_final(template.to(p.device).float()[None], p)
return (p * value.float()).sum(1) / 127.5 - 1
def mean(self):
template = torch.arange(256)
p = self.p
value = broadcast_final(template.to(p.device).float()[None], p)
return (p * value.float()).sum(1) / 127.5 - 1