from ffcv.writer import DatasetWriter
from ffcv.fields import RGBImageField, IntField
from torchvision.datasets import CIFAR100
from torchvision import datasets, transforms

# Your dataset (`torch.utils.data.Dataset`) of (image, label) pairs
DATASET_DIR = '/d1/xxx/DBQ/'
train_dataset = CIFAR100(DATASET_DIR, train=True, download=True)
test_dataset = CIFAR100(DATASET_DIR, train=False, download=True)

write_path = '/d1/xxx/DBQ/CIFAR100_train.beton'
writer = DatasetWriter(write_path, {
    # Tune options to optimize dataset size, throughput at train-time
    'image': RGBImageField(
        max_resolution=32,
    ),
    'label': IntField()
})

# Write dataset
writer.from_indexed_dataset(train_dataset)

write_path = '/d1/xxx/DBQ/CIFAR100_test.beton'
writer = DatasetWriter(write_path, {
    # Tune options to optimize dataset size, throughput at train-time
    'image': RGBImageField(
        max_resolution=32,
    ),
    'label': IntField()
})
writer.from_indexed_dataset(test_dataset)
