Fix slow performance for confusion matrix based metrics

[SkafteNicki]

What does this PR do?

Fixes #1267
Fixes #1277
After the refactor the underlying _bincount function that does a lot of the computations for confusion matrix based metrics was changed. I did some testing at the time and everything seemed to be fine. However, for large inputs the
implementation is really slow.

Here is a direct comparison in colab:
https://colab.research.google.com/drive/18tGZj_dPria6NSwVOIgwPXO8mJBr21kc?usp=sharing

The results:

The TLDR:

• Regardless of batch size the new implementation is approx 2x slower on cpu
• For batch size 100 the new implementation is approx ~4x FASTER on gpu
• For batch size 10.000 the new implementation is approx ~30x SLOWER on gpu

The simple fix currently in this PR is to change it back to the old implementation. The alternative is that we have something like this in the code:

def _bincount(x, minlength=minlength):
    ...
    if len(x) > N or not x.is_cuda:
        return torch.bincount(x, minlength=minlenght)
    else:
        z = torch.zeros(minlength, device=x.device, dtype=x.dtype)
        return z.index_add_(0, x, torch.ones_like(x))

where we have to set N based on some experimentation.
@Borda, @justusschock opinions?

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 🙃

[Borda]

LGTM 🎉

[codecov]

Codecov Report

Merging #1302 (fb3a387) into master (604ed80) will decrease coverage by 0%.
The diff coverage is 100%.

Additional details and impacted files

@@          Coverage Diff           @@
##           master   #1302   +/-   ##
======================================
- Coverage      87%     87%   -0%     
======================================
  Files         190     190           
  Lines       11121   11120    -1     
======================================
- Hits         9621    9620    -1     
  Misses       1500    1500           

[justusschock]

@SkafteNicki For now, I think sticking to the default here makes sense. I would be open for changing this conditionally in the future, but this would require more experimenting :)

[Borda]

I would be open for changing this conditionally in the future, but this would require more experimenting :)

do you want to open an issue for it so we won't forget? 🦦

[Callidior]

@SkafteNicki I stumbled upon this issue and did some investigation.

The new implementation used torch.Tensor.index_add_, which is known to be slow on integer tensors (see pytorch/pytorch#42109). That issue is open now for over 2 years, so a fix in the near future seems unlikely. However, the issue mentions using a float64 tensor as a work-around.

I ran a quick performance test with different tensor sizes.

Integer implementation:

z = torch.zeros(4, device=unique_mapping.device, dtype=unique_mapping.dtype)
return z.index_add_(0, unique_mapping, torch.ones_like(unique_mapping))

Float implementation:

z = torch.zeros(4, device=unique_mapping.device, dtype=torch.float64)
return z.index_add_(0, unique_mapping, torch.ones_like(unique_mapping, dtype=torch.float64))

Timing results on RTX 3090 (in milliseconds):

Number of Samples	Integer index_add	Float index_add	torch.bincount
1,000	0.3	0.1	0.2
100,000	144.7	0.1	0.2
10,000,000	18,497.1	9.3	4.2

Timing results on CPU (in milliseconds):

Number of Samples	Integer index_add	Float index_add	torch.bincount
1,000	0.20	0.03	0.01
100,000	0.48	0.39	0.19
10,000,000	18.0	16.3	8.4

Looks like torch.bincount is the best option, aside from other issues with it like #1413.