Data

This page contains the API reference documentation for the timm.data sub-package, and an overview of the input pipeline it provides: resolving a model’s preprocessing config, building transforms, creating datasets and loaders, Mixup / CutMix, and the NaFlex variable resolution loader.

Usage

A pipeline is built in four steps: resolve_model_data_config reads the model’s preprocessing config, create_transform builds the train or eval transform, create_dataset creates a dataset by name and root, and create_loader builds the transform and a data loader for a dataset.

import os
import tempfile

import torch
from PIL import Image

import timm
from timm.data import create_dataset, create_loader, create_transform, resolve_model_data_config

model = timm.create_model('resnet50', pretrained=False).eval()

# Preprocessing config from the model's pretrained_cfg, the default pretrained tag's config is attached to the model
# with or without pretrained weights.
data_config = resolve_model_data_config(model)
print(data_config)
# {'input_size': (3, 224, 224), 'interpolation': 'bicubic', 'mean': (0.485, 0.456, 0.406),
#  'std': (0.229, 0.224, 0.225), 'crop_pct': 0.95, 'crop_mode': 'center'}

# Eval transform for a PIL image, resize to 224 / crop_pct, center crop 224, to tensor and normalize.
transform = create_transform(**data_config, is_training=False)
img = Image.new('RGB', (320, 240), color=(200, 120, 40))
with torch.no_grad():
    output = model(transform(img).unsqueeze(0))  # (1, 1000)

# A folder dataset, <root>/<split>/<class>/<image>, two classes w/ four generated images each.
root = tempfile.mkdtemp()
for i, cls in enumerate(('cat', 'dog')):
    os.makedirs(os.path.join(root, 'validation', cls))
    for j in range(4):
        path = os.path.join(root, 'validation', cls, f'{j}.jpg')
        Image.new('RGB', (300, 200), color=(60 * i, 40 * j, 128)).save(path)

dataset = create_dataset('', root=root, split='validation')  # finds the 'validation' sub-folder
print(len(dataset), dataset.reader.class_to_idx)  # 8 {'cat': 0, 'dog': 1}

# The loader creates the transform from the same config and assigns it to the dataset.
loader = create_loader(dataset, **data_config, batch_size=4, use_prefetcher=False, num_workers=0)
with torch.no_grad():
    for images, targets in loader:
        output = model(images)  # images (4, 3, 224, 224), targets (4,)

The quickstart walks through the config to transform steps for inference.

The training script (train.py) exposes these options for each part of the pipeline.

Input config overrides (resolve_input_data_config):

--input-size sets the input shape (C, H, W).
--img-size is the square shortcut for --input-size, keeping the model’s channel count.
--in-chans sets the input channels.
--interpolation sets the resize interpolation.
--mean sets the normalization mean.
--std sets the normalization standard deviation.
--crop-pct sets the evaluation crop fraction.

Dataset creation (create_dataset) and target setup:

--dataset selects the dataset and reader.
--data-dir sets the dataset root.
--train-split selects the training split.
--val-split selects the validation split.
--class-map supplies the class mapping.
--input-img-mode sets the image color conversion mode.
--input-key selects the image field.
--target-key selects the label field or fields.
--target-format selects class indices or dense multi-hot targets.

Loader and augmentation options (create_loader) include:

--aa selects the auto-augmentation policy.
--reprob sets the random erasing probability.
--remode selects the random erasing fill mode.
--scale sets the random resized crop scale range.
--ratio sets the random resized crop aspect ratio range.
--hflip sets the horizontal flip probability.
--color-jitter sets color jitter strength.
--train-interpolation sets the training resize interpolation.
--workers sets the number of loader workers.
--no-prefetcher disables the prefetcher.

Mixup and CutMix (Mixup):

--mixup sets the Mixup alpha.
--cutmix sets the CutMix alpha.
--cutmix-minmax sets the CutMix box size range.
--mixup-prob sets the probability of mixing.
--mixup-switch-prob sets the probability of choosing CutMix when both methods are enabled.
--mixup-mode selects batch, pair or element-wise mixing.
--smoothing sets label smoothing.

See Scripts for examples.

Data config

resolve_data_config returns a dict with the keys below, which match the arguments of create_transform and create_loader. Values set in args (a dict, e.g. vars(args) of parsed command line args) take precedence over the pretrained_cfg argument, which defaults to model.pretrained_cfg. resolve_model_data_config is the same with the model as first argument.

Both read the input size and channels from the pretrained cfg, which describes the source checkpoint, not a model created or changed to use a different img_size or in_chans. resolve_input_data_config, used by the training, validation and inference scripts, resolves against the model’s current input instead: a fixed input size model always gets its current size, explicit args that conflict with the model’s channels or fixed size raise an error, and a mean / std that doesn’t match the channel count is replaced by its average (with a warning).

Key	`args` keys	Description	Default
`input_size`	`input_size`, `img_size`, `in_chans`	`(C, H, W)` input size. With `use_test_size=True` the config’s `test_input_size` is used if it has one	`(3, 224, 224)`
`interpolation`	`interpolation`	Resize interpolation, `bilinear`, `bicubic`, …	`bicubic`
`mean`, `std`	`mean`, `std`	Normalization, a single value is repeated for each channel	ImageNet mean / std
`crop_pct`	`crop_pct`	Eval crop, the image is resized to `size / crop_pct` before cropping to `size` (`test_crop_pct` with `use_test_size=True`)	`0.875`
`crop_mode`	`crop_mode`	Eval crop mode, `center`, `squash` or `border`	`center`

The eval crop modes, for an output size and scale_size = floor(size / crop_pct):

`crop_mode`	Behaviour
`center`	Resize the shortest edge to `scale_size` (aspect ratio preserved), center crop to `size`
`squash`	Resize both edges to `scale_size` (aspect ratio not preserved), center crop to `size`
`border`	Resize the longest edge to `scale_size`, center crop and pad with the mean color, the whole image is kept at `crop_pct=1.0`

Transforms

create_transform returns a torchvision Compose. The eval transform resizes and crops as above, then converts to a tensor and normalizes. The training transform is a random resized crop (scale=(0.08, 1.0), ratio=(3/4, 4/3) by default), horizontal flip (hflip=0.5), color jitter (color_jitter=0.4), the optional auto_augment policy, tensor conversion and normalization, and random erasing.

train_crop_mode='rkrc' or 'rkrr' replaces the random resized crop with a random resize that keeps the aspect ratio, followed by a center (rkrc) or random (rkrr) crop or pad.
interpolation='random' (the train.py default for training) picks bilinear or bicubic per image.
no_aug=True reduces training to a resize and center crop.

`auto_augment`	Augmentation
`rand-m9-mstd0.5-inc1`	`RandAugment` via `rand_augment_transform`
`augmix-m5-w4-d2`	AugMix via `augment_and_mix_transform`
`original`, `originalr`, `v0`, `v0r`, `3a`	AutoAugment policy via `auto_augment_transform`

RandAugment option tokens:

m sets magnitude (default 10).
n sets the number of ops per image (default 2).
p sets the probability per op (default 0.5).
mstd sets magnitude noise std.
inc1 selects ops that increase in severity w/ magnitude.
t<name> selects the op set (weights, 3a, 3aw).

AugMix option tokens:

m sets severity (default 3).
w sets width (default 3).
d sets depth (default -1, random 1-3).
b1 enables blended augmentation.
mstd sets magnitude noise std.

AutoAugment accepts mstd, e.g. original-mstd0.5.

Transform behavior:

Color jitter is disabled when auto_augment is set, unless the string contains 3a.
RandomErasing (re_prob, re_count) runs on the normalized tensor. re_mode fills with zeros (const, the default), a random normal color per region (rand) or random normal pixels (pixel, the train.py default).
use_prefetcher=True ends the transform at a uint8 tensor, normalization and random erasing are then done by the loader on the device.
create_transform and create_loader default to bilinear interpolation and ImageNet mean / std, not the model’s values. Pass the resolved data config to match the model’s pretraining.

Datasets

create_dataset selects the dataset implementation and reader from a prefix of name:

`name`	Dataset	`root`	Requires
`''` or no prefix	`ImageDataset`, images in class sub-folders (searched recursively, the leaf folder name is the class), or a single `.tar` file of the same	Dataset folder or `.tar` file
`torch/<name>`	torchvision `cifar10`, `cifar100`, `mnist`, `kmnist`, `fashion_mnist`, `qmnist`, `inaturalist` (`inat`), `places365`, `imagenet` or `image_folder` (`folder`)	torchvision dataset root
`hfds/<name>`	`ImageDataset` on a Hugging Face Hub dataset (or local path) via `datasets.load_dataset`	Cache dir (optional)	`datasets`
`hfids/<name>`	`IterableImageDataset`, the Hugging Face dataset streamed, or downloaded and prepared first with `download=True`	Cache dir (optional)	`datasets`
`tfds/<name>`	`IterableImageDataset` on TensorFlow Datasets, `download=True` downloads and builds	TFDS data dir	`tensorflow`, `tensorflow-datasets`
`wds/<name>`	`IterableImageDataset` on WebDataset tar shards	Shard folder w/ an `_info.json` (or `info.json`) describing the splits	`webdataset`

Dataset options:

split selects the dataset split (default validation). For WebDataset it can also be a shard pattern with a sample count, 'imagenet-train-{0000..1023}.tar|1281167'.
search_split=True (default) makes folder datasets look for the sub-folder of root named split or, failing that, a synonym of a train (train, training) or eval (val, valid, validation, eval, evaluation) split.
add_img_extensions adds to the default image extensions (.png, .jpg and .jpeg).
class_map remaps class names to indices with a dict, a .txt (one class per line) or a .pkl file.

Iterable datasets (hfids, tfds, wds) also take:

is_training enables shuffling.
batch_size rounds each worker’s share of the samples up to whole batches when training.
num_samples sets the sample count, required when the reader can’t determine it (hfids without split info).
seed sets the random seed.
repeats multiplies the samples per epoch, the readers pass over the data that many times.

Call dataset.set_epoch() each epoch to reshuffle.

Other keyword arguments pass through to the reader. For the Hugging Face, TFDS and WebDataset readers:

input_key selects the image field (image by default, for wds a jpg, png or webp image).
target_key selects the label field (label by default, for wds the .cls file or a .json sidecar field named by target_key). A comma separated value reads one binary field per class.
target_format='multihot' reads dense multi-hot vectors instead of lists of class indices for multi-label data.
target_transform=MultiLabelTarget(num_classes) encodes targets as float multi-hot tensors. Pass dense=True to MultiLabelTarget for multi-hot fields.

The folder and tar readers yield single class indices. See multi-label classification for details.

Loaders

create_loader creates the transform with create_transform from its own arguments, assigns it to dataset.transform (replacing any existing transform) and returns a DataLoader. Training loaders (is_training=True) shuffle and drop the last incomplete batch.

Prefetcher. With use_prefetcher=True (default) the dataset yields uint8 tensors, the batches are moved to device (default cuda, on a side stream for CUDA), converted to img_dtype, normalized and randomly erased there. Set device to the training device, or use_prefetcher=False for a plain DataLoader of normalized tensors.
Distributed. distributed=True uses a DistributedSampler for map-style training datasets (or a repeated augmentation sampler w/ num_aug_repeats) and an ordered sampler for eval that pads each process to the same number of samples, which can slightly alter validation results. Iterable datasets shard themselves across processes and workers.
Collate. collate_fn replaces the default (fast_collate w/ the prefetcher, default_collate without), e.g. with a FastCollateMixup.
Augmentation splits. For AugMix w/ the JSD loss (--aug-splits), wrap the dataset in AugMixDataset(dataset, num_splits=n) and pass num_aug_splits=n. Each batch holds a clean split and n - 1 augmented splits, concatenated along the batch dimension by the prefetcher.
Scheduled resolution. input_size_choices (--train-img-sizes) picks the input size per batch, optionally with a batch size per choice (batch_size_choices) and a progressive schedule from the first to the last choice.

Mixup and CutMix

Mixup mixes a batch of normalized images on the device, after the loader. FastCollateMixup takes the same arguments and mixes uint8 images while collating in the loader workers, it requires the prefetcher. Both return dense (batch_size, num_classes) targets with label_smoothing (default 0.1) applied, to be used with a soft target loss such as timm.loss.SoftTargetCrossEntropy, and require an even batch size.

from timm.data import Mixup, FastCollateMixup

# Continuing the example above.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
mixup_args = dict(mixup_alpha=0.8, cutmix_alpha=1.0, label_smoothing=0.1, num_classes=1000)

# Prefetcher enabled (default), mixed in the loader's collate.
loader = create_loader(
    dataset,
    **data_config,
    batch_size=4,
    is_training=True,
    collate_fn=FastCollateMixup(**mixup_args),
    device=device,
)
for images, targets in loader:
    ...  # images normalized on device, targets (4, 1000)

# Prefetcher disabled, mixed in the training loop.
loader = create_loader(dataset, **data_config, batch_size=4, is_training=True, use_prefetcher=False)
mixup_fn = Mixup(**mixup_args)
for images, targets in loader:
    images, targets = mixup_fn(images.to(device), targets.to(device))

Mixup and CutMix controls:

mixup_alpha > 0 enables Mixup.
cutmix_alpha > 0 enables CutMix.
cutmix_minmax enables CutMix with a box size range.
switch_prob sets the probability of choosing CutMix when both methods are enabled.
prob sets the probability of mixing at all.
mode draws the parameters once per batch, per pair of elements or per elem (element).
multi_label=True mixes dense multi-label targets, smoothing each label towards 0.5.
mixup_enabled = False turns mixing off, e.g. for the final epochs (--mixup-off-epoch). For FastCollateMixup it is read in the worker processes, persistent workers keep the value they started with.

NaFlex variable resolution

create_naflex_loader loads images as variable length sequences of patches, w/o resizing to a fixed square size, for models that take NaFlex patch inputs (e.g. the naflexvit_* models). Each batch is a dict of patches, patch_coord, patch_valid and seq_len, padded to the batch’s sequence length.

Training. The dataset is wrapped in a NaFlexMapDatasetWrapper (map-style datasets only), which picks one of train_seq_lens per batch and random resize crops each image to fit that many patches. batch_size is the batch size at the longest sequence length, shorter lengths get proportionally larger batches (rounded down to a multiple of 8) so the tokens per batch stay roughly constant. patch_size_choices (w/ patch_size_choice_probs) also varies the patch size per batch, for models that support it.
Eval. Images are resized to fit max_seq_len patches at the given patch_size (required for eval) and padded to max_seq_len.
Mixup. Pass a NaFlexMixup as mixup_fn, it mixes within each batch and does not support mode or cutmix_minmax.

The train.py options are:

--naflex-loader enables the NaFlex loader.
--naflex-train-seq-lens sets the training sequence lengths.
--naflex-max-seq-len sets the maximum evaluation sequence length.
--naflex-patch-sizes sets the patch sizes for variable patch size training.
--naflex-patch-size-probs sets the probabilities for those patch sizes.
--naflex-loss-scale selects loss scaling by batch size (none, sqrt or linear).

timm

Data

Usage

Data config

Transforms

Datasets

Loaders

Mixup and CutMix

NaFlex variable resolution

API

Data config

timm.data.resolve_data_config

timm.data.resolve_model_data_config

timm.data.resolve_input_data_config

Transforms

timm.data.create_transform

timm.data.rand_augment_transform

timm.data.auto_augment_transform

timm.data.auto_augment.augment_and_mix_transform

class timm.data.random_erasing.RandomErasing

Datasets

timm.data.create_dataset

class timm.data.ImageDataset

class timm.data.IterableImageDataset

class timm.data.AugMixDataset

class timm.data.MultiLabelTarget

Loaders

timm.data.create_loader

Mixup and CutMix

class timm.data.Mixup

class timm.data.FastCollateMixup

NaFlex

timm.data.create_naflex_loader

class timm.data.NaFlexMapDatasetWrapper

set_epoch

class timm.data.NaFlexMixup

set_epoch