ENSTA-U2IS-AI / torch-uncertainty

Open-source framework for uncertainty and deep learning models in PyTorch :seedling:
https://torch-uncertainty.github.io
Apache License 2.0
222 stars 16 forks source link

Pytorch Lightning Model Question #107

Open calvinp0 opened 1 week ago

calvinp0 commented 1 week ago

Hi!

I have built my model using Pytorch Lighntning, thus it has the functions training_step, validation_step etc. I attempted to follow the tutorial here: https://torch-uncertainty.github.io/auto_tutorials/tutorial_der_cubic.html#gathering-everything-and-training-the-model

But it errors with NotImplementedError: Module [CMPNNModel] is missing the required "forward" function (which I guess may be obvious). So does this mean to utilise this package I will need to change my model from a PyTorch Lightning one to a Torch one? Or have I done something incorrect.

Thank you!

o-laurent commented 1 week ago

Hi @calvinp0!

Thank you for your feedback!

You need to define the __forward__(self, x: Tensor) -> Tensor method of your Lightning module (as shown in the starter example). A Lightning module is an extension of an nn.Module, and therefore should include a __forward__. If you do so, you will be able to test and train your model with the RegressionRoutine.

However, if you use trainer.fit or trainer.test, it will not use your own loops but those of the RegressionRoutine, so it may not work depending on your model. In the general supervised case, we would advise wrapping a simple torch.nn.Module in the RegressionRoutine.

:construction: Our implementation for regression is still unstable but we have made progress in the soon-to-come 0.2.1 version that we will merge in the following days. Reach out and raise issues if you have other questions or concerns. :construction:

To read if you want to keep your LightningModule:

You won't have the computation of the metrics that come with the RegressionRoutine - you can directly use the DERLoss from torch_uncertainty.losses and the NormalInverseGammaLayer from torch_uncertainty.layers.distributions. But anyway, you will need the __forward__(self, x: Tensor) -> Tensor method.

In any case, don't hesitate to give us more details here or contact us through Discord

(written with @alafage)