Implementation:Microsoft DeepSpeedExamples Dataset Loading With Augmentation

Knowledge Sources	Microsoft_DeepSpeedExamples
Domains	Computer Vision, Data Augmentation
Last Updated	2026-02-07 12:00 GMT

Overview

Provides dataset loading utilities for CIFAR-10, CIFAR-100, and ImageNet with configurable image augmentation strategies for Vision Transformer fine-tuning.

Description

This module implements a suite of functions for loading vision datasets with appropriate preprocessing pipelines for both standard CNN training and Vision Transformer (ViT) fine-tuning. The top-level get_dataset function dispatches to dataset-specific loaders based on the dataset name, automatically detecting ViT-specific variants (e.g., cifar10vit224, cifar100vit384) by parsing the image size from the name suffix.

Two augmentation strategies are supported through the get_aug function: a 'large' strategy designed for ViT models (default 224px with resize, random resized crop, and horizontal flip for training) and a 'small' strategy for standard CNN models (default 32px with random crop and horizontal flip). The get_transform function composes these augmentations with appropriate normalization constants -- ImageNet statistics for large augmentations and CIFAR statistics for small augmentations.

The ViT-specific loader get_cifar_vit applies specialized preprocessing pipelines with CIFAR-specific normalization values and supports arbitrary image sizes. The get_imagenet_vit function handles ImageNet directory-based datasets with train/val splits. All dataset loaders return torchvision datasets enriched with nchannels and imsize attributes for downstream model configuration.

Usage

Use this module when setting up data pipelines for ViT fine-tuning experiments on CIFAR or ImageNet datasets. It is designed to be called from ViT training scripts that need configurable image resolution and augmentation strategies.

Code Reference

Source Location

• Repository: Microsoft_DeepSpeedExamples
• File: training/data_efficiency/vit_finetuning/utils/get_data.py
• Lines: 1-167

Signature

def get_dataset(dataset_name, data_dir, split, rand_fraction=None,
                clean=False, transform=None, imsize=None,
                bucket='pytorch-data', **kwargs):
    ...

def get_aug(split, imsize=None, aug='large'):
    ...

def get_transform(split, normalize=None, transform=None,
                  imsize=None, aug='large'):
    ...

def get_cifar10(dataset_name, data_dir, split, transform=None,
                imsize=None, bucket='pytorch-data', **kwargs):
    ...

def get_cifar100(dataset_name, data_dir, split, transform=None,
                 imsize=None, bucket='pytorch-data', **kwargs):
    ...

def get_cifar_vit(dataset_name, data_dir, split, transform=None,
                  imsize=None, bucket='pytorch-data', **kwargs):
    ...

def get_imagenet_vit(dataset_name, data_dir, split, transform=None,
                     imsize=None, bucket='pytorch-data', **kwargs):
    ...

Import

from utils.get_data import get_dataset, get_transform, get_aug

I/O Contract

Inputs

Name	Type	Required	Description
dataset_name	str	Yes	Name of the dataset (e.g., 'cifar10', 'cifar100', 'cifar10vit224', 'cifar100vit384')
data_dir	str	Yes	Path to the data directory for downloading/loading
split	str	Yes	Dataset split: 'train' or 'test'
rand_fraction	float	No	Fraction of random labels (used in noisy label experiments)
transform	transforms.Compose	No	Custom transform to override defaults
imsize	int	No	Target image size (inferred from dataset_name if not specified)
bucket	str	No	Data bucket name (default: 'pytorch-data')

Outputs

Name	Type	Description
dataset	torchvision.datasets.*	A dataset object with added nchannels and imsize attributes

Usage Examples

from utils.get_data import get_dataset

# Load CIFAR-10 with ViT-224 augmentation
train_dataset = get_dataset('cifar10vit224', data_dir='./data', split='train')
print(f"Channels: {train_dataset.nchannels}, Image size: {train_dataset.imsize}")

# Load standard CIFAR-100
test_dataset = get_dataset('cifar100', data_dir='./data', split='test')

# Use in DataLoader
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=64, shuffle=True, num_workers=4
)