Make BenchmarkModule compatible with PyTorch Lightning 2.0

[guarin]

What has changed and why?

• Make BenchmarkModule compatible with PyTorch Lightning 2.0
• Create KNN features at start of validation epoch instead of end of training epoch

The benchmarks are not yet completely compatible with PyTorch Lightning 2.0 as we use the LARS optimizer from PyTorch Lightning Bolts which is not compatible with PyTorch Lightning 2.0: Lightning-Universe/lightning-bolts#962

I guess the simplest solution for this is to copy one of the LARS implementations into our package as the code is pretty simple and we can avoid additional dependencies.

How was it tested?

• Ran imagenette benchmark on simclr and got same accuracy as before
• Tested on PyTorch Lightning 1.8.0 and 2.0.1

[codecov]

Codecov Report

Patch coverage: 88.57% and project coverage change: +1.32 🎉

Comparison is base (cbeb363) 89.35% compared to head (0cf8316) 90.68%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1136      +/-   ##
==========================================
+ Coverage   89.35%   90.68%   +1.32%     
==========================================
  Files         129      131       +2     
  Lines        5477     5623     +146     
==========================================
+ Hits         4894     5099     +205     
+ Misses        583      524      -59     

Impacted Files	Coverage Δ
lightly/utils/benchmarking.py	94.28% <88.57%> (+94.28%)	⬆️

... and 4 files with indirect coverage changes

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

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.

[MalteEbner]

I don't understand the design of this fully yet and why which code is distributed across
on_validation_epoch_start
validation_step
on_validation_epoch_end

I am also a bit unsure about the usage.
E.g. before the change, the top1 was returned by the. validation_step.
Now this self.max_accuracy is set by on_validation_epoch_end.
Is this required by the new pytorch 2.0?

Furthermore, I am wondering whether we have a unittest for parts of it.

[guarin]

I don't understand the design of this fully yet and why which code is distributed across
on_validation_epoch_start
validation_step
on_validation_epoch_end

• on_validation_epoch_start -> Generate KNN features
• validation_step -> Get nearest neighbors for current validation batch
• on_validation_epoch_end -> Calculate accuracy

We use on_validation_epoch_start instead of training_epoch_end or training_epoch_end because we only want to calculate the knn features when we actually do validation.

Furthermore, I am wondering whether we have a unittest for parts of it.

Added unit tests. I kept tests quite minimal because we know from the benchmarks that it works correctly.

[MalteEbner]

Makes sense to me, thank you for explaining it.