CPU Optimizer

Optimize models for CPU inference

Overview

optimize_for_cpu prepares a model for efficient CPU inference by combining:

  1. Channels-last memory format — optimizes layout for CNN operations on CPU
  2. Compilation — torch.compile (default) or torch.jit.trace for operator fusion
Backend Speed Compatibility Best For
"compile" Faster Most models Default choice
"trace" Good Requires static shapes Legacy / mobile

source

accelerate_model_for_cpu

def accelerate_model_for_cpu(
    model:torch.nn.modules.module.Module, example_input:torch.Tensor
):

Deprecated: use optimize_for_cpu() instead


source

optimize_for_cpu

def optimize_for_cpu(
    model:torch.nn.modules.module.Module, sample:torch.Tensor, # Sample input for tracing (with batch dim)
    *, backend:str='compile', # "compile" (torch.compile) or "trace" (torch.jit.trace)
    compile_mode:str='default'
)->torch.nn.modules.module.Module:

Optimize model for CPU inference via channels-last layout + compilation

from fasterai.misc.cpu_optimizer import optimize_for_cpu

model = resnet18(pretrained=True)
sample = torch.randn(1, 3, 224, 224)

# Default: torch.compile
optimized = optimize_for_cpu(model, sample)

# Or JIT trace for mobile/static shapes
traced = optimize_for_cpu(model, sample, backend="trace")

Note: accelerate_model_for_cpu is deprecated. Use optimize_for_cpu instead.


See Also

Tests live in nbs/tests/test_cpu_optimizer.ipynb.