Learn AI Create the Future — with Pranay
Main site

GPU Memory and Speed

channels_last for convolution networks

channels_last stores the same image tensor in a different order in memory — the order convolution hardware prefers — and buys a real speedup for CNNs with two lines of code.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

channels_last keeps the same numbers but arranges them in memory in the order the hardware likes to read them.

Think of a warehouse storing shirts. One worker files them colour by colour: all blues together, all reds together. Another files them size by size: everything for one customer's order sits side by side. Same shirts, same shelves — but if orders always come per customer, the second arrangement means one walk down one aisle instead of twenty short trips.

Computer memory is one long shelf. How you interleave the data decides how much walking the hardware does.

Why it exists

An image tensor holds values for each position and each channel — the colour components, or the many feature maps deeper in a CNN. PyTorch's traditional order groups by channel: the whole red plane, then the whole green plane, then blue.

Convolution hardware — especially the tensor cores that do fast half-precision maths — wants the opposite: for one pixel, all channels side by side. When data arrives in the traditional order, the hardware reorganises it on the fly, again and again. channels_last stores it pre-organised, so convolutions read straight through.

How it works

same 2x2 image, 3 channels, one long shelf of memory:

channels first:  R R R R  G G G G  B B B B     (a plane at a time)
channels last:   R G B  R G B  R G B  R G B    (a pixel at a time)

Nothing about the values changes. The tensor's shape does not even change. Only the walking order does.

A real example you have seen

A well-organised kitchen does this. Ingredients used together — for one dish — sit in one tray, not sorted by ingredient type across ten cupboards. Restaurants that pre-tray their dishes serve visibly faster with the same ingredients and the same stove.

Remember this

  • channels_last changes memory order, never values or shapes.
  • Convolution hardware reads that order without reorganising, so CNNs speed up.
  • It pairs with half-precision maths — that is when the gain is biggest.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

The stride demo runs anywhere and is exact. The speed numbers were captured with torch 2.5.1 on an NVIDIA RTX A6000 — the gain depends heavily on GPU generation, model and dtype.

Proof that only the order changes

Strides — how many memory steps separate neighbouring positions along each dimension — expose the arrangement:

cl_strides.py
import torch

x = torch.randn(2, 3, 4, 4)                       # batch, channels, height, width
print("standard layout stride:", x.stride())

y = x.to(memory_format=torch.channels_last)
print("channels_last stride:  ", y.stride())
print("same shape:            ", x.shape == y.shape)
print("same values:           ", torch.equal(x, y))
print("channels_last?         ", y.is_contiguous(memory_format=torch.channels_last))
Output
standard layout stride: (48, 16, 4, 1)
channels_last stride:   (48, 1, 12, 3)
same shape:             True
same values:            True
channels_last?          True

Read the strides. Standard: moving one step in width moves 1 slot — width is innermost. channels_last: moving one channel moves 1 slot — channels are innermost, exactly the interleaving from the diagram. Shape and values: untouched. (Strides are the same machinery behind views and contiguity.)

The two lines, and the payoff

cl_speed.py
import torch
import torch.nn as nn
import time

if not torch.cuda.is_available():
    raise SystemExit("needs a GPU with tensor cores to show the gain")

def bench(memory_format):
    torch.manual_seed(0)
    model = nn.Sequential(
        nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(),
        nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
        nn.Conv2d(128, 128, 3, padding=1),
    ).cuda().to(memory_format=memory_format)
    x = torch.randn(32, 3, 224, 224, device="cuda").to(memory_format=memory_format)
    with torch.autocast("cuda", dtype=torch.float16):
        for _ in range(3):
            model(x)                       # warm-up lets cuDNN pick kernels
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(20):
            model(x)
        torch.cuda.synchronize()
    return (time.perf_counter() - start) / 20 * 1000

print(f"contiguous:    {bench(torch.contiguous_format):6.2f} ms per forward")
print(f"channels_last: {bench(torch.channels_last):6.2f} ms per forward")
Output
contiguous:     18.96 ms per forward
channels_last:  12.73 ms per forward

A third faster, from model.to(memory_format=torch.channels_last) plus the same call on the input. Convert both: a channels_last model fed standard-layout inputs converts them per layer, per step, and can end up slower than doing nothing. Note the autocast — in float32 the gain shrinks a lot, because the fastest tensor-core convolution kernels are the half-precision ones. The benchmark discipline (warm-up, synchronize) is from timing GPU code correctly.

Common mistakes

Converting the model but not the data. The most common failure. Convert the input batch in the training loop, right where it moves to the GPU: x = x.to("cuda", memory_format=torch.channels_last).

Expecting gains on transformers or MLPs. channels_last is a convolution story. Linear layers and attention neither need nor notice it — for those, look at SDPA and torch.compile.

A stray .contiguous() undoing everything. .contiguous() with no argument means standard order — one call mid-model silently converts back, and every later convolution pays reorganisation again. Use .contiguous(memory_format=torch.channels_last) when you need the explicit call.

Judging it on float32 CPU runs. The demo's gap comes from tensor-core kernel selection. On CPU or in float32 the effect ranges from small to nothing; measure on your target setup before adopting.

Try it yourself

Remove the torch.autocast context from both benchmarks and rerun in float32. Then swap the Conv2d stack for nn.Linear layers and confirm the effect disappears — layout is a convolution concern.

What to learn next

Researcher — Mathematics and papers.

Layout, vectorisation and tensor cores

NCHW ("channels first") makes the spatial dimension innermost; NHWC ("channels last") makes $C$ innermost. Implicit-GEMM convolution lowers to matrix multiply with the channel dimension contiguous in the reduction axis — NHWC lets 128-bit vector loads fetch consecutive channels directly into tensor-core fragments. cuDNN's NCHW tensor-core path inserts transpose kernels around the maths; NHWC removes them. NVIDIA's deep-learning performance guidance is blunt: tensor-core convolutions prefer NHWC, and frameworks that are NHWC-native (TensorFlow) made it their default long ago.

PyTorch keeps NCHW as the logical shape contract and encodes NHWC purely in strides: shape $(N, C, H, W)$ with strides $(CHW, 1, WC, C)$. Operator kernels propagate the format through pointwise ops and normalisations, so one conversion at the model boundary survives the whole forward — until an op without a channels_last kernel silently normalises the layout, which is why measuring, not assuming, is the rule. torch.compile folds layout conversions into its generated kernels and will itself choose layouts in some cases; the explicit conversion remains the reliable baseline.

The batch-norm interaction is worth knowing: NHWC reduction over $(N, H, W)$ for per-channel statistics reads contiguous channel vectors, fusing cleanly with the preceding convolution's epilogue — part of why the measured gain exceeds what convolution kernels alone predict. See BatchNorm.

Cost model

The conversion itself is one full tensor copy, $O(NCHW)$, amortised over all layers of all steps if done once per batch. The per-layer penalty for mixed layouts is two transposes per convolution — $O(NCHW)$ each — which is exactly what the two-line recipe deletes. Expected end-to-end training speedup on convolution-dominated networks in AMP: commonly reported in the 1.2–1.9x range depending on architecture and card generation; verify per workload.

References

  • NVIDIA, Deep Learning Performance Documentation — "Tensor Layouts in Memory: NCHW vs NHWC", the kernel-selection rules.
  • PyTorch, channels_last memory format tutorial and RFC (github.com/pytorch/pytorch/issues/19092) — the stride-encoding design and operator coverage.
  • Chetlur et al. (2014), cuDNN: Efficient Primitives for Deep Learning — the implicit-GEMM lowering that makes layout matter.

What to learn next