Memory bank distributed training support

[guarin]

Changes

• Add feature dimension when setting memory bank size. This is required for distributed training because the buffer size cannot be changed dynamically.
• Cleanup memory bank implementation.
• Add feature_dim_first argument to memory bank. This is required for backwards compatibility. In the future we should be able to set feature_dim_first=False and remove the requirement for transposing features coming from the feature bank.
• Update classes/functions using memory bank from size=num_features to size=(num_features, dim).

How was it tested?

• Added unit tests

[codecov]

Codecov Report

Attention: 7 lines in your changes are missing coverage. Please review.

Comparison is base (18f13b2) 85.22% compared to head (d1c0063) 85.27%.

Files	Patch %	Lines
lightly/models/utils.py	0.00%	7 Missing ⚠️

Additional details and impacted files

@@                                Coverage Diff                                @@
##           guarin-lig-3062-add-mocov2-imagenet-benchmark    #1293      +/-   ##
=================================================================================
+ Coverage                                          85.22%   85.27%   +0.04%     
=================================================================================
  Files                                                130      130              
  Lines                                               5522     5533      +11     
=================================================================================
+ Hits                                                4706     4718      +12     
+ Misses                                               816      815       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[philippmwirth]

Approved, let's see how the MoCo benchmarks works now.

[guarin]

Merging this as the model is training well, see #1291