New argument compute_on_cpu

[SkafteNicki]

What does this PR do?

Rework of #851
Fixes #848
General solution that allows all metrics where the metric state is a list to store the state on CPU instead of GPU through a new keyword argument compute_on_cpu=True/False (default False). This can beneficial in situations where the user do not have that much VRAM. Consequence of enabling this feature is that code in compute will be executed on CPU (code in update will still be executed on GPU).

Fixes #866
Related as the MAP metric currently moves all metric states to CPU for faster computations. As the issue points out this can overload the CPU. Feature in this PR essentially also fixes this by making it an option to move to CPU by setting MeanAveragePrecision(compute_on_cpu=True).

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #867 (7da164c) into master (d587be9) will increase coverage by 25%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #867     +/-   ##
=======================================
+ Coverage      70%    95%    +25%     
=======================================
  Files         175    175             
  Lines        7468   7476      +8     
=======================================
+ Hits         5249   7102   +1853     
+ Misses       2219    374   -1845     

[maximsch2]

A few questions:

• If I'm reading code correctly, self.device will still be gpu with this during compute? I'm wondering if this will break some metrics.
• The implicit assumption here is that a List state is only appended during update as otherwise we'll get device mismatches, right?
• This doesn't help with other metrics that can occupy significant memory (e.g. binned curve metrics) on GPU as those are allocated as fixed size tensors. In that case though, I guess, a user can manually place them on CPU and transfer data to CPU before update. Is that what we would be recommending people?

[justusschock]

@maximsch2 I think the device should indicate the input and output device of a metric. E.g. even if the metic states are on cpu, it would require the devices of the inputs to be cuda and handle the device-transfer logic accordingly and the output would also be on gpu.

[Borda]

@SkafteNicki mind resolve conflicts 🐰
@justusschock mind review? 🐟

[Borda]

@SkafteNicki @justusschock pls ^^ 🐰

[Borda]

seems there is a DDP issue... :/

[SkafteNicki]

@Borda all unrelated to the PR, trying to run tests again :]