ffcv_imagenet_loaders_single_gpu.py

BATCH_SIZE = args.batch_size
NUM_WORKERS = 4
MEAN = torch.tensor([0.485, 0.456, 0.406]) * 255
STD = torch.tensor([0.229, 0.224, 0.225]) * 255

train_pipeline = [
    RandomResizedCropRGBImageDecoder((224, 224)),
    RandomHorizontalFlip(),
    ToTensor(),
    ToDevice('cuda:0', non_blocking=True),
    ToTorchImage(),
    Convert(torch.float32),
    torchvision.transforms.Normalize(MEAN, STD),  # Normalize using image statistics
]

val_pipeline = [
    CenterCropRGBImageDecoder((224, 224), 1),
    RandomHorizontalFlip(),
    ToTensor(),
    ToDevice('cuda:0', non_blocking=True),
    ToTorchImage(),
    Convert(torch.float32),
    torchvision.transforms.Normalize(MEAN, STD),  # Normalize using image statistics
]

val_loader = Loader(
    '/gpfs/gpfs0/k.fedyanin/space/imagenet_val.beton',
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    order=OrderOption.SEQUENTIAL,
    pipelines={'image': val_pipeline, 'label': [IntDecoder(), ToTensor(), ToDevice('cuda:0', non_blocking=True)]}
)

train_loader = Loader(
    '/gpfs/gpfs0/k.fedyanin/space/imagenet_train.beton',
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    order=OrderOption.RANDOM,
    pipelines={'image': train_pipeline, 'label': [IntDecoder(), ToTensor(), ToDevice('cuda:0', non_blocking=True)]}
)