Guarin lig 1572 implement msn paper

[guarin]

TODO:

• [x] projection head
• [x] collate function
• [x] loss
• [ ] examples
• [ ] benchmark
• [ ] docs
• [ ] add to api docs
• [ ] update readme

Closes #878

[codecov[bot]]

Codecov Report

Merging #895 (fded3d7) into master (097927c) will increase coverage by 0.01%. The diff coverage is 90.42%.

@@            Coverage Diff             @@
##           master     #895      +/-   ##
==========================================
+ Coverage   89.36%   89.38%   +0.01%     
==========================================
  Files          97       98       +1     
  Lines        4335     4426      +91     
==========================================
+ Hits         3874     3956      +82     
- Misses        461      470       +9     

Impacted Files	Coverage Δ
lightly/models/modules/masked_autoencoder.py	92.00% <70.00%> (-8.00%)	:arrow_down:
lightly/loss/msn_loss.py	95.00% <95.00%> (ø)
lightly/data/__init__.py	100.00% <100.00%> (ø)
lightly/data/collate.py	97.45% <100.00%> (+0.18%)	:arrow_up:
lightly/loss/__init__.py	100.00% <100.00%> (ø)
lightly/models/modules/heads.py	98.78% <100.00%> (+0.04%)	:arrow_up:

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[IgorSusmelj]

On which dataset are the benchmarks? And which GPU? The GPU memory usage seems waaay off. If you want to add the benchmark results, please make sure you use the same config (batch size, epochs etc.) as we use by default such that all the benchmarks can be run together and we should get the same result :)

[guarin]

On which dataset are the benchmarks? And which GPU? The GPU memory usage seems waaay off. If you want to add the benchmark results, please make sure you use the same config (batch size, epochs etc.) as we use by default such that all the benchmarks can be run together and we should get the same result :)

Good point :) I ran the benchmarks with the augmentation settings from the paper (image size = 224). In total there are a couple reasons why GPU usage is higher compared to other models:

• ViT-S backbone instead of Res18
• Many views in single forward pass (12 views)
• Larger image size than for other benchmarks (224 instead of 128)

I ran an additional benchmark with a smaller image size (128) and got the following results:

------------------------------------------------------------------------------------------
| Model         | Batch Size | Epochs |  KNN Test Accuracy |       Time | Peak GPU Usage |
------------------------------------------------------------------------------------------
| MSN           |        256 |    200 |              0.741 |   92.7 Min |     16.3 GByte |
------------------------------------------------------------------------------------------

I already updated the results in the benchmark file but not yet the docs. I am currently running the 800 epoch benchmark with smaller image size as well. Will update all the benchmark results once this one is done.

[guarin]

I'll run the 800 epoch benchmark, update the benchmark docs and then merge it. Should be done by tomorrow.