pytorch / pytorch/vision

Make `to_image` reading return tensors in CHW

Open
#9,547 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
17.9k
Forks
7.3k
Avg merge
1d 15h
Merged PRs (30d)
13

Description

🚀 The feature

Make to_image reading return tensors in CHW

Motivation, pitch

https://github.com/pytorch/vision/blob/da446275eccae6cd00923080d9cc91f1f3a73744/torchvision/transforms/v2/functional/_type_conversion.py#L10-L25

import pyvips
import torchvision,torchvision.transforms.v2

img = pyvips.Image.new_from_file(r"F:\test.jpg")
tensor = torchvision.transforms.v2.functional.to_image(img.numpy())
tensor.shape
#torch.Size([3, 1830, 1370])

tensor.shape, tensor.stride(), tensor.is_contiguous(memory_format=torch.channels_last), tensor.unsqueeze(0).is_contiguous(memory_format=torch.channels_last)

# (torch.Size([3, 1830, 1370]), (2507100, 1370, 1), False, False) # Not memory_format=torch.channels_last

new_tensor = torch.from_numpy(tensor.numpy()).permute((2,0,1))
new_tensor.shape, new_tensor.stride(), new_tensor.is_contiguous(memory_format=torch.channels_last), new_tensor.unsqueeze(0).is_contiguous(memory_format=torch.channels_last)

# (torch.Size([3, 1830, 1370]), (1, 4110, 3), False, True) # This is memory_format=torch.channels_last

https://docs.pytorch.org/vision/main/transforms.html#range-and-dtype
...
Transforms tend to be sensitive to the input strides / memory format. Some transforms will be faster with channels-first images while others prefer channels-last. Like torch operators, most transforms will preserve the memory format of the input, but this may not always be respected due to implementation details. You may want to experiment a bit if you’re chasing the very best performance. Using torch.compile() on individual transforms may also help factoring out the memory format variable (e.g. on Normalize). Note that we’re talking about memory format, not tensor shape.

Note that resize transforms like Resize and RandomResizedCrop typically prefer channels-last input and tend not to benefit from torch.compile() at this time.
...

As torchvision.transforms.v2.functional.to_image usually being the starting point of the processing pipeline, the latter process of it would always be float().div_(255) or to_dtype then resize and etc., these all prefer channels_last memory_format for better performance.
Therefore the starting point of the processing pipeline should keep its NHWC memory format. And which make the memcopy by contiguous() avoid that always happen (since image is always HWC when read in numpy)

Thanks for your reading.
I could make a PR

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 torchvision/transforms/v2/functional/_type_conversion.py, especially the linked to_image implementation and its handling of NumPy image strides. Clarify whether the requested result is CHW shape, channels-last memory format, or both, then inspect the surrounding transform behavior. Done means the returned tensor has the agreed shape and memory layout without unnecessary copying.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python, pytorch
Domain
computer-vision, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.