def posterior()

in models/s2s_big_hier.py [0:0]

• 45 lines of code
• 9 McCabe index (conditional complexity)

    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 == 4:
                        out = layer(out)

                    elif layer_idx == 3:
                        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=4)
                        z1 = z1.view(b, t, c, z1.shape[-2], z1.shape[-1])
                        out = torch.cat([out, z1], 2)
                        out = layer(out)

                    elif layer_idx == 2:
                        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=8)
                        z1 = z1.view(b, t, c, z1.shape[-2], z1.shape[-1])
                        z2 = dists[1][-2]
                        b, t, c, h, w = z2.shape
                        z2 = z2.view(b*t, c, h, w)
                        z2 = F.interpolate(z2, scale_factor=2)
                        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