Fix lightning issue

[LorenzLamm]

This PR should fix the issue mentioned in https://github.com/teamtomo/membrain-seg/issues/40

Thanks @alisterburt for providing the solution -- I just needed to add the strict=False argument, because somehow the information stored in the models was not compatible anymore (some weight of a Pytorch loss function).

[codecov-commenter]

Codecov Report

Merging #41 (8d3bea7) into main (c132a79) will increase coverage by 0.00%. The diff coverage is 0.00%.

:exclamation: Your organization needs to install the Codecov GitHub app to enable full functionality.

@@          Coverage Diff          @@
##            main     #41   +/-   ##
=====================================
  Coverage   6.61%   6.62%           
=====================================
  Files         38      38           
  Lines       1315    1314    -1     
=====================================
  Hits          87      87           
+ Misses      1228    1227    -1     

Files	Coverage Δ
src/membrain_seg/segmentation/segment.py	0.00% <0.00%> (ø)

[alisterburt]

@LorenzLamm this will likely break again at some point if they update the checkpoint metadata format again, if you search for how to upgrade model checkpoints to new lightning versions this is the more correct solution 🙂

[LorenzLamm]

@LorenzLamm this will likely break again at some point if they update the checkpoint metadata format again, if you search for how to upgrade model checkpoints to new lightning versions this is the more correct solution 🙂

Hi @alisterburt , Thanks a lot for the feedback. I looked into the error message after changing the model loading a bit further. FYI, this is the error message I receive: RuntimeError: Error(s) in loading state_dict for SemanticSegmentationUnet: Missing key(s) in state_dict: "loss_function.loss_fn.dice_loss.loss.class_weight".

The problem here is not on the Pytorch lightning side, as I had initially thought, but on the MONAI side, which also got an update 3 weeks ago (I feel every package got an update during my vacation :D). They added support for different class weights with the new update: https://github.com/Project-MONAI/MONAI/releases

One workaround is to manually set the class_weight variable, e.g. pl_model.loss_function.loss_fn.dice_loss.loss.class_weight = None but I feel this is not very elegant and may also cause problems in the future.

I did manage to convert the model file to a new version, containing the class_weight variable. With this, it runs smoothly with the new MONAI version. However, this is not backward-compatible: When using a previous MONAI version with the new model, it cannot handle this additional key in the state_dict.

In order to not overcomplicate this, I still vote for the initial fix with setting strict=False. What do you think?

[alisterburt]

Oh, strange! Sounds like you've got a handle on things, maybe updating the weights and setting strict=False is the best path :-)

[alisterburt]

To be clear, I don't think backwards compatibility is super important, just be clear about which version of our package is compatible with which version of monai/lightning etc

Maybe we should version monai strictly

[LorenzLamm]

To be clear, I don't think backwards compatibility is super important, just be clear about which version of our package is compatible with which version of monai/lightning etc

Maybe we should version monai strictly

Yes, I agree. Most important is that our most recent (and ideally best) model is working properly. Probably makes sense to version monai (and maybe also other packages?) strictly at some point?

But then let's for now go with the strict=False option. I have now also added the download link to the model compatible with the new MONAI version, even though everything should be working even without it. I'll make sure to have all upcoming models in the new MONAI format.

[alisterburt]

You're a star @LorenzLamm ! Hope you enjoyed vacation