Training CIFAR-10 in 36 seconds on a single A100

In this example, we’ll show how to use FFCV and the ResNet-9 architecture in order to train a CIFAR-10 classifier to 92.6% accuracy in 36 seconds on a single NVIDIA A100 GPU.

We also provide the code here and the corresponding script here.

Here, we show a step by step walkthrough. First, we import torch and necessary components from ffcv.

from typing import List

import torch as ch
import torchvision

from ffcv.fields import IntField, RGBImageField
from ffcv.fields.decoders import IntDecoder, SimpleRGBImageDecoder
from ffcv.loader import Loader, OrderOption
from ffcv.pipeline.operation import Operation
from ffcv.transforms import RandomHorizontalFlip, Cutout, \
    RandomTranslate, Convert, ToDevice, ToTensor, ToTorchImage
from ffcv.transforms.common import Squeeze
from ffcv.writer import DatasetWriter

Step 1: Create an FFCV-compatible CIFAR-10 dataset

First, we’ll use DatasetWriter to convert torchvision.datasets.CIFAR10 to FFCV format. (See Writing datasets for more details.) We use a single RGBImageField to store the image and a single IntField to store the label.

datasets = {
    'train': torchvision.datasets.CIFAR10('/tmp', train=True, download=True),
    'test': torchvision.datasets.CIFAR10('/tmp', train=False, download=True)

for (name, ds) in datasets.items():
    writer = DatasetWriter(f'/tmp/cifar_{name}.beton', {
        'image': RGBImageField(),
        'label': IntField()

Step 2: Create data loaders

Next, we construct FFCV dataloaders from the .beton dataset file created above. (See Making an FFCV dataloader for more details.)

For the training set, we use a set of standard data augmentations: random horizontal flip, random translation, and Cutout. Note that the transformation pipeline can consist of both standard transforms from ffcv and other sources such as any torch.nn.Module.

# Note that statistics are wrt to uin8 range, [0,255].
CIFAR_MEAN = [125.307, 122.961, 113.8575]
CIFAR_STD = [51.5865, 50.847, 51.255]


loaders = {}
for name in ['train', 'test']:
    label_pipeline: List[Operation] = [IntDecoder(), ToTensor(), ToDevice('cuda:0'), Squeeze()]
    image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]

    # Add image transforms and normalization
    if name == 'train':
            Cutout(8, tuple(map(int, CIFAR_MEAN))), # Note Cutout is done before normalization.
        ToDevice('cuda:0', non_blocking=True),
        torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD),

    # Create loaders
    loaders[name] = Loader(f'/tmp/cifar_{name}.beton',
                            drop_last=(name == 'train'),
                            pipelines={'image': image_pipeline,
                                       'label': label_pipeline})

Step 3: Setup model architecture and optimization parameters

For the model, we use a custom ResNet-9 architecture from KakaoBrain.

class Mul(ch.nn.Module):
    def __init__(self, weight):
        super(Mul, self).__init__()
        self.weight = weight
    def forward(self, x): return x * self.weight

class Flatten(ch.nn.Module):
    def forward(self, x): return x.view(x.size(0), -1)

class Residual(ch.nn.Module):
    def __init__(self, module):
        super(Residual, self).__init__()
        self.module = module
    def forward(self, x): return x + self.module(x)

def conv_bn(channels_in, channels_out, kernel_size=3, stride=1, padding=1, groups=1):
    return ch.nn.Sequential(
            ch.nn.Conv2d(channels_in, channels_out,
                         kernel_size=kernel_size, stride=stride, padding=padding,
                         groups=groups, bias=False),

model = ch.nn.Sequential(
    conv_bn(3, 64, kernel_size=3, stride=1, padding=1),
    conv_bn(64, 128, kernel_size=5, stride=2, padding=2),
    Residual(ch.nn.Sequential(conv_bn(128, 128), conv_bn(128, 128))),
    conv_bn(128, 256, kernel_size=3, stride=1, padding=1),
    Residual(ch.nn.Sequential(conv_bn(256, 256), conv_bn(256, 256))),
    conv_bn(256, 128, kernel_size=3, stride=1, padding=0),
    ch.nn.AdaptiveMaxPool2d((1, 1)),
    ch.nn.Linear(128, NUM_CLASSES, bias=False),
model = model.to(memory_format=ch.channels_last).cuda()

Note the ch.channels_last option when we put the model on GPU.

Next, we define the optimizer and hyperparameters. We use standard SGD on the cross entropy loss with label smoothing and a cyclic learning rate schedule (triangular).

import numpy as np
from torch.cuda.amp import GradScaler, autocast
from torch.nn import CrossEntropyLoss
from torch.optim import SGD, lr_scheduler


opt = SGD(model.parameters(), lr=.5, momentum=0.9, weight_decay=5e-4)
iters_per_epoch = 50000 // BATCH_SIZE
lr_schedule = np.interp(np.arange((EPOCHS+1) * iters_per_epoch),
                        [0, 5 * iters_per_epoch, EPOCHS * iters_per_epoch],
                        [0, 1, 0])
scheduler = lr_scheduler.LambdaLR(opt, lr_schedule.__getitem__)
scaler = GradScaler()
loss_fn = CrossEntropyLoss(label_smoothing=0.1)

Step 4: Train and evaluate the model

Finally, we’re ready to train our model.

from tqdm import tqdm

for ep in range(EPOCHS):
    for ims, labs in tqdm(loaders['train']):
        with autocast():
            out = model(ims)
            loss = loss_fn(out, labs)


with ch.no_grad():
    total_correct, total_num = 0., 0.
    for ims, labs in tqdm(loaders['test']):
        with autocast():
            out = (model(ims) + model(ch.fliplr(ims))) / 2. # Test-time augmentation
            total_correct += out.argmax(1).eq(labs).sum().cpu().item()
            total_num += ims.shape[0]

    print(f'Accuracy: {total_correct / total_num * 100:.1f}%')

Wrapping up

It’s that simple! In this tutorial, we used FFCV to train a CIFAR-10 classifier to 92.6% accuracy in 36 seconds.

For a different example using FFCV to speed up training, see Large-Scale Linear Regression.