Repository: Lightning-AI/pytorch-lightning

Issue:
Slow accuracy metric
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

Accuracy class from metrics.classification is too slow, compared to pytorch vanilla implementation.


### To Reproduce

Try the following demo code.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue. 
Minimal means having the shortest code but still preserving the bug. -->


```python3
from time import time
import torch
from pytorch_lightning.metrics.classification import Accuracy


def test_acc(pred, targ, get_acc, steps=1000, rep=10):
    tsum, mean_res = 0, []
    for i in range(rep):
        res = []
        ts = time()
        for j in range(steps):
            res.append(get_acc(pred, targ))
        te = time()
        tsum += te - ts
        mean_acc = torch.stack(res).mean().item()
        mean_res.append(mean_acc)
    
    return (tsum / rep), mean_res


def demo(batch_size=32, num_classes=10, device='cuda', steps=100, rep=5):
    print("Testing  bs: %d,  num_cls: %d,  device: %s,  steps: %d,  rep: %d" % (batch_size, num_classes, device, steps, rep))

    pred = torch.randint(0, num_classes, (batch_size, )).to(device=device)
    targ = torch.randint(0, num_classes, (batch_size, )).to(device=device)

    pl_acc = Accuracy(num_classes=num_classes)
    def pt_acc(pred, targ):
        return pred.eq(targ).to(dtype=torch.float).mean()

    pl_time, pl_res = test_acc(pred, targ, pl_acc, steps, rep)
    pt_time, pt_res = test_acc(pred, targ, pt_acc, steps, rep)

    if pl_res != pt_res:
        print("Results mismatch!")
    print("Lightning: %2.6fs" % pl_time)
    print("Vanilla  : %2.6fs" % pt_time)
    print("")

if __name__ == "__main__":
    demo(32, 10)
    demo(512, 10)
    demo(32, 100)
    demo(512, 100)
    demo(32, 1000)
    demo(512, 1000)

    demo(32, 10, 'cpu')
    demo(512, 10, 'cpu')
    demo(32, 100, 'cpu')
    demo(512, 100, 'cpu')
    demo(32, 1000, 'cpu')
    demo(512, 1000, 'cpu')
```

On my machine, the results are like the following.

```
Testing  bs: 32,  num_cls: 10,  device: cuda,  steps: 100,  rep: 5
Lightning: 0.378017s
Vanilla  : 0.002766s

Testing  bs: 512,  num_cls: 10,  device: cuda,  steps: 100,  rep: 5
Lightning: 0.372391s
Vanilla  : 0.002776s

Testing  bs: 32,  num_cls: 100,  device: cuda,  steps: 100,  rep: 5
Lightning: 3.346141s
Vanilla  : 0.002812s

Testing  bs: 512,  num_cls: 100,  device: cuda,  steps: 100,  rep: 5
Lightning: 3.335049s
Vanilla  : 0.002660s

Testing  bs: 32,  num_cls: 1000,  device: cuda,  steps: 100,  rep: 5
Lightning: 32.751692s
Vanilla  : 0.002666s

Testing  bs: 512,  num_cls: 1000,  device: cuda,  steps: 100,  rep: 5
Lightning: 32.967594s
Vanilla  : 0.002762s

Testing  bs: 32,  num_cls: 10,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.128131s
Vanilla  : 0.001317s

Testing  bs: 512,  num_cls: 10,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.128481s
Vanilla  : 0.001403s

Testing  bs: 32,  num_cls: 100,  device: cpu,  steps: 100,  rep: 5
Lightning: 1.117403s
Vanilla  : 0.001275s

Testing  bs: 512,  num_cls: 100,  device: cpu,  steps: 100,  rep: 5
Lightning: 1.166188s
Vanilla  : 0.001330s

Testing  bs: 32,  num_cls: 1000,  device: cpu,  steps: 100,  rep: 5
Lightning: 10.927701s
Vanilla  : 0.001229s

Testing  bs: 512,  num_cls: 1000,  device: cpu,  steps: 100,  rep: 5
Lightning: 11.560013s
Vanilla  : 0.001310s
```

100x ~ 10000x more time is needed for pl's accuracy, depending on the num_classes argument.

### Expected behavior

Pytorch-lightning's accuracy function should take similar time compared to the pytorch vanilla implementation, and it should not depend on num_classes too much.


<!-- A clear and concise description of what you expected to happen. -->

### Environment

```
* CUDA:
	- GPU:
		- GeForce GTX 1050 Ti
	- available:         True
	- version:           10.2
* Packages:
	- numpy:             1.19.0
	- pyTorch_debug:     False
	- pyTorch_version:   1.5.1
	- pytorch-lightning: 0.8.1
	- tensorboard:       2.2.2
	- tqdm:              4.47.0
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- ELF
	- processor:         x86_64
	- python:            3.6.9
	- version:           #46~18.04.1-Ubuntu SMP Fri Jul 10 07:21:24 UTC 2020
```

I tested the same code from two other machines (each with 2080 ti and with V100), and they give similar results.

### Additional context

<!-- Add any other context about the problem here. -->

I'm not sure this issue falls into the bug report or the feature request category. Please relocate this as you like. :)



Discussion:
Hi! thanks for your contribution!, great first issue!
Eh, now this issue looks a lot like a enhancement issue. Sorry :sweat_smile: 

I saw the implementation of the accuracy function, and I noticed a few metric functions depend on the [stat_scores_multiple_classes](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/metrics/functional/classification.py#L137) function.

I think this function is way too slow, as it doesn't run parallel on num_classes. It would be great if we can reimplement this function only with vanilla pytorch operations, but I can't find a way for that.

Considering the huge complexity gap, I think implementing on ad-hoc basis for each metric would be better, though it might be a bit ugly.

Hi @Diuven , I was able to reproduce your results. 

I agree, that this is a huge slowdown and we should fix it. Unfortunately, I am not quite sure, why this happens.

I have profiled the code and by far the most time is taken to compute the stat scores.  Although I don't know why this takes that long in PyTorch, I am afraid, that we need a more general version. Do you have a proposal on how to parallelise the computation of stat scores among classes?

To be honest I don't like to duplicate stuff there. But I agree, that we should speed the current implementation up :)
Hi @justusschock , thanks for your reply.

Yeah, I agree on not having duplicate codes in pl. It would be great if we simply replace `stat_scores_multiple_classes` with faster version.

Though I haven't profiled the code, I think the main reason is that there is no parallelism on `num_classes`. In current implementation of `stat_scores_multiple_classes`, we simply run for loops to compute states for each classes.
In the results of the demo code, the elapsed time also heavily depends on `num_classes`, so it supports the claim above.

So if we replace the for loop, `for c in range(num_classes):`,  in the function with the parallelized code, it would be much faster. Or at least it wouldn't depend on `num_classes`.


Making a CUDA library just for this function would work, but that would be our last option, considering the pl's philosophy and current state. (I think there's no CUDA code in this repo, so it would be an overkill to make a CUDA library only for this)

I'm not sure about whether this will work well or not, but I think parallelized [scatter](https://pytorch.org/docs/stable/tensors.html?highlight=scatter#torch.Tensor.scatter) function in pytorch might help.
You're right that this would be great. Mind putting together a draft? I would for sure help you with that. 

Also it's worth mentioning, that PLs version will always be a bit slower than the pure pytorch version, since it does a bit more. I also checks all incoming and outgoing arguments for device and type and converts them if necessary :)
Here's my draft:

```python3
def stat_scores_multiple_classes(pred: LongTensor, targ: LongTensor):
  match_true = (pred == targ).long()
  match_false = match_true.logical_not()
  # tps, fps, tns, fns, sups initialized as same
  
  tps.scatter_add_(dim=0, pred, match_true)
  fps.scatter_add_(dim=0, pred, match_false)
  fns.scatter_add_(dim=0, targ, match_false)
  tns = pred.size(0) - tps - fps - fns
  sups.scatter_add_(dim=0, targ, torch.ones_like(match_true))

  return tps, fps, tns, fns, sups
```
edit: I erased unused two lines. It was a draft of the draft :sweat_smile: 

Please check if I missed something big there. If it seems right, I'll work on it and make pull request.


> Also it's worth mentioning, that PLs version will always be a bit slower than the pure pytorch version, since it does a bit more. I also checks all incoming and outgoing arguments for device and type and converts them if necessary :)

Yeah, that's a really great part of PL, and why I like it. :smile_cat: I agree on compromising a bit of speediness to get huge convenience. All we need is just comparable performance to vanilla case, in my opinion.
I checked the code below with the demo code at the top, and it was 10x faster than the previous code, but still slower than the vanilla code and still depends on the `num_classes`.

```python3
def stat_scores_multiple_classes(
        pred: torch.Tensor,
        target: torch.Tensor,
        num_classes: Optional[int] = None,
        argmax_dim: int = 1,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    
    
    #num_classes = get_num_classes(pred=pred, target=target,
    #                              num_classes=num_classes)

    #if pred.ndim == target.ndim + 1:
    #    pred = to_categorical(pred, argmax_dim=argmax_dim)
    
    placeholder = torch.zeros((num_classes,), device=pred.device)

    match_true = (pred == target).float()
    match_false = 1 - match_true
    # tps, fps, tns, fns, sups initialized as same

    tps = placeholder.scatter_add(0, pred, match_true)
    fps = placeholder.scatter_add(0, pred, match_false)
    fns = placeholder.scatter_add(0, target, match_false)
    tns = pred.size(0) - (tps + fps + fns)
    sups = placeholder.scatter_add(0, target, torch.ones_like(match_true))

    return tps, fps, tns, fns, sups
```

I think if we don't give up the `(num_classes, )` sized tensor results, this is as far as we can get. If we make another version of `stat_scores_multiple_classes` returning only sums of the stats, it would be far better, in my opinion.

I'm not sure about whether we can bear with those two similar codes (sum version and tensor version of `stat_scores_multiple_classes` function). What's your thoughts?


Results:

```
Testing  bs: 32,  num_cls: 10,  device: cuda,  steps: 100,  rep: 5
Lightning: 0.053924s
Vanilla  : 0.002690s

Testing  bs: 512,  num_cls: 10,  device: cuda,  steps: 100,  rep: 5
Lightning: 0.052600s
Vanilla  : 0.002591s

Testing  bs: 32,  num_cls: 100,  device: cuda,  steps: 100,  rep: 5
Lightning: 0.234717s
Vanilla  : 0.002604s

Testing  bs: 512,  num_cls: 100,  device: cuda,  steps: 100,  rep: 5
Lightning: 0.237395s
Vanilla  : 0.002614s

Testing  bs: 32,  num_cls: 1000,  device: cuda,  steps: 100,  rep: 5
Lightning: 1.996236s
Vanilla  : 0.002550s

Testing  bs: 512,  num_cls: 1000,  device: cuda,  steps: 100,  rep: 5
Lightning: 2.008798s
Vanilla  : 0.002579s

Testing  bs: 32,  num_cls: 10,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.019533s
Vanilla  : 0.001189s

Testing  bs: 512,  num_cls: 10,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.020818s
Vanilla  : 0.001329s

Testing  bs: 32,  num_cls: 100,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.080327s
Vanilla  : 0.001167s

Testing  bs: 512,  num_cls: 100,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.081321s
Vanilla  : 0.001278s

Testing  bs: 32,  num_cls: 1000,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.682839s
Vanilla  : 0.001193s

Testing  bs: 512,  num_cls: 1000,  device: cpu,  steps: 100,  rep: 5
Lightning: 0.704203s
Vanilla  : 0.001301s
```
That's a great improvement :) 

As long as we only end up with these two alternatives and not one version per metric this should be fine, I think. But doesn't that also depend on the reduction type?
Yeah you're right.

But as long as we use only `mean`, `sum` and `none` reduction methods, it the `stat_scores_multiple_classes` function won't be too long.

If we want to support additional reduction methods in the future, we probably can use `none` version of the function (i.e. function returning `(num_classes, )` sized tensors) and simply reduce after that. But if we have a lot faster version for certain reduction types, such as `sum` and `mean`, it would be very useful to have faster version of it.

Resolving patch:
diff --git a/pytorch_lightning/metrics/functional/classification.py b/pytorch_lightning/metrics/functional/classification.py
--- a/pytorch_lightning/metrics/functional/classification.py
+++ b/pytorch_lightning/metrics/functional/classification.py
@@ -138,10 +138,10 @@ def stat_scores_multiple_classes(
         target: torch.Tensor,
         num_classes: Optional[int] = None,
         argmax_dim: int = 1,
+        reduction: str = 'none',
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """
-    Calls the stat_scores function iteratively for all classes, thus
-    calculating the number of true postive, false postive, true negative
+    Calculates the number of true postive, false postive, true negative
     and false negative for each class
 
     Args:
@@ -150,6 +150,12 @@ def stat_scores_multiple_classes(
         num_classes: number of classes if known
         argmax_dim: if pred is a tensor of probabilities, this indicates the
             axis the argmax transformation will be applied over
+        reduction: method for reducing result values (default: none)
+            Available reduction methods:
+
+            - elementwise_mean: takes the mean
+            - none: pass array
+            - sum: add elements
 
     Return:
         True Positive, False Positive, True Negative, False Negative, Support
@@ -173,16 +179,58 @@ def stat_scores_multiple_classes(
     if pred.ndim == target.ndim + 1:
         pred = to_categorical(pred, argmax_dim=argmax_dim)
 
-    num_classes = get_num_classes(pred=pred, target=target,
-                                  num_classes=num_classes)
+    num_classes = get_num_classes(pred=pred, target=target, num_classes=num_classes)
 
-    tps = torch.zeros((num_classes,), device=pred.device)
-    fps = torch.zeros((num_classes,), device=pred.device)
-    tns = torch.zeros((num_classes,), device=pred.device)
-    fns = torch.zeros((num_classes,), device=pred.device)
-    sups = torch.zeros((num_classes,), device=pred.device)
-    for c in range(num_classes):
-        tps[c], fps[c], tns[c], fns[c], sups[c] = stat_scores(pred=pred, target=target, class_index=c)
+    if pred.dtype != torch.bool:
+        pred.clamp_max_(max=num_classes)
+    if target.dtype != torch.bool:
+        target.clamp_max_(max=num_classes)
+
+    possible_reductions = ('none', 'sum', 'elementwise_mean')
+    if reduction not in possible_reductions:
+        raise ValueError("reduction type %s not supported" % reduction)
+
+    if reduction == 'none':
+        pred = pred.view((-1, )).long()
+        target = target.view((-1, )).long()
+
+        tps = torch.zeros((num_classes + 1,), device=pred.device)
+        fps = torch.zeros((num_classes + 1,), device=pred.device)
+        tns = torch.zeros((num_classes + 1,), device=pred.device)
+        fns = torch.zeros((num_classes + 1,), device=pred.device)
+        sups = torch.zeros((num_classes + 1,), device=pred.device)
+
+        match_true = (pred == target).float()
+        match_false = 1 - match_true
+
+        tps.scatter_add_(0, pred, match_true)
+        fps.scatter_add_(0, pred, match_false)
+        fns.scatter_add_(0, target, match_false)
+        tns = pred.size(0) - (tps + fps + fns)
+        sups.scatter_add_(0, target, torch.ones_like(match_true))
+
+        tps = tps[:num_classes]
+        fps = fps[:num_classes]
+        tns = tns[:num_classes]
+        fns = fns[:num_classes]
+        sups = sups[:num_classes]
+
+    elif reduction == 'sum' or reduction == 'elementwise_mean':
+        count_match_true = (pred == target).sum().float()
+        oob_tp, oob_fp, oob_tn, oob_fn, oob_sup = stat_scores(pred, target, num_classes, argmax_dim)
+
+        tps = count_match_true - oob_tp
+        fps = pred.nelement() - count_match_true - oob_fp
+        fns = pred.nelement() - count_match_true - oob_fn
+        tns = pred.nelement() * (num_classes + 1) - (tps + fps + fns + oob_tn)
+        sups = pred.nelement() - oob_sup.float()
+
+        if reduction == 'elementwise_mean':
+            tps /= num_classes
+            fps /= num_classes
+            fns /= num_classes
+            tns /= num_classes
+            sups /= num_classes
 
     return tps, fps, tns, fns, sups
 
@@ -218,16 +266,13 @@ def accuracy(
         tensor(0.7500)
 
     """
-    tps, fps, tns, fns, sups = stat_scores_multiple_classes(
-        pred=pred, target=target, num_classes=num_classes)
-
     if not (target > 0).any() and num_classes is None:
         raise RuntimeError("cannot infer num_classes when target is all zero")
 
-    if reduction in ('elementwise_mean', 'sum'):
-        return reduce(sum(tps) / sum(sups), reduction=reduction)
-    if reduction == 'none':
-        return reduce(tps / sups, reduction=reduction)
+    tps, fps, tns, fns, sups = stat_scores_multiple_classes(
+        pred=pred, target=target, num_classes=num_classes, reduction=reduction)
+
+    return tps / sups
 
 
 def confusion_matrix(


Test patch:
diff --git a/tests/metrics/functional/test_classification.py b/tests/metrics/functional/test_classification.py
--- a/tests/metrics/functional/test_classification.py
+++ b/tests/metrics/functional/test_classification.py
@@ -121,15 +121,19 @@ def test_stat_scores(pred, target, expected_tp, expected_fp, expected_tn, expect
     assert sup.item() == expected_support
 
 
-@pytest.mark.parametrize(['pred', 'target', 'expected_tp', 'expected_fp',
+@pytest.mark.parametrize(['pred', 'target', 'reduction', 'expected_tp', 'expected_fp',
                           'expected_tn', 'expected_fn', 'expected_support'], [
-    pytest.param(torch.tensor([0., 2., 4., 4.]), torch.tensor([0., 4., 3., 4.]),
+    pytest.param(torch.tensor([0., 2., 4., 4.]), torch.tensor([0., 4., 3., 4.]), 'none',
+                 [1, 0, 0, 0, 1], [0, 0, 1, 0, 1], [3, 4, 3, 3, 1], [0, 0, 0, 1, 1], [1, 0, 0, 1, 2]),
+    pytest.param(to_onehot(torch.tensor([0., 2., 4., 4.])), torch.tensor([0., 4., 3., 4.]), 'none',
                  [1, 0, 0, 0, 1], [0, 0, 1, 0, 1], [3, 4, 3, 3, 1], [0, 0, 0, 1, 1], [1, 0, 0, 1, 2]),
-    pytest.param(to_onehot(torch.tensor([0., 2., 4., 4.])), torch.tensor([0., 4., 3., 4.]),
-                 [1, 0, 0, 0, 1], [0, 0, 1, 0, 1], [3, 4, 3, 3, 1], [0, 0, 0, 1, 1], [1, 0, 0, 1, 2])
+    pytest.param(to_onehot(torch.tensor([0., 2., 4., 4.])), torch.tensor([0., 4., 3., 4.]), 'sum',
+                 torch.tensor(2), torch.tensor(2), torch.tensor(14), torch.tensor(2), torch.tensor(4)),
+    pytest.param(to_onehot(torch.tensor([0., 2., 4., 4.])), torch.tensor([0., 4., 3., 4.]), 'elementwise_mean',
+                 torch.tensor(0.4), torch.tensor(0.4), torch.tensor(2.8), torch.tensor(0.4), torch.tensor(0.8))
 ])
-def test_stat_scores_multiclass(pred, target, expected_tp, expected_fp, expected_tn, expected_fn, expected_support):
-    tp, fp, tn, fn, sup = stat_scores_multiple_classes(pred, target)
+def test_stat_scores_multiclass(pred, target, reduction, expected_tp, expected_fp, expected_tn, expected_fn, expected_support):
+    tp, fp, tn, fn, sup = stat_scores_multiple_classes(pred, target, reduction=reduction)
 
     assert torch.allclose(torch.tensor(expected_tp).to(tp), tp)
     assert torch.allclose(torch.tensor(expected_fp).to(fp), fp)