🎓 A1Learner
🔍

PyTorch Channels Last

Channels last is a layout that keeps each pixel's channel values side by side in memory. Convolution code on CPU can read that layout faster than the usual one.

A plain tensor is not channels last

A fresh image tensor uses the usual layout. The check prints False.

import torch

torch.manual_seed(0)
x = torch.randn(2, 3, 8, 8)

print("Plain input is channels last:", x.is_contiguous(memory_format=torch.channels_last))

Converting keeps the values

The conversion rearranges memory and touches no values. Both prints confirm the change and the sameness.

x_fast = x.to(memory_format=torch.channels_last)
print("Converted input is channels last:", x_fast.is_contiguous(memory_format=torch.channels_last))
print("Values unchanged:", torch.equal(x, x_fast))

A converted convolution agrees

A convolution in channels last must answer like the plain one. The outputs agree within a small tolerance.

conv = torch.nn.Conv2d(3, 4, 3, padding=1)
conv_fast = torch.nn.Conv2d(3, 4, 3, padding=1)
conv_fast.load_state_dict(conv.state_dict())
conv_fast = conv_fast.to(memory_format=torch.channels_last)
with torch.no_grad():
    a = conv(x)
    b = conv_fast(x_fast)
print("Conv outputs agree:", torch.allclose(a, b, atol=1e-5))

Example

import time

import torch


class ConvNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = torch.nn.Conv2d(16, 16, 3, padding=1)
        self.head = torch.nn.Linear(16, 5)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        return self.head(x.mean(dim=(2, 3)))


def timed(model, x, reps=10):
    with torch.no_grad():
        for _ in range(3):
            model(x)
        start = time.perf_counter()
        out = None
        for _ in range(reps):
            out = model(x)
        ms = (time.perf_counter() - start) * 1000 / reps
    return out, ms


torch.manual_seed(0)
plain = ConvNet()
plain.eval()
fast = ConvNet()
fast.load_state_dict(plain.state_dict())
fast.eval()
fast = fast.to(memory_format=torch.channels_last)

x = torch.randn(16, 3, 32, 32)
x_fast = x.to(memory_format=torch.channels_last)

plain_out, plain_ms = timed(plain, x)
fast_out, fast_ms = timed(fast, x_fast)

print("Input uses channels last:", x_fast.is_contiguous(memory_format=torch.channels_last))
print("Outputs agree:", torch.allclose(plain_out, fast_out, atol=1e-5))
print(f"Plain format per run: {plain_ms:.0f} ms")
print(f"Channels last per run: {fast_ms:.0f} ms")

Key points

  • Memory format: the order a tensor's values are laid out in memory
  • Use .to(memory_format=torch.channels_last) to convert a model or a tensor
  • Both the model and its input need the format for convolutions to use it
  • Conversion changes the layout only, and the values stay the same
  • is_contiguous(memory_format=torch.channels_last) reports whether a tensor uses the layout
  • At the small size in the closing example, the two formats timed within noise of each other
  • In short, channels last rearranges memory so CPU convolutions can read it in a friendlier order
  • The gain grows with larger convolution workloads

In this example, one small conv net is timed in both memory formats. The second copy loads the same weights, then converts to channels last. The input converts too, and the layout check prints True. The outputs agree within a small tolerance, so the format changed no answers. On the recorded run, channels last was slightly slower, and a second run flipped the order. Both times vary on every run. At this small size the honest verdict is a tie, and the format earns its keep on larger conv workloads.

Note

A small model on a small image may show no gain at all, as the closing example shows.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does converting a tensor to channels last change?