Simplify models

[dario-coscia]

Hi! I open a discussion here for simplifying the model structure. Now for compositional models (i.e. models that are composed or stacked together), we perform many checks on the shapes by creating fictitious inputs. Examples of these models are DeepONet, FNO, LowRankNeuralOperator, MIONET, ....

Here is an example of what we do with FNO:

PINA/pina/model/fno.py

Line 235 in f0bff24

lifting_operator_out = lifting_net(

I think that this approach has a significant advantage: we can control the type of passed input. In version 0.1 where we only allowed for LabelTensor, this is what we needed to do to allow LabelTensor to work correctly. However, with the 0.2, this type of checking is limiting the usage. For example, in FNO let's suppose I work with a costume network as follows:

class LiftingNet(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.fc0 = nn.Linear(in_features=in_features, out_features=out_features)
    def forward(self, x):
        x = x.permute(0, 2, 3, 1)
        x = self.fc0(x)
        x = x.permute(0, 3, 1, 2)
        return x

this triggers the error but it is completely valid (I have been there, done that already 🙁).

I suggest removing the network shape checks in compositional models such to leave more freedom to the users to use them. What do you think?

[FilippoOlivo]

Hi, agree in removing the checks if they blocks some options that are in practice feasible and they may also improve the overall performance of the models

[dario-coscia]

@ndem0 @GiovanniCanali what do you think?

[GiovanniCanali]

I am not sure whether checks have to be removed or redesigned, but I agree on changes if the current usage is limited only to LabelTensors.