Lightning-AI / Lightning-AI/torchmetrics

Almost All Image Metrics Save All Images In a Buffer

Open
#1,224 13 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement help wanted topic: Image
Dominant language
Python
Stars
2.5k
Forks
526
Avg merge
6d 11h
Merged PRs (30d)
5

Description

## 🐛 Bug

Currently the SSIM module seems to cache *all* of the predicted and reference images for the entire dataset, computing the metric only at the end of validation.

That's probably not a good idea for a couple of reasons:

1. Images tend to be quite large, video frames make this even larger. This is a great way to run out of memory.
2. The code crashes if the images are of different sizes. In other libraries this is solvable with a batch size of 1, but not in torchmetrics

I think most people would compute SSIM over each image in the batch and sum the results over all batches and divide by the total number of images. As far as I can tell there is *no* advantage to keeping the entire prediction and reference sets in working memory.

This is also true for the following metrics

* SpectralDistortionIndex
* ErrorRelativeGlobalDimensionlessSynthesis
* InceptionScore
* KernelInceptionDistance
* SpectralAngleMapper
* SSIM
* MS-SSIM
* UniversalImageQualityIndex

Additionally a number of these modules have incorrect help that mentions the wrong metric, almost like someone was just blindly copying and pasting code when the implementing them (far be it from me to make that accusation though)

This module needs a lot of love

### To Reproduce

1. Take two images of different sizes
3. Call update on them separately
4. Call compute

#### Code sample

```py
import torch
from torchmetrics import StructuralSimilarityIndexMeasure

ssim = StructuralSimilarityIndexMeasure()

a = torch.rand(5, 3, 10, 10)
b = torch.rand(5, 3, 20, 20)

ssim.update(a, a)
ssim.update(b, b)

print(ssim.compute())
```

output:

```
Traceback (most recent call last):
File "repro.py", line 12, in
print(ssim.compute())
File "/mnt/fsx-outputs-chipdesign/mehrlich/poetry_cache/metabit-wl5Ts0fA-py3.8/lib/python3.8/site-packages/torchmetrics/metric.py", line 531, in wrapped_func
value = compute(*args, **kwargs)
File "/mnt/fsx-outputs-chipdesign/mehrlich/poetry_cache/metabit-wl5Ts0fA-py3.8/lib/python3.8/site-packages/torchmetrics/image/ssim.py", line 117, in compute
preds = dim_zero_cat(self.preds)
File "/mnt/fsx-outputs-chipdesign/mehrlich/poetry_cache/metabit-wl5Ts0fA-py3.8/lib/python3.8/site-packages/torchmetrics/utilities/data.py", line 42, in dim_zero_cat
return torch.cat(x, dim=0)
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 10 but got size 20 for tensor number 1 in the list.
```

### Expected behavior

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): git master, installed with pip
- Python & PyTorch Version (e.g., 1.0): 3.8/1.12
- Any other relevant information such as OS (e.g., Linux):

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with torchmetrics/image/ssim.py, especially compute where the traceback reaches dim_zero_cat, then inspect the other listed image metrics for the same buffering behavior and incorrect help text. Reproduce the two differently sized updates from the issue and add focused coverage; done means supported metrics no longer require concatenating all images before compute and the help text names the correct metric.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.