FrEIA.modules

[psteinb]

What?

There are many modules that are 2-3 years old, introduced for some specific purpose.
Sometimes it turned out that the concept didn't even work and the module is useless.
In some cases, they overlap in functionality with other modules.
Some of them also don't have docstrings or explanations what it does.

How?

Ensure that

1. all the modules are useful
2. the same functionality isn't implemented in two different modules
3. all modules have a useful docstring explaining the arguments
4. the modules work correctly

Go through all the modules listed in FrEIA/modules/__init__.py.
Check points 1-3 above, and perhaps 4 (overlap with unittests).

Then deprecate modules that are old or unnecessary, or can be combined with similar ones,
and add docstrings.
Delete the ones that were already deprecated with v0.2 (see FrEIA/modules/coupling_layers.py, lines 430-434)

[psteinb]

it might be worth considering to create a PR for this and to drop tags in the code, so people know where to direct their attention to.

[ardizzone]

Thank you! I will consider that.

[ardizzone]

Another point that came up:

• Adding an absttract base class for FrEIA.modules, to make it clear what methods have to be implemented if one writes their own invertible function.

[fdraxler]

I really like the ReversibleSequential returning a tuple of output, jacobian (set to None if expensive and not requested). Can we extend this to all modules?

So we could aim for a signature like this:

TensorOrTensorList = Union[Tensor, List[Tensor]]

def forward(self, x: TensorOrTensorList, c=None, rev=False, jac=True) -> Tuple[TensorOrTensorList, Union[TensorOrTensorList, None]]:
    ...
    if jac:
        return z, jac
    return z, None

Pro: This would also solve #7 and simplify #8, because there is no state to be stored between forward() and jacobian().
Con: This is a breaking change and people have to modify their self-implemented modules.

[ardizzone]

Let's think about this, it seems like a very good idea, except for ruining existing code if people update FrEIA.
(Also because the return of ReversibleGraphNet will be different, maybe even worse than just the FrEIA.modules changing).

But I admit the current system is pretty janky. (with the run_forward kwarg in .log_jacobian() etc.) It was the best we could do at the time without breaking code. But maybe it's the right time to just do it.

[ardizzone]

Also:

• add the GLOW invertible 1x1 conv:
  Missing invertible 1x1 conv from GLOW FrEIA#50

[wapu]

I am in favor of this breaking change. We are super inconsistent with the log-Jacobian functions and buffered values anyway. Also it seems that most other normalizing flow implementations chose that signature, so might make it easier for people to understand our API.

[fdraxler]

Any opinion on how do we handle multiple in- & outputs, i.e. split and merge modules? Do we make all modules to work with tuples of tensors everywhere? Then ReversibleSequential is inconsistent with the new signature.

Compromise: We enforce that a module spits out a tuple exactly when it got a tuple in (unit tests #4 for the win!)? This makes minimal breaking changes for users.
Then I would propose what I requested in vislearn/FrEIA#54

Tuples of Jacobian make no sense from a mathematic stand point, I think.

[ardizzone]

No, I think we keep tuples for in- and outut of the individual nodes i.e. return (out1, out2), jac (as you said, there can be only one jacobian). That's the way it is now, except that the (out1, out2) and the jac are returned by different methods.

I don't think the internal stuff necessarily has to be consistent with the Reversible*-interface.
And I am always weary of the kinds of solutions à la "maybe it's a tuple, maybe not, you just have to check and do something different depending" (danger for confusion, need to explain more instead of just saying "ART: Always Return Tuples", and so on)

[ardizzone]

Also, I don't understand:
What if i'm a module with multiple outputs.
I only get a single tensor from the module before me.
Therefore, a single tensor will be output? But I have multiple outputs and need to return a tuple?

[fdraxler]

We should also have the output_dims(input_dims) abstract method in InvertibleModule. We should add this, right?

[ardizzone]

Yes, that is true!

[fdraxler]

Good, I added it in my branch.

Maybe even slicker: Modules should not have a method output_dims, put should tell their output dimension to the InvertibleModule in the constructor. The output_dims method gets exactly the same shapes passed in as the constructor.
Disadvantage: All shape computations have to be done as the first thing in each constructor and we have to do another pass through all modules.

[wapu]

If it can't be missed by users implementing their own modules how and where to do that, I like the idea.

[ardizzone]

I think this is a very good idea, just let me know before you go ahead so that i can it around in all the modules.

[fdraxler]

Ok, the PR vislearn/FrEIA#69 is a proposal.

[ardizzone]

Did I understand correctly, the proposal is rejected?