vislearn / FrEIA-community

Community-driven effort to improve and extend FrEIA
0 stars 0 forks source link

FrEIA.modules #3

Closed psteinb closed 3 years ago

psteinb commented 3 years ago

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 commented 3 years ago

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 commented 3 years ago

Thank you! I will consider that.

ardizzone commented 3 years ago

Another point that came up:

fdraxler commented 3 years ago

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 commented 3 years ago

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 commented 3 years ago

Also:

wapu commented 3 years ago

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 commented 3 years ago

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 https://github.com/VLL-HD/FrEIA/pull/54

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

ardizzone commented 3 years ago

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 commented 3 years ago

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 commented 3 years ago

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

ardizzone commented 3 years ago

Yes, that is true!

fdraxler commented 3 years ago

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 commented 3 years ago

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

ardizzone commented 3 years ago

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 commented 3 years ago

Ok, the PR https://github.com/VLL-HD/FrEIA/pull/69 is a proposal.

ardizzone commented 3 years ago

Did I understand correctly, the proposal is rejected?