SemanticSegmentationTask: fix ignore_index weighting

[adamjstewart]

Fixes the second half of #1245. The following script can be used to confirm the bug and the fix:

import torch
from torchmetrics.classification import MulticlassAccuracy

metric = MulticlassAccuracy(num_classes=3, ignore_index=0, average="weighted")
prediction = torch.ones(20, 10)

truth = torch.zeros_like(prediction)
truth[9] = 1
truth[10:] = 2

metric(prediction[:10], truth[:10])
metric(prediction[10:], truth[10:])
print(metric.compute())
metric.reset()

truth = torch.zeros_like(prediction)
truth[10] = 1
truth[11:] = 2

metric(prediction[:10], truth[:10])
metric(prediction[10:], truth[10:])
print(metric.compute())

The default averaging technique "macro" produces the wrong values. Both "micro" and "weighted" compute the correct values. I believe "micro" stores the entire batch in memory and computes the metric at the end of each epoch, while "weighted" computes the metric and percent of non-ignored pixels for each batch and computes the metric at the end of each epoch, so "weighted" should have less memory requirements.

[isaaccorley]

We should still use micro instead since this represents overall mIoU. Most papers do not show weighted mIoU. Weighted means it's a weighted average across the individual classes, not across batches.

[isaaccorley]

Great catch on this actually. The majority of other torchmetrics default to average='micro' but MultiClassJaccardIndex does not.

[adamjstewart]

Would be good if someone could test this and make sure memory usage doesn't explode.