NCHW to NHWC without copying

An image batch often uses one of these logical axis orders:

NCHW: (batch, channels, height, width)
NHWC: (batch, height, width, channels)

Many PyTorch image layers use NCHW. Other interfaces use channels last; check the receiving function's contract rather than assuming an order. To change a strided tensor from NCHW to NHWC, use permute:

import torch

x = torch.arange(12).reshape(1, 3, 2, 2)
y = x.permute(0, 2, 3, 1)
print(list(y.shape))  # [1, 2, 2, 3]
print(y[0, 0, 0].tolist())  # [0, 4, 8]
print(y.untyped_storage().data_ptr() == x.untyped_storage().data_ptr())  # True

The arguments list input axes in their desired output order. Here batch stays first, height and width follow, and channels move last. Sizes and strides change together, so y[n, h, w, c] refers to x[n, c, h, w]. The data buffer is shared.

When a copy is needed

A view avoids copying values, but still creates tensor metadata. The new layout can affect the cost of later operations.

Calling .contiguous() on a tensor that is not contiguous in the requested memory format copies its values. If it already has that layout, the call returns the same tensor. For example, permuting a (1, 1, 1, 1) tensor leaves it contiguous, so the starter does not copy that particular input.

Do not add .contiguous() automatically after every permutation. Check whether the next operation needs a particular layout. Even view can accept some noncontiguous inputs; what matters is whether the requested shape can be represented by the existing strides.

A NumPy round trip is not a reliable way to force a copy either. For a compatible CPU tensor, numpy(), NumPy's transpose and torch.from_numpy can all share the same storage. Using permute expresses this task directly and avoids NumPy conversion restrictions.

Changing logical NCHW axes to NHWC is distinct from requesting PyTorch's channels_last memory format, which retains the logical NCHW shape.

def to_nhwc(x: torch.Tensor) -> torch.Tensor