search

teamtomo / membrain-seg

membrane segmentation in 3D for cryo-ET
Other
48 stars 12 forks source link

Training adjustments #4

Closed LorenzLamm closed 1 year ago

LorenzLamm commented 1 year ago

I tried to reproduce the performance of our best model, which is based on nnUNet. To this end, I added a lot of data augmentations, adjusted the model architecture, including deep supervision, and refined the training procedure (e.g. PolyLR scheduler, slight weight decay). Now I'm happy with the performance. Please have a look and see if everything is fine :)

LorenzLamm commented 1 year ago

Hey, @LorenzLamm ! This is great. Thanks for this. I have made some minor comments below. I think we merge soon so that we can keep moving.

Things I think we should change before merging:

  • where possible, I think we should remove the ___main__ code from the various modules. I can see these are useful for testing, but I think we will want users to use our CLI (as we develop it) rather than calling the module files directly.
  • Can we keep the SemanticSegmentationUnet in the networks module? Users/devs may want to import it for inference as well as training, so I think it would be more intuitive there.

Additionally, I think there are a few things we can do in follow-up PRs

  • convert print statements to logging that way users and developers can easily suppress them.
  • add docstrings with parameters and return values explained

Thanks a lot for your feedback, Kevin! I tried to work in your suggested changes and removed most of the main() functions + changed the unet folder back to "networks". Also cleaned up some lines of code & removed some redundant functions.

As you suggest, probably it's best to replace the print statements and add more information to the docstring in later PRs. Need to invest some time for this I think :-D