ColossalAI/tests/test_config/sample_config.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-

train_data = dict(
    dataset=dict(
        type="CIFAR10Dataset",
        root="/path/to/data",
        download=True,
        transform_pipeline=[
            dict(type="RandomResizedCrop", size=224),
            dict(type="RandomHorizontalFlip"),
            dict(type="ToTensor"),
            dict(type="Normalize", mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
        ],
    ),
    dataloader=dict(
        batch_size=64,
        pin_memory=True,
        num_workers=4,
        sampler=dict(
            type="DataParallelSampler",
            shuffle=True,
        ),
    ),
)