Mean teacher model training

[plbenveniste]

This issue investigates the training of a model using the mean-teacher algorithm. The method is inspired from https://github.com/perone/mean-teacher.

It uses the following steps (copied from the above repo):

• Take a supervised architecture and make a copy of it. Let's call the original model the student and the new one the teacher.
• At each training step, use the same minibatch as inputs to both the student and the teacher but add random augmentation or noise to the inputs separately.
• Add an additional consistency cost between the student and teacher outputs (after softmax).
• Let the optimizer update the student weights normally.
• Let the teacher weights be an exponential moving average (EMA) of the student weights. That is, after each training step, update the teacher weights a little bit toward the student weights.

[plbenveniste]

Here are sources of inspiration which could prove useful in the future:

• https://github.com/lemoshu/AC-MT/tree/main : Ambiguity-selective consistency regularization for mean-teacher semi-supervised medical image segmentation
• https://github.com/CuriousAI/mean-teacher : basic mean teacher implementation (they seem to have commits more recent then Perone)
• https://github.com/HiLab-git/SSL4MIS : a benchmark with a lot of code on semi-supervised learning
• https://github.com/himashi92/Co-BioNet : Co-BioNet: Uncertainty-Guided Dual-Views for Semi-Supervised Volumetric Medical Image Segmentation

[plbenveniste]

I am working on branch plb/ssl to train the semi-supervised model with the mean-teacher. For now, it only uses labelled images.
For now my code is having the following behaviors:

-> I see this weird behaviour which shows that at each epoch end, the DiceCE loss rises.
-> The consistency loss doesn't change.

Looking into it now.