`torchvision.transforms.v2.functional.pad` should have maintained the `channels_last` format of the input

[NewUserHa]

🐛 Describe the bug

import torch
import torchvision.transforms.v2.functional
from torchcodec.decoders import VideoDecoder

a=VideoDecoder(r"F:\....gif",device='cuda')

# When input is CHW channels_last 
c=torchvision.transforms.v2.functional.crop(a[0], 10,10,200,282)
c.shape, c.stride()
# (torch.Size([3, 200, 200]), (1, 876, 3))

d=torchvision.transforms.v2.functional.resize(a[0], (300,300))
d.shape, d.stride()
# (torch.Size([3, 300, 300]), (1, 900, 3))

e=torchvision.transforms.v2.functional.pad(a[0], (500,500), 255, 'constant')
e.shape, e.stride()
# (torch.Size([3, 1241, 1292]), (1603372, 1292, 1))  # channels_last broken!

f=torch.nn.functional.pad(a[0], (500,500,500,500), 'constant', 255)
f.shape, f.stride()
# (torch.Size([3, 1241, 1292]), (1603372, 1292, 1))  # channels_last broken!

g=torchvision.transforms.v2.functional.crop(a[0], 10,10,500,500)
g.shape, g.stride()
# (torch.Size([3, 500, 500]), (250000, 500, 1))


# When input is NCHW channels_last 
aa=a[0].unsqueeze(0)
aa.shape,aa.stride()

cc=torchvision.transforms.v2.functional.crop(aa, 10,10,200,200)
print(cc.shape, cc.stride())

dd=torchvision.transforms.v2.functional.resize(aa, (300,300))
print(dd.shape, dd.stride())

ee=torchvision.transforms.v2.functional.pad(aa, (500,500), 255, 'constant')
print(ee.shape, ee.stride())

ff=torch.nn.functional.pad(aa, (500,500,500,500), 'constant', 255)
print(ff.shape, ff.stride())

gg=torchvision.transforms.v2.functional.crop(aa, 10,10,500,500)
print(gg.shape, gg.stride())

# torch.Size([1, 3, 200, 200]) (3, 1, 876, 3)
# torch.Size([1, 3, 300, 300]) (3, 1, 900, 3)
# torch.Size([1, 3, 1241, 1292]) (4810116, 1603372, 1292, 1)  # channels_last broken!
# torch.Size([1, 3, 1241, 1292]) (4810116, 1603372, 1292, 1)  # channels_last broken!
# torch.Size([1, 3, 500, 500]) (750000, 250000, 500, 1)


# When input is NCHW channels_last for sure! perfect and stictly channels_last stride.
aa=a[0].unsqueeze(0).to(memory_format=torch.channels_last)
aa.shape,aa.stride()

cc=torchvision.transforms.v2.functional.crop(aa, 10,10,200,200)
print(cc.shape, cc.stride())

dd=torchvision.transforms.v2.functional.resize(aa, (300,300))
print(dd.shape, dd.stride())

ee=torchvision.transforms.v2.functional.pad(aa, (500,500), 255, 'constant')
print(ee.shape, ee.stride())

ff=torch.nn.functional.pad(aa, (500,500,500,500), 'constant', 255)
print(ff.shape, ff.stride())

gg=torchvision.transforms.v2.functional.crop(aa, 10,10,500,500)
print(gg.shape, gg.stride())

# torch.Size([1, 3, 200, 200]) (211116, 1, 876, 3)
# torch.Size([1, 3, 300, 300]) (3, 1, 900, 3)
# torch.Size([1, 3, 1241, 1292]) (4810116, 1603372, 1292, 1)  # channels_last broken!
# torch.Size([1, 3, 1241, 1292]) (4810116, 1, 3876, 3)
# torch.Size([1, 3, 500, 500]) (750000, 250000, 500, 1)

As we can see,

• resize and crop they both support "CHW" and "NCHW", "channels_last", "strict stride channels_last" memory_format.
• torch.nn.functional.pad supports only "NCHW", "strict stride channels_last" memory_format, not "CHW" nor "(unstrict) channels_last".
• torchvision.transforms.v2.functional.pad does not even support what torch.nn.functional.pad supported.

The "channels_last" is needed and has higher performance in image process where R G B are together so filters can process them at onece, and gpu need it to accelarate in its tensor cores.

So torchvision.transforms.v2.functional.pad should have the ablity to keep the channels_last format of the input like resize and crop, instead of breaking it within the pipeline process in torchvision.

Thank your time for reading.
I would like to make a PR.

Versions

torchvision.version
'0.27.1+cu126'

[NicolasHug]

Hi @NewUserHa , transforms try to preserve memory layout but it's not a guarantee - sometimes it would cause an extra undesirable copy, so there's no good one-size-fits-all policy.

The recommendation is for user to explicitly convert to channels-last or to the memory layout of their choice.

[NewUserHa]

However, the torchvision.transforms.v2.functional.pad and torchvision.transforms.v2.functional.to_image can be easy to be updated to solve to maintain channels_last memory format,
since the channels_last boost performance just ~10%-30%, forcing user to waste it on memcopy, it could cancel the performance it gained.