cannot convert to channels last format for conv2d conv3d hybrid model

[SeedKunY]

🐛 Describe the bug

cannot convert to channels last format for conv2d conv3d hybrid model

import torch.nn as nn
import torch

cpu_device = torch.device("cpu")


class Hybrid_model(nn.Module):
  def __init__(self):
    super().__init__()
    self.layer1 = nn.Conv2d(2, 4, kernel_size=3, stride=1, padding=1, bias=False)
    self.layer2 = nn.Conv3d(4, 8, kernel_size=3, stride=1, padding=1, bias=False)

  def forward(self, inputs):
    x = self.layer1(inputs)
    x = torch.reshape(x, (x.size(0), x.size(1), x.size(2), 16, 2))
    x = self.layer2(x)
    return x


if __name__ == "__main__":
  test_model = Hybrid_model().to(memory_format=torch.channels_last)
  x = torch.randn([3, 2, 32, 32], dtype=torch.float, requires_grad=True)
  y = test_model(x)

python test_conv2dconv3d_hybrid_model.py
Traceback (most recent call last):
  File "test_conv2dconv3d_hybrid_model.py", line 21, in <module>
    test_model = Hybrid_model().to(memory_format=torch.channels_last)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 899, in to
    return self._apply(convert)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 570, in _apply
    module._apply(fn)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 593, in _apply
    param_applied = fn(param)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 896, in convert
    non_blocking, memory_format=convert_to_format)
RuntimeError: required rank 4 tensor to use channels_last format

Even we use test_model = Hybrid_model().to(memory_format=torch.channels_last_3d), it still report error as below:

python test_conv2dconv3d_hybrid_model.py
Traceback (most recent call last):
  File "test_conv2dconv3d_hybrid_model.py", line 21, in <module>
    test_model = Hybrid_model().to(memory_format=torch.channels_last_3d)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 899, in to
    return self._apply(convert)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 570, in _apply
    module._apply(fn)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 593, in _apply
    param_applied = fn(param)
  File "/home/gta/miniconda3/envs/xxx/lib/python3.7/site-packages/torch/nn/modules/module.py", line 896, in convert
    non_blocking, memory_format=convert_to_format)
RuntimeError: required rank 5 tensor to use channels_last_3d format

Versions

Collecting environment information...
PyTorch version: 1.10.0a0+gitcb9f926
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.17

Python version: 3.7.11 (default, Jul 27 2021, 14:32:16) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.10.54+prerelease2927-x86_64-with-debian-bullseye-sid
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] torch==1.10.0a0+gitcb9f926
[conda] mkl 2022.0.1 h06a4308_117 defaults
[conda] mkl-include 2022.0.1 h06a4308_117 defaults
[conda] numpy 1.21.2 py37hd8d4704_0 defaults
[conda] numpy-base 1.21.2 py37h2b8c604_0 defaults
[conda] torch 1.10.0a0+gitcb9f926 pypi_0 pypi

[SeedKunY]

@ezyang @VitalyFedyunin

[ezyang]

Unfortunately under the current way that memory format is implemented it is not possible to handle this case correctly. You should manually transfer the weights for each of the inner convolutions separately.

[SeedKunY]

OK. Any support plan to provide python API like Hybrid_model().to(memory_format=torch.channels_last) to support both conv2d and conv3d at the same time?

[ezyang]

There isn't a plan but there is some related discussion at #78050

[N-Friederich]

@ezyang Thx for your work! Is there still no plan to support it? I think that could be a great feature in terms of speed :)