def forward(self, context=None, x_prime=None, more_context=None, z=None):
"""
:param x: observation at current step
:param context: to be fed at each timestep
:param x_prime: observation at next step
:param more_context: also to be fed at each timestep.
:param z: (optional) if not None z is used directly and not sampled
:return:
"""
# TODO to get rid of more_context, make an interface that allows context structures
output = AttrDict()
output.p_z = self.prior(torch.zeros_like(
x_prime)) # the input is only used to read the batchsize atm
if x_prime is not None:
output.q_z = self.inf(
self.inf_lstm(x_prime, context, more_context).output)
if z is None:
if self._sample_prior:
z = Gaussian(output.p_z).sample()
else:
z = Gaussian(output.q_z).sample()
pred_input = [z, context, more_context]
output.x = self.gen_lstm(*pred_input).output
return output
def forward(self,
x,
context=None,
x_prime=None,
more_context=None,
z=None):
"""
:param x: observation at current step
:param context: to be fed at each timestep
:param x_prime: observation at next step
:param more_context: also to be fed at each timestep.
:param z: (optional) if not None z is used directly and not sampled
:return:
"""
# TODO to get rid of more_context, make an interface that allows context structures
output = AttrDict()
if x_prime is None:
x_prime = torch.zeros_like(x) # Used when sequence isn't available
output.q_z = self.inf(
self.inf_lstm(x_prime, context, more_context).output)
output.p_z = self.prior(
x) # the input is only used to read the batchsize atm
if z is not None:
if self._hp.prior_type == 'learned':
z = output.p_z.reparametrize(z)
pass # use z directly
elif self._sample_prior:
z = output.p_z.sample()
else:
z = output.q_z.sample()
# Note: providing x might be unnecessary if it is provided in the init_state
pred_input = [x, z, context, more_context]
# x_t is fed back in as input (technically, it is unnecessary, however, the lstm is setup to observe a frame
# every step because it observes one in the beginning).
output.x = self.gen_lstm(*pred_input).output
return output
def forward(self, context=None, more_context=None, z=None, batch_size=None):
"""
:param x: observation at current step
:param context: to be fed at each timestep
:param x_prime: observation at next step
:param more_context: also to be fed at each timestep.
:param z: (optional) if not None z is used directly and not sampled
:return:
"""
# TODO to get rid of more_context, make an interface that allows context structures
outputs = AttrDict()
outputs.p_z = self.prior(self.gen_lstm.output.weight.new_zeros(batch_size)) # the input is only used to read the batchsize atm
if z is None:
z = Gaussian(outputs.p_z).sample()
pred_input = [z, context, more_context]
outputs.x = self.gen_lstm(*pred_input).output
return outputs
def produce_subgoal(self,
inputs,
layerwise_inputs,
start_ind,
end_ind,
left_parent,
right_parent,
depth=None):
"""
Divides the subsequence by producing a subgoal inside it.
This function represents one step of recursion of the model
"""
subgoal = AttrDict()
e_l = left_parent.e_g_prime
e_r = right_parent.e_g_prime
subgoal.p_z = self.prior(e_l, e_r)
if 'z' in layerwise_inputs:
z = layerwise_inputs.z
if self._hp.prior_type == 'learned': # reparametrize if learned prior is used
z = subgoal.p_z.reparametrize(z)
elif self._sample_prior:
z = subgoal.p_z.sample()
else:
## Inference
if self._hp.attentive_inference:
subgoal.update(
self.inference(inputs, e_l, e_r, start_ind, end_ind))
else:
subgoal.match_timesteps = self.binding.comp_timestep(
left_parent.match_timesteps, right_parent.match_timesteps)
subgoal.update(
self.inference(inputs, e_l, e_r, start_ind, end_ind,
subgoal.match_timesteps.float()))
z = subgoal.q_z.sample()
## Predict the next node
pred_input = [e_l, e_r, z]
if self._hp.context_every_step:
mult = int(z.shape[0] / inputs.e_0.shape[0])
pred_input += [
inputs.e_0.repeat_interleave(mult, 0),
inputs.e_g.repeat_interleave(mult, 0)
]
if self._hp.tree_lstm:
if left_parent.hidden_state is None and right_parent.hidden_state is None:
left_parent.hidden_state, right_parent.hidden_state = self.lstm_initializer(
e_l, e_r, z)
subgoal.hidden_state, subgoal.e_g_prime = \
self.subgoal_pred(left_parent.hidden_state, right_parent.hidden_state, *pred_input)
else:
subgoal.e_g_prime_preact = self.subgoal_pred(*pred_input)
subgoal.e_g_prime = torch.tanh(subgoal.e_g_prime_preact)
subgoal.ind = (
start_ind + end_ind
) / 2 # gets overwritten w/ argmax of matching at training time (in loss)
return subgoal, left_parent, right_parent
def produce_subgoal(self, inputs, layerwise_inputs, start_ind, end_ind, left_parent, right_parent, depth=None):
"""
Divides the subsequence by producing a subgoal inside it.
This function represents one step of recursion of the model
"""
subgoal = AttrDict()
batch_size = start_ind.shape[0]
e_l = left_parent.e_g_prime
e_r = right_parent.e_g_prime
subgoal.p_z = self.prior(e_l, e_r)
if 'z' in layerwise_inputs:
z = layerwise_inputs.z
if self._hp.prior_type == 'learned': # reparametrize if learned prior is used
z = subgoal.p_z.reparametrize(z)
elif self._sample_prior:
z = subgoal.p_z.sample()
else:
## Inference
if (self._hp.timestep_cond_attention or self._hp.forced_attention):
subgoal.fraction = self.fraction_pred(e_l, e_r)[..., -1] if self.predict_fraction else None
subgoal.match_timesteps = self.matcher.comp_timestep(left_parent.match_timesteps,
right_parent.match_timesteps,
subgoal.fraction[:,
None] if subgoal.fraction is not None else None)
subgoal.update(self.inference(inputs, e_l, e_r, start_ind, end_ind, subgoal.match_timesteps.float()))
else:
subgoal.update(self.inference(
inputs, e_l, e_r, start_ind, end_ind, attention_weights=layerwise_inputs.safe.attention_weights))
z = subgoal.q_z.sample()
## Predict the next node
pred_input = [e_l, e_r, z]
if self._hp.context_every_step:
mult = int(z.shape[0] / inputs.enc_e_0.shape[0])
pred_input += [inputs.enc_e_0.repeat_interleave(mult, 0),
inputs.enc_e_g.repeat_interleave(mult, 0)]
if self._hp.tree_lstm:
if left_parent.hidden_state is None and right_parent.hidden_state is None:
left_parent.hidden_state, right_parent.hidden_state = self.lstm_initializer(e_l, e_r, z)
if self._hp.lstm_warmup_cycles > 0:
for _ in range(self._hp.lstm_warmup_cycles):
left_parent.hidden_state, __ = \
self.subgoal_pred(left_parent.hidden_state, right_parent.hidden_state, e_l, e_r, z)
right_parent.hidden_state = left_parent.hidden_state.clone()
subgoal.hidden_state, subgoal.e_g_prime = \
self.subgoal_pred(left_parent.hidden_state, right_parent.hidden_state, *pred_input)
else:
subgoal.e_g_prime_preact = self.subgoal_pred(*pred_input)
subgoal.e_g_prime = torch.tanh(subgoal.e_g_prime_preact)
## Additional predicted values
if self.predict_fraction and not self._hp.timestep_cond_attention:
subgoal.fraction = self.fraction_pred(e_l, e_r, subgoal.e_g_prime)[..., -1] # remove unnecessary dim
# add attention target if trained with attention supervision
if self._hp.supervise_attn_weight > 0.0:
frac = subgoal.fraction[:, None] if 'fraction' in subgoal and subgoal.fraction is not None else None
subgoal.match_timesteps = self.matcher.comp_timestep(left_parent.match_timesteps,
right_parent.match_timesteps,
frac)
subgoal.ind = (start_ind + end_ind) / 2 # gets overwritten w/ argmax of matching at training time (in loss)
return subgoal, left_parent, right_parent
