in models/s2s_big_hier_128.py [0:0]
def posterior(self, x, ctx, use_mean=False):
dists = []
sto_branches = sorted(self.sto_branches.keys(), reverse=True)
for layer_idx in sto_branches:
# print(layer_idx)
# Find the corresopnding activations
out = x[layer_idx][:, self.n_ctx:].contiguous()
cur_ctx = ctx[layer_idx][:, :self.n_ctx].contiguous()
branch_layers = self.posterior_branches['layer_{}'.format(layer_idx)]
# Process the current branch
for branch_layer_idx, layer in enumerate(branch_layers):
#print(branch_layer_idx)
if isinstance(layer, layers.ConvLSTM):
# Get initial condition
cur_ctx = cur_ctx.view(cur_ctx.shape[0], -1, cur_ctx.shape[-2], cur_ctx.shape[-1])
cur_ctx = cur_ctx.unsqueeze(1)
cur_ctx = self.posterior_init_nets['layer_{}'.format(layer_idx)](cur_ctx)
cur_ctx = cur_ctx.squeeze(1)
# Forward LSTM
out = layer(out, torch.chunk(cur_ctx, 2, 1))
# Handcrafted rules for integrating the different z's
elif branch_layer_idx == 3:
if layer_idx == 19:
out = layer(out)
elif layer_idx == 10:
z1 = dists[0][-2]
b, t, c, h, w = z1.shape
z1 = z1.view(b*t, c, h, w)
z1 = F.interpolate(z1, scale_factor=16)
z1 = z1.view(b, t, c, z1.shape[-2], z1.shape[-1])
out = torch.cat([out, z1], 2)
out = layer(out)
elif layer_idx == 4:
z1 = dists[0][-2]
z2 = dists[1][-2]
b, t, c, h, w = z1.shape
z1 = z1.view(b*t, c, h, w)
z1 = F.interpolate(z1, scale_factor=64)
z1 = z1.view(b, t, c, z1.shape[-2], z1.shape[-1])
b, t, c, h, w = z2.shape
z2 = z2.view(b*t, c, h, w)
z2 = F.interpolate(z2, scale_factor=4)
z2 = z2.view(b, t, c, z2.shape[-2], z2.shape[-1])
out = torch.cat([out, z1, z2], 2)
out = layer(out)
else:
out = layer(out)
# Compute distribution stats
mean, var = torch.chunk(out, 2, 2)
# Softplus var
logvar = F.softplus(var).log()
# Generate sample from this distribution
z0 = flows.gaussian_rsample(mean, logvar, use_mean=use_mean)
dists.append([mean, logvar, z0, z0, None])
return dists