Can you tell me glow's splitting about freia

[ardizzone]

Hi!

I don't find the GLOW code very easy to understand, but I hope the following clarifies things:
Either way, I should first say that the objective will be the same mathematically, splitting the latent vector is just for convenience, special sampling schemes, style transfer, etc.
It will not change the result, or the computation speed, as far as I know.

not merging the z_i at the end can be done, to have multiple outputs.
E.g. for the example from the tutorial (https://github.com/VLL-HD/FrEIA#convolutional-inn),
you could put

nodes.append(Ff.OutputNode(nodes[-1].out0, name='output Z2'))
nodes.append(Ff.OutputNode(split_node.out1, name='output Z1'))

conv_inn = Ff.ReversibleGraphNet(nodes)

instead of the last three lines. Then the model will return a list containing Z1 and Z2.
However, for the jacobian this is not possible, it cannot be decomposed this way, and is kept track of internally.
So the split objective would be something like

z1, z2 = conv_inn(input)
log_jac = conv_inn.log_jacobian(run_forward=False)
total_log_likelihood = z1.sum(1) + z2.sum(1) - log_jac
total_log_likelihood  = total_log_likelihood.mean() / n_dimensions

Does this answer the question? If I misunderstood the issue, please let me know!

[MokkeMeguru]

I want to say, we may need the layer like this.

class GlowSplitChannel(nn.Module):
    '''Splits along channels to produce two separate outputs (for skip connections
    and such).'''
    def __init__(self, dims_in):
        super().__init__()
        assert len(dims_in) == 1, "Use channel_merge_layer instead"
        self.channels = dims_in[0][0]

    def forward(self, x, rev=False):
        if rev:
            return [torch.cat(x, dim=1)]
        else:
            return [x[0][:, :self.channels//2], x[0][:, self.channels//2:]]

    def jacobian(self, x, rev=False):
        # TODO batch size
        # return 0
        # Here is Glow's Split Layer, I think...
        return gaussianize(x[0][:, :self.channels//2], x[0][:, self.channels//2:])

    def output_dims(self, input_dims):
        assert len(input_dims) == 1, "Use channel_merge_layer instead"
        return [[input_dims[0][0]//2, *input_dims[0][1:]],
                [input_dims[0][0] - input_dims[0][0]//2, *input_dims[0][1:]]]

[ardizzone]

Hi! Sorry for the late answer, it has been a bumpy year.

The only difference to the SplitChannel module (https://github.com/VLL-HD/FrEIA/blob/master/FrEIA/modules/graph_topology.py#L7), is the gaussianize function, right?
I guess this is in reference to https://github.com/openai/glow/blob/master/model.py#L577 ?

But it is a bit of a miunderstanding, their code works differently:
What they call 'mean' and 'logs' is our 's' and 't'. In our case, that is already handled elsewhere and included in the jacobian that comes back from the GlowCouplingBlock.
In the glow code, the call to return Z.gaussian_diag(mean, logs) at the end of split2d is the affine transformation.
The jacobian of the splitting operation itself it zero.
(We have checked the jacobian of a small network numerically, it seems correctly implemented.)

I'll close the issue for now, but feel free to re-open if you disagree of have more questions!
We are working a lot on FrEIA in the next few weeks, improving the documentation and API.