Lightning-AI / Lightning-AI/torchmetrics
Almost All Image Metrics Save All Images In a Buffer
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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