FFCV

https://docs.ffcv.io
Admin
FFCV is a drop-in data loading system designed to dramatically increase data throughput and...
Tokens:29,865
Snippets:254
Trust Score:7
Update:1 month ago
Show doc for...
Tokens:
Raw
### Configuring FFCV Loader with Image Decoders

Source: https://docs.ffcv.io/_sources/working_with_images.rst.txt

Example of initializing an FFCV Loader with specific image decoding pipelines for different fields.

```python
writer = Loader('my_file.beton',
    batch_size=15,
    num_workers=10
    pipelines = {
        'image': [RandomResizedCropRGBImageDecoder((224, 224))]
        'other_image_field': [CenterCropRGBImageDecoder((224, 224), 224/256)]
    }
)
```

--------------------------------

### Benchmark output

Source: https://docs.ffcv.io/ffcv_examples/custom_transforms.html

Example output showing the performance difference between pipelines with and without the custom augmentation.

```text
Method: with | Shape: torch.Size([512, 16, 16, 3]) | Time per epoch: 0.06596s
Method: without | Shape: torch.Size([512, 32, 32, 3]) | Time per epoch: 0.02828s
```

--------------------------------

### Implementing declare_state_and_memory

Source: https://docs.ffcv.io/ffcv_examples/custom_transforms.html

Example implementation for defining output state and pre-allocating memory for a transform that modifies image dimensions.

```python
from ffcv.pipeline.allocation_query import AllocationQuery
from dataclasses import replace

# Inside the MaybeBrighten class:
def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:
    h, w, c = previous_state.shape
    new_shape = (h // 2, w // 2, c)

    # Everything in the state stays the same other than the shape
    # States are immutable, so we have to edit them using the
    # dataclasses.replace function
    new_state = replace(previous_state, shape=new_shape)

    # We need to allocate memory for the new images
    # so below, we ask for a memory allocation whose width and height is
    # half the original image, with the same type
    # (shape=(,)) of the same type as the image data
    mem_allocation = AllocationQuery(new_shape, previous_state.dtype)
    return (new_state, mem_allocation)
```

--------------------------------

### Install FFCV via Conda

Source: https://docs.ffcv.io/_sources/index.rst.txt

Use this command to create a dedicated environment and install FFCV with its necessary dependencies.

```bash
conda create -y -n ffcv python=3.9 cupy pkg-config compilers libjpeg-turbo opencv pytorch torchvision cudatoolkit=11.3 numba -c pytorch -c conda-forge
conda activate ffcv
pip install ffcv
```

--------------------------------

### Optimized benchmark output

Source: https://docs.ffcv.io/ffcv_examples/custom_transforms.html

Example output demonstrating the performance improvement after applying parallel compilation.

```text
Method: with | Shape: torch.Size([512, 16, 16, 3]) | Time per epoch: 0.03404s
Method: without | Shape: torch.Size([512, 32, 32, 3]) | Time per epoch: 0.02703s
```

--------------------------------

### Write indexable dataset to .beton

Source: https://docs.ffcv.io/_sources/writing_datasets.rst.txt

Execute the conversion process after initializing the writer.

```python
writer.from_indexed_dataset(my_dataset)
```

--------------------------------

### Initialize FFCV Loader

Source: https://docs.ffcv.io/_sources/making_dataloaders.rst.txt

Instantiate the Loader class with dataset path, batch size, workers, ordering, and pipelines.

```python
loader = Loader('/path/to/dataset.beton',
                batch_size=BATCH_SIZE,
                num_workers=NUM_WORKERS,
                order=ORDERING,
                pipelines=PIPELINES)
```

--------------------------------

### Configure Optimizer and Scheduler

Source: https://docs.ffcv.io/_sources/ffcv_examples/cifar10.rst.txt

Sets up SGD with a triangular cyclic learning rate schedule and label smoothing for training.

```python
import numpy as np
from torch.cuda.amp import GradScaler, autocast
from torch.nn import CrossEntropyLoss
from torch.optim import SGD, lr_scheduler

EPOCHS = 24

opt = SGD(model.parameters(), lr=.5, momentum=0.9, weight_decay=5e-4)
iters_per_epoch = 50000 // BATCH_SIZE
lr_schedule = np.interp(np.arange((EPOCHS+1) * iters_per_epoch),
                        [0, 5 * iters_per_epoch, EPOCHS * iters_per_epoch],
                        [0, 1, 0])
scheduler = lr_scheduler.LambdaLR(opt, lr_schedule.__getitem__)
scaler = GradScaler()
loss_fn = CrossEntropyLoss(label_smoothing=0.1)
```

--------------------------------

### Initialize an FFCV Loader

Source: https://docs.ffcv.io/_sources/making_dataloaders.rst.txt

Configure the final Loader instance with specific pipelines for covariates and labels.

```python
loader = Loader('/path/to/dataset.beton',
                batch_size=BATCH_SIZE,
                num_workers=NUM_WORKERS,
                order=OrderOption.RANDOM,
                pipelines={
                  'covariate': [NDArrayDecoder(), ToTensor(), AddGaussianNoise(0.1)],
                  'label': [FloatDecoder(), ToTensor()]
                })
```

--------------------------------

### Initialize transformation function implementation

Source: https://docs.ffcv.io/_sources/ffcv_examples/custom_transforms.rst.txt

Begin the implementation of the transformation logic using the generate_code factory function.

```python
import numpy as np
```

--------------------------------

### Configure DatasetWriter for indexable data

Source: https://docs.ffcv.io/_sources/writing_datasets.rst.txt

Initialize the writer with a path and a dictionary mapping keys to FFCV fields, ensuring the order matches the dataset's __getitem__ output.

```python
from ffcv.fields import NDArrayField, FloatField

writer = DatasetWriter(write_path, {
    'covariate': NDArrayField(shape=(d,), dtype=np.dtype('float32')),
    'label': FloatField(),

}, num_workers=16)
```

--------------------------------

### Write Webdataset to .beton

Source: https://docs.ffcv.io/_sources/writing_datasets.rst.txt

Configure the writer with appropriate fields and process the shards using the defined pipeline.

```python
from ffcv.fields import RGBImageField, IntField

writer = DatasetWriter(write_path, {
    'image': RGBImageField()
    'label': IntField(),

}, num_workers=40)

writer.from_webdataset(my_shards, pipeline)
```

--------------------------------

### Configure FFCV data loaders

Source: https://docs.ffcv.io/_sources/ffcv_examples/cifar10.rst.txt

Sets up training and testing loaders with specific image pipelines, including augmentations like Cutout and normalization.

```python
# Note that statistics are wrt to uin8 range, [0,255].
CIFAR_MEAN = [125.307, 122.961, 113.8575]
CIFAR_STD = [51.5865, 50.847, 51.255]

BATCH_SIZE = 512

loaders = {}
for name in ['train', 'test']:
    label_pipeline: List[Operation] = [IntDecoder(), ToTensor(), ToDevice('cuda:0'), Squeeze()]
    image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]

    # Add image transforms and normalization
    if name == 'train':
        image_pipeline.extend([
            RandomHorizontalFlip(),
            RandomTranslate(padding=2),
            Cutout(8, tuple(map(int, CIFAR_MEAN))), # Note Cutout is done before normalization.
        ])
    image_pipeline.extend([
        ToTensor(),
        ToDevice('cuda:0', non_blocking=True),
        ToTorchImage(),
        Convert(ch.float16),
        torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
    ])

    # Create loaders
    loaders[name] = Loader(f'/tmp/cifar_{name}.beton',
                            batch_size=BATCH_SIZE,
                            num_workers=8,
                            order=OrderOption.RANDOM,
                            drop_last=(name == 'train'),
                            pipelines={'image': image_pipeline,
                                       'label': label_pipeline})
```

--------------------------------

### RGBImageField.__init__

Source: https://docs.ffcv.io/_modules/ffcv/fields/rgb_image.html

Initializes a new RGBImageField instance with specific configuration for image storage and compression.

```APIDOC
## RGBImageField.__init__(write_mode='raw', max_resolution=None, smart_threshold=None, jpeg_quality=90, compress_probability=0.5)

### Description
Initializes the field configuration for RGB images, defining how images are processed and stored in the dataset.

### Parameters
- **write_mode** (str, optional) - How to write image data: 'raw', 'jpg', 'smart', or 'proportion'. Default: 'raw'.
- **max_resolution** (int, optional) - Maximum side length for resizing images before saving.
- **smart_threshold** (int, optional) - Threshold for deciding between raw and JPEG compression when mode is 'smart'.
- **jpeg_quality** (int, optional) - Quality parameter for JPEG encoding (1-100). Default: 90.
- **compress_probability** (float, optional) - Probability of JPEG compression when mode is 'proportion'. Default: 0.5.
```

--------------------------------

### Configure Dataset Ordering

Source: https://docs.ffcv.io/_sources/making_dataloaders.rst.txt

Define the loading order using OrderOption to manage RAM usage and randomness.

```python
from ffcv.loader import OrderOption

# Truly random shuffling (shuffle=True in PyTorch)
ORDERING = OrderOption.RANDOM

# Unshuffled (i.e., served in the order the dataset was written)
ORDERING = OrderOption.SEQUENTIAL

# Memory-efficient but not truly random loading
# Speeds up loading over RANDOM when the whole dataset does not fit in RAM!
ORDERING = OrderOption.QUASI_RANDOM
```

--------------------------------

### Import DatasetWriter

Source: https://docs.ffcv.io/_sources/writing_datasets.rst.txt

Include the DatasetWriter class in your script to begin the conversion process.

```python
from ffcv.writer import DatasetWriter
```

--------------------------------

### Initialize DatasetWriter with RGBImageField

Source: https://docs.ffcv.io/_sources/working_with_images.rst.txt

Configures a DatasetWriter to store images using a hybrid compression strategy with specified quality and resolution constraints.

```python
writer = DatasetWriter('my_file.beton', {
        # Roughly 25% of the images will be stored in raw and the other in jpeg
        'image': RGBImageField(
            write_mode='proportion', # Randomly compress
            compress_probability=0.25, # Compress a random 1/4 of the dataset
            max_resolution=(256, 256), # Resize anything above 256 to 256
            jpeg_quality=50  # Use 50% quality when compressing an image using JPG
        ),
        'label': IntField()
    },
)
```

--------------------------------

### ModuleWrapper(module: torch.nn.Module)

Source: https://docs.ffcv.io/_modules/ffcv/transforms/module.html

Initializes a wrapper for a torch.nn.Module to be used as an FFCV operation.

```APIDOC
## ModuleWrapper

### Description
Wraps a torch.nn.Module to be used as a transformation within an FFCV pipeline.

### Constructor
`ModuleWrapper(module: torch.nn.Module)`

### Parameters
- **module** (torch.nn.Module) - Required - The PyTorch module to apply during the transformation process.

### Methods
- **generate_code() -> Callable**: Returns a function that applies the wrapped module to the input.
- **declare_state_and_memory(previous_state: State) -> Tuple[State, Optional[AllocationQuery]]**: Declares the state and memory requirements for the operation.
```

--------------------------------

### Import FFCV and PyTorch components

Source: https://docs.ffcv.io/_sources/ffcv_examples/cifar10.rst.txt

Initial imports required for FFCV dataset handling, pipeline operations, and transformations.

```python
from typing import List

import torch as ch
import torchvision

from ffcv.fields import IntField, RGBImageField
from ffcv.fields.decoders import IntDecoder, SimpleRGBImageDecoder
from ffcv.loader import Loader, OrderOption
from ffcv.pipeline.operation import Operation
from ffcv.transforms import RandomHorizontalFlip, Cutout, \
    RandomTranslate, Convert, ToDevice, ToTensor, ToTorchImage
from ffcv.transforms.common import Squeeze
from ffcv.writer import DatasetWriter
```

--------------------------------

### RandomResizedCrop(scale, ratio, size)

Source: https://docs.ffcv.io/_modules/ffcv/transforms/random_resized_crop.html

Initializes the RandomResizedCrop transformation with specified scale, aspect ratio, and output size parameters.

```APIDOC
## class RandomResizedCrop(scale, ratio, size)

### Description
Crop a random portion of an image with a random aspect ratio and resize it to a given size. Note that it is generally recommended to use the ResizedCropRGBImageDecoder class instead for better performance.

### Parameters
- **scale** (Tuple[float, float]) - Required - Lower and upper bounds for the ratio of random area of the crop.
- **ratio** (Tuple[float, float]) - Required - Lower and upper bounds for random aspect ratio of the crop.
- **size** (int) - Required - Side length of the output.
```

--------------------------------

### Define an indexable dataset

Source: https://docs.ffcv.io/_sources/writing_datasets.rst.txt

Create a class that implements __len__ and __getitem__ to provide data samples as tuples.

```python
import numpy as np

class LinearRegressionDataset:
    def __init__(self, N, d):
        self.X = np.random.randn(N, d)
        self.Y = np.random.randn(N)

    def __getitem__(self, idx):
        return (self.X[idx].astype('float32'), self.Y[idx])

    def __len__(self):
        return len(self.X)

N, d = (100, 6)
dataset = LinearRegressionDataset(N, d)
```

--------------------------------

### Replace PyTorch loader with FFCV loader

Source: https://docs.ffcv.io/_sources/quickstart.rst.txt

Configure data pipelines with decoders and transforms, then initialize the FFCV Loader to replace standard PyTorch data loading.

```python
from ffcv.loader import Loader, OrderOption
from ffcv.transforms import ToTensor, ToDevice, ToTorchImage, Cutout
from ffcv.fields.decoders import IntDecoder, RandomResizedCropRGBImageDecoder

# Random resized crop
decoder = RandomResizedCropRGBImageDecoder((224, 224))

# Data decoding and augmentation
image_pipeline = [decoder, Cutout(), ToTensor(), ToTorchImage(), ToDevice(0)]
label_pipeline = [IntDecoder(), ToTensor(), ToDevice(0)]

# Pipeline for each data field
pipelines = {
    'image': image_pipeline,
    'label': label_pipeline
}

# Replaces PyTorch data loader (`torch.utils.data.Dataloader`)
loader = Loader(write_path, batch_size=bs, num_workers=num_workers,
                order=OrderOption.RANDOM, pipelines=pipelines)

# rest of training / validation proceeds identically
for epoch in range(epochs):
    ...
```

--------------------------------

### Writing FFCV Dataset

Source: https://docs.ffcv.io/ffcv_examples/linear_regression.html

Converts existing data into the FFCV .beton format using DatasetWriter.

```python
from ffcv.fields import NDArrayField, FloatField

class LinearRegressionDataset:
    def __getitem__(self, idx):
        return (X[idx], np.array(Y[idx]).astype('float32'))

    def __len__(self):
        return len(X)

writer = DatasetWriter('/tmp/linreg_data.beton', {
    'covariate': NDArrayField(shape=(D,), dtype=np.dtype('float32')),
    'label': NDArrayField(shape=(1,), dtype=np.dtype('float32')),
}, num_workers=16)

writer.from_indexed_dataset(LinearRegressionDataset())
```

--------------------------------

### Convert dataset to FFCV format

Source: https://docs.ffcv.io/_sources/quickstart.rst.txt

Use DatasetWriter to convert an existing PyTorch dataset into the .beton format, specifying field types for optimization.

```python
from ffcv.writer import DatasetWriter
from ffcv.fields import RGBImageField, IntField

# Your dataset (`torch.utils.data.Dataset`) of (image, label) pairs
my_dataset = make_my_dataset()
write_path = '/output/path/for/converted/ds.beton'

# Pass a type for each data field
writer = DatasetWriter(write_path, {
    # Tune options to optimize dataset size, throughput at train-time
    'image': RGBImageField(
        max_resolution=256
    ),
    'label': IntField()
})

# Write dataset
writer.from_indexed_dataset(my_dataset)
```

--------------------------------

### Naive PyTorch TensorDataset Implementation

Source: https://docs.ffcv.io/_sources/ffcv_examples/linear_regression.rst.txt

Standard data loading approach that keeps the entire dataset in RAM.

```python
from torch.utils.data import TensorDataset, DataLoader

X, W, b, Y = pkl.load(open('/tmp/linreg_data.pkl', 'rb'))
dataset = TensorDataset(ch.tensor(X), ch.tensor(Y))
train_loader = DataLoader(dataset, num_workers=8, shuffle=True)
# ... rest of code as above
```

--------------------------------

### Generate Fake Dataset

Source: https://docs.ffcv.io/_sources/ffcv_examples/linear_regression.rst.txt

Creates a large synthetic dataset of 1,000,000 inputs and saves it to a pickle file.

```python
import numpy as np

# 1,000,000 inputs each of dimension 10,000 = 40GB of data
N, D = 1000000, 10000
X = np.random.rand(N, D).astype('float32')
# Ground-truth vector
W, b = np.random.rand(D).astype('float32'), np.random.rand()
# Response variables
Y = X @ W + np.random.randn(N)
# Save the dataset:
pkl.dump((X, W, b, Y), open('/tmp/linreg_data.pkl', 'wb'))
```

--------------------------------

### Test custom augmentation performance

Source: https://docs.ffcv.io/ffcv_examples/custom_transforms.html

A script to benchmark the performance impact of a custom transformation within an FFCV loader pipeline.

```python
ds = torchvision.datasets.CIFAR10('/tmp', train=True, download=True)
writer = DatasetWriter('/tmp/cifar.beton', {'image': RGBImageField(),
                                            'label': IntField()})
writer.from_indexed_dataset(ds)

BATCH_SIZE = 512
image_pipelines = {
    'with': [SimpleRGBImageDecoder(), PickACorner(), ToTensor()],
    'without': [SimpleRGBImageDecoder(), ToTensor()]
}

for name, pipeline in image_pipelines.items():
    loader = Loader(f'/tmp/cifar.beton', batch_size=BATCH_SIZE,
                    num_workers=8, order=OrderOption.RANDOM,
                    drop_last=True, pipelines={'image': pipeline})

    # First epoch includes compilation time
    for ims, labs in loader: pass
    start_time = time.time()
    for _ in range(100):
        for ims, labs in loader: pass
    print(f'Method: {name} | Shape: {ims.shape} | Time per epoch: {(time.time() - start_time) / 100:.4f}s')
```

--------------------------------

### Cutout(crop_size, fill)

Source: https://docs.ffcv.io/_modules/ffcv/transforms/cutout.html

Initializes the Cutout augmentation operation.

```APIDOC
## class Cutout(crop_size, fill=(0, 0, 0))

### Description
Cutout data augmentation (https://arxiv.org/abs/1708.04552). This operation randomly masks out a square region of the input images.

### Parameters
- **crop_size** (int) - Required - Size of the random square to cut out.
- **fill** (Tuple[int, int, int]) - Optional - An RGB color ((0, 0, 0) by default) to fill the cutout square with. Useful for when a normalization layer follows cutout.
```

--------------------------------

### Write Linear Regression Dataset to FFCV

Source: https://docs.ffcv.io/_sources/ffcv_examples/linear_regression.rst.txt

Initializes a DatasetWriter to store linear regression data in a .beton file format.

```python
writer = DatasetWriter('/tmp/linreg_data.beton', {
    'covariate': NDArrayField(shape=(D,), dtype=np.dtype('float32')),
    'label': NDArrayField(shape=(1,), dtype=np.dtype('float32')),
}, num_workers=16)

writer.from_indexed_dataset(LinearRegressionDataset())
```

--------------------------------

### Convert CIFAR-10 to FFCV format

Source: https://docs.ffcv.io/_sources/ffcv_examples/cifar10.rst.txt

Uses DatasetWriter to convert torchvision CIFAR-10 datasets into .beton files for efficient loading.

```python
datasets = {
    'train': torchvision.datasets.CIFAR10('/tmp', train=True, download=True),
    'test': torchvision.datasets.CIFAR10('/tmp', train=False, download=True)
}

for (name, ds) in datasets.items():
    writer = DatasetWriter(f'/tmp/cifar_{name}.beton', {
        'image': RGBImageField(),
        'label': IntField()
    })
    writer.from_indexed_dataset(ds)
```

--------------------------------

### View

Source: https://docs.ffcv.io/_modules/ffcv/transforms/ops.html

Views an array using numpy.view or torch.view.

```APIDOC
## View

### Description
Provides a view of the array using the underlying numpy or torch view functionality.

### Constructor
`View(target_dtype)`

### Parameters
- **target_dtype** (numpy.dtype or torch.dtype) - Required - The target data type for the view.
```

--------------------------------

### Import FFCV Loader and Decoders

Source: https://docs.ffcv.io/_sources/making_dataloaders.rst.txt

Required imports for initializing the FFCV loader and specifying field decoders.

```python
from ffcv.loader import Loader, OrderOption
from ffcv.fields.decoders import NDArrayDecoder, FloatDecoder
```

--------------------------------

### SimpleRGBImageDecoder Implementation

Source: https://docs.ffcv.io/_modules/ffcv/fields/rgb_image.html

Decoder for constant-resolution images, supporting JPEG decompression and raw memory copying.

```python
class SimpleRGBImageDecoder(Operation):
    """Most basic decoder for the :class:`~ffcv.fields.RGBImageField`.

    It only supports dataset with constant image resolution and will simply read (potentially decompress) and pass the images as is.
    """
    def __init__(self):
        super().__init__()


    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, AllocationQuery]:
        widths = self.metadata['width']
        heights = self.metadata['height']
        max_width = widths.max()
        max_height = heights.max()
        min_height = heights.min()
        min_width = widths.min()
        if min_width != max_width or max_height != min_height:
            msg = """SimpleRGBImageDecoder only supports constant image,
consider RandomResizedCropRGBImageDecoder or CenterCropRGBImageDecoder
instead."""
            raise TypeError(msg)

        biggest_shape = (max_height, max_width, 3)
        my_dtype = np.dtype('<u1')

        return (
            replace(previous_state, jit_mode=True,
                    shape=biggest_shape, dtype=my_dtype),
            AllocationQuery(biggest_shape, my_dtype)
        )


    def generate_code(self) -> Callable:
        mem_read = self.memory_read
        imdecode_c = Compiler.compile(imdecode)

        jpg = IMAGE_MODES['jpg']
        raw = IMAGE_MODES['raw']
        my_range = Compiler.get_iterator()
        my_memcpy = Compiler.compile(memcpy)

        def decode(batch_indices, destination, metadata, storage_state):
            for dst_ix in my_range(len(batch_indices)):
                source_ix = batch_indices[dst_ix]
                field = metadata[source_ix]
                image_data = mem_read(field['data_ptr'], storage_state)
                height, width = field['height'], field['width']

                if field['mode'] == jpg:
                    imdecode_c(image_data, destination[dst_ix],
                               height, width, height, width, 0, 0, 1, 1, False, False)
                else:
                    my_memcpy(image_data, destination[dst_ix])

            return destination[:len(batch_indices)]

        decode.is_parallel = True
        return decode
```

--------------------------------

### Basic Custom Transform Implementation

Source: https://docs.ffcv.io/_sources/ffcv_examples/custom_transforms.rst.txt

A basic implementation of a custom transform function that crops images to a corner.

```python
def generate_code(self) -> Callable:
    def pick_a_corner(images, dst):
        which_corner = np.random.randint(low=0, high=4, size=(images.shape[0]))
        for i in range(images.shape[0]):
            if which_corner[i] == 0:
                dst[i] = images[i,:images.shape[1]//2, :images.shape[2]//2]
            else:
                dst[i] = images[i,-images.shape[1]//2:,
                -images.shape[2]//2:]

        return dst
    return pick_a_corner
```

--------------------------------

### Implement Cutout Operation

Source: https://docs.ffcv.io/_modules/ffcv/transforms/cutout.html

Defines the Cutout class inheriting from Operation, providing parameters for crop size and fill color, and implementing the JIT-compiled transformation logic.

```python
"""
Cutout augmentation (https://arxiv.org/abs/1708.04552)
"""
import numpy as np
from typing import Callable, Optional, Tuple
from dataclasses import replace

from ffcv.pipeline.compiler import Compiler
from ..pipeline.allocation_query import AllocationQuery
from ..pipeline.operation import Operation
from ..pipeline.state import State


[docs]class Cutout(Operation):
    """Cutout data augmentation (https://arxiv.org/abs/1708.04552).

    Parameters
    ----------
    crop_size : int
        Size of the random square to cut out.
    fill : Tuple[int, int, int], optional
        An RGB color ((0, 0, 0) by default) to fill the cutout square with.
        Useful for when a normalization layer follows cutout, in which case
        you can set the fill such that the square is zero
        post-normalization.
    """
    def __init__(self, crop_size: int, fill: Tuple[int, int, int] = (0, 0, 0)):
        super().__init__()
        self.crop_size = crop_size
        self.fill = np.array(fill)


[docs]    def generate_code(self) -> Callable:
        my_range = Compiler.get_iterator()
        crop_size = self.crop_size
        fill = self.fill

        def cutout_square(images, *_):
            for i in my_range(images.shape[0]):
                # Generate random origin
                coord = (
                    np.random.randint(images.shape[1] - crop_size + 1),
                    np.random.randint(images.shape[2] - crop_size + 1),
                )
                # Black out image in-place
                images[i, coord[0]:coord[0] + crop_size, coord[1]:coord[1] + crop_size] = fill

            return images
        cutout_square.is_parallel = True

        return cutout_square


[docs]    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:
        return replace(previous_state, jit_mode=True), None
```

--------------------------------

### ToDevice

Source: https://docs.ffcv.io/_modules/ffcv/transforms/ops.html

Moves a tensor to a specified device (e.g., CPU or GPU).

```APIDOC
## ToDevice

### Description
Moves a tensor to a specified device. Supports asynchronous copying when moving from CPU to GPU.

### Constructor
`ToDevice(device, non_blocking=True)`

### Parameters
- **device** (torch.device) - Required - The target device to move the tensor to.
- **non_blocking** (bool) - Optional - Whether to perform an asynchronous copy (default: True).
```

--------------------------------

### Collect Webdataset shards

Source: https://docs.ffcv.io/_sources/writing_datasets.rst.txt

Use glob to retrieve a list of shard files from the specified directory.

```python
from glob import glob
from os import path

my_shards = glob(path.join(FOLDER, '*'))
```

--------------------------------

### Implement a basic custom transform

Source: https://docs.ffcv.io/ffcv_examples/custom_transforms.html

Defines a transformation function using generate_code that processes image batches without explicit performance optimizations.

```python
import numpy as np

def generate_code(self) -> Callable:
    def pick_a_corner(images, dst):
        which_corner = np.random.randint(low=0, high=4, size=(images.shape[0]))
        for i in range(images.shape[0]):
            if which_corner[i] == 0:
                dst[i] = images[i,:images.shape[1]//2, :images.shape[2]//2]
            else:
                dst[i] = images[i,-images.shape[1]//2:,
                -images.shape[2]//2:]

        return dst
    return pick_a_corner
```

--------------------------------

### Train and Evaluate Model

Source: https://docs.ffcv.io/_sources/ffcv_examples/cifar10.rst.txt

Executes the training loop with mixed precision and evaluates accuracy using test-time augmentation.

```python
from tqdm import tqdm

for ep in range(EPOCHS):
    for ims, labs in tqdm(loaders['train']):
        opt.zero_grad(set_to_none=True)
        with autocast():
            out = model(ims)
            loss = loss_fn(out, labs)

        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()
        scheduler.step()

model.eval()
with ch.no_grad():
    total_correct, total_num = 0., 0.
    for ims, labs in tqdm(loaders['test']):
        with autocast():
            out = (model(ims) + model(ch.fliplr(ims))) / 2. # Test-time augmentation
            total_correct += out.argmax(1).eq(labs).sum().cpu().item()
            total_num += ims.shape[0]

    print(f'Accuracy: {total_correct / total_num * 100:.1f}%')
```

--------------------------------

### ToDevice Operation

Source: https://docs.ffcv.io/_modules/ffcv/transforms/ops.html

Moves a tensor to a specified device, supporting asynchronous copying.

```python
[docs]class ToDevice(Operation):
    """Move tensor to device.

    Parameters
    ----------
    device: torch.device
        Device to move to.
    non_blocking: bool
        Asynchronous if copying from CPU to GPU.
    """
    def __init__(self, device, non_blocking=True):
        super().__init__()
        self.device = device
        # assert isinstance(device, ch.device), \
        #    f'Make sure device is a ch.device (not a {type(device)})'
        self.non_blocking = non_blocking

[docs]    def generate_code(self) -> Callable:
        def to_device(inp, dst):
            if len(inp.shape) == 4:
                if inp.is_contiguous(memory_format=ch.channels_last):
                    dst = dst.reshape(inp.shape[0], inp.shape[2], inp.shape[3], inp.shape[1])
                    dst = dst.permute(0, 3, 1, 2)
            dst = dst[:inp.shape[0]]
            dst.copy_(inp, non_blocking=self.non_blocking)
            return dst

        return to_device

[docs]    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:
        return replace(previous_state, device=self.device), AllocationQuery(previous_state.shape, dtype=previous_state.dtype, device=self.device)
```

--------------------------------

### Linear Regression Training Template

Source: https://docs.ffcv.io/_sources/ffcv_examples/linear_regression.rst.txt

Implements SGD for linear regression, requiring a pre-defined train_loader.

```python
import torch as ch
from tqdm import tqdm
import time

train_loader = None # TODO!

# Calculate data mean and variance for normalization
def calculate_stats(loader, N):
    mean, stdev = 0., 0.
    for x_batch, _ in tqdm(loader):
        mean += x_batch.sum(0) / N
        stdev += x_batch.pow(2).sum(0) / N
    return mean, ch.sqrt(stdev - mean.pow(2))

start_time = time.time()
mean, stdev = calculate_stats(train_loader, N)
mean, stdev = mean.cuda(), stdev.cuda()
w_est, b_est = ch.zeros(D).cuda(), ch.zeros(1).cuda() # Initial guess for W
num_epochs = 10 # Number of full passes over the data to do

lr = 5e-2
for _ in range(num_epochs):
    total_loss, num_examples = 0., 0.
    start_time = time.time()
    for (x_batch, y_batch) in tqdm(train_loader):
        x_batch = x_batch.cuda()
        y_batch = y_batch.cuda()
        # Normalize the data for stability
        x_batch = (x_batch - mean) / stdev
        residual = x_batch @ w_est + b_est - y_batch
        # Gradients
        w_grad = x_batch.T @ residual / x_batch.shape[0]
        b_grad = ch.mean(residual, dim=0)
        # import ipdb; ipdb.set_trace()
        w_est = w_est - lr * w_grad
        b_est = b_est - lr * b_grad
        total_loss += residual.pow(2).sum()
        num_examples += x_batch.shape[0]
    print('Epoch time:', time.time() - start_time)
    print(f'Average loss: {total_loss / num_examples:.3f} | ',
        f'Norm diff', ch.norm(w_est / stdev - ch.tensor(W).cuda()).cpu().item())

print(f'Total script running time: {time.time() - start_time:.2f}s')
```

--------------------------------

### Implement IntDecoder and FloatDecoder

Source: https://docs.ffcv.io/_modules/ffcv/fields/basics.html

Concrete decoder implementations for 64-bit integers and 64-bit floating-point scalars.

```python
class IntDecoder(BasicDecoder):
    """Decoder for signed integers scalars (int64)
    """
    dtype = np.dtype('<i8')

class FloatDecoder(BasicDecoder):
    """Decoder for floating point scalars (float64)
    """
    dtype = np.dtype('<f8')
```

--------------------------------

### class Poison(mask, alpha, indices, clamp=(0, 255))

Source: https://docs.ffcv.io/_modules/ffcv/transforms/poisoning.html

Initializes the Poison transformation to apply a mask to specified image indices.

```APIDOC
## class Poison

### Description
Poison specified images by adding a mask with given opacity. Operates on raw arrays (not tensors).

### Parameters
- **mask** (ndarray) - Required - The mask to apply to each image.
- **alpha** (float) - Required - The opacity of the mask.
- **indices** (Sequence[int]) - Required - The indices of images that should have the mask applied.
- **clamp** (Tuple[int, int]) - Optional - Clamps the final pixel values between these two values (default: (0, 255)).
```

--------------------------------

### Define Data Pipelines

Source: https://docs.ffcv.io/_sources/making_dataloaders.rst.txt

Specify how to read fields and apply transformations using a dictionary of decoders and transforms.

```python
from ffcv.transforms import ToTensor

PIPELINES = {
  'covariate': [NDArrayDecoder(), ToTensor()],
  'label': [FloatDecoder(), ToTensor()]
}
```

--------------------------------

### Define ResNet-9 architecture components

Source: https://docs.ffcv.io/_sources/ffcv_examples/cifar10.rst.txt

Helper modules for the ResNet-9 architecture, including custom Mul, Flatten, and Residual layers.

```python
class Mul(ch.nn.Module):
    def __init__(self, weight):
        super(Mul, self).__init__()
        self.weight = weight
    def forward(self, x): return x * self.weight

class Flatten(ch.nn.Module):
    def forward(self, x): return x.view(x.size(0), -1)

class Residual(ch.nn.Module):
    def __init__(self, module):
        super(Residual, self).__init__()
        self.module = module
    def forward(self, x): return x + self.module(x)
```

--------------------------------

### RandomResizedCrop Implementation

Source: https://docs.ffcv.io/_modules/ffcv/transforms/random_resized_crop.html

The class definition for RandomResizedCrop, including initialization, code generation for the pipeline, and state declaration.

```python
"""
Random resized crop, similar to torchvision.transforms.RandomResizedCrop
"""
from dataclasses import replace
from .utils import fast_crop
import numpy as np
from typing import Callable, Optional, Tuple
from ..pipeline.allocation_query import AllocationQuery
from ..pipeline.operation import Operation
from ..pipeline.state import State
from ..pipeline.compiler import Compiler


[docs]class RandomResizedCrop(Operation):
    """Crop a random portion of image with random aspect ratio and resize it to
    a given size. Chances are you do not want to use this augmentation and
    instead want to include RRC as part of the decoder, by using the 
    :cla:`~ffcv.fields.rgb_image.ResizedCropRGBImageDecoder` class.

    Parameters
    ----------
    scale : Tuple[float, float]
        Lower and upper bounds for the ratio of random area of the crop.
    ratio : Tuple[float, float]
        Lower and upper bounds for random aspect ratio of the crop.
    size : int
        Side length of the output.
    """
    def __init__(self, scale: Tuple[float, float], ratio: Tuple[float, float], size: int):
        super().__init__()
        self.scale = scale
        self.ratio = ratio
        self.size = size
        

[docs]    def generate_code(self) -> Callable:
        scale, ratio = self.scale, self.ratio
        if isinstance(scale, tuple):
            scale = np.array(scale)
        if isinstance(ratio, tuple):
            ratio = np.array(ratio)
        my_range = Compiler.get_iterator()
        def random_resized_crop(images, dst):
            for idx in my_range(images.shape[0]):
                i, j, h, w = fast_crop.get_random_crop(images[idx].shape[0],
                                                    images[idx].shape[1],
                                                    scale,
                                                    ratio)
                fast_crop.resize_crop(images[idx], i, i + h, j, j + w, dst[idx])
            return dst

        random_resized_crop.is_parallel = True
        return random_resized_crop


[docs]    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:
        return replace(previous_state, jit_mode=True, shape=(self.size, self.size, 3)), \
               AllocationQuery((self.size, self.size, 3), dtype=previous_state.dtype)
```

--------------------------------

### View Operation

Source: https://docs.ffcv.io/_modules/ffcv/transforms/ops.html

Provides a view of an array using NumPy or PyTorch view methods.

```python
[docs]class View(Operation):
    """View array using np.view or torch.view.

    Parameters
    ----------
    target_dtype: numpy.dtype or torch.dtype
        Target data type.
    """
    def __init__(self, target_dtype):
        super().__init__()
        self.target_dtype = target_dtype
```

--------------------------------

### Image Resizing Utility

Source: https://docs.ffcv.io/_modules/ffcv/fields/rgb_image.html

Resizes an image based on a target resolution while maintaining aspect ratio.

```python
def resizer(image, target_resolution):
    if target_resolution is None:
        return image
    original_size = np.array([image.shape[1], image.shape[0]])
    ratio = target_resolution / original_size.max()
    if ratio < 1:
        new_size = (ratio * original_size).astype(int)
        image = cv2.resize(image, tuple(new_size), interpolation=cv2.INTER_AREA)
    return image
```

--------------------------------

### RandomHorizontalFlip

Source: https://docs.ffcv.io/_modules/ffcv/transforms/flip.html

Initializes the RandomHorizontalFlip transformation to be used within an FFCV pipeline.

```APIDOC
## class RandomHorizontalFlip(flip_prob=0.5)

### Description
Flip the image horizontally with a given probability. This operation works on raw arrays rather than tensors.

### Parameters
- **flip_prob** (float) - Optional - The probability with which to flip each image in the batch horizontally. Defaults to 0.5.
```

--------------------------------

### Implement FloatField and IntField

Source: https://docs.ffcv.io/_modules/ffcv/fields/basics.html

Field classes for scalar data, defining metadata types, binary serialization, and associated decoder classes.

```python
class FloatField(Field):
    """
    A subclass of :class:`~ffcv.fields.Field` supporting (scalar) floating-point (float64)
    values.
    """
    def __init__(self):
        pass

    @property
    def metadata_type(self) -> np.dtype:
        return np.dtype('<f8')

    @staticmethod
    def from_binary(binary: ARG_TYPE) -> Field:
        return FloatField()

    def to_binary(self) -> ARG_TYPE:
        return np.zeros(1, dtype=ARG_TYPE)[0]

    def encode(self, destination, field, malloc):
        destination[0] = field

    def get_decoder_class(self) -> Type[Operation]:
        return FloatDecoder

class IntField(Field):
    """
    A subclass of :class:`~ffcv.fields.Field` supporting (scalar) integer
    values.
    """
    @property
    def metadata_type(self) -> np.dtype:
        return np.dtype('<i8')

    @staticmethod
    def from_binary(binary: ARG_TYPE) -> Field:
        return IntField()

    def to_binary(self) -> ARG_TYPE:
        return np.zeros(1, dtype=ARG_TYPE)[0]

    def encode(self, destination, field, malloc):
        # We just allocate 1024bytes for fun
        destination[0] = field

    def get_decoder_class(self) -> Type[Operation]:
        return IntDecoder
```

--------------------------------

### Define Webdataset decoding pipeline

Source: https://docs.ffcv.io/_sources/writing_datasets.rst.txt

Create a function to decode and format data from Webdataset shards for FFCV compatibility.

```python
def pipeline(dataset):
    return dataset.decode('rgb8').to_tuple('jpg:png;jpeg cls')
```

--------------------------------

### FFCV Transformation Operation Methods

Source: https://docs.ffcv.io/_modules/ffcv/transforms/ops.html

Implementation of the generate_code and declare_state_and_memory methods for custom FFCV transformation operations.

```python
    def generate_code(self) -> Callable:
        def convert(inp, dst):
            return inp.view(self.target_dtype)

        convert.is_parallel = True

        return convert



    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:
        return replace(previous_state, dtype=self.target_dtype, jit_mode=False), None
```

--------------------------------

### ResizedCropRGBImageDecoder Implementation

Source: https://docs.ffcv.io/_modules/ffcv/fields/rgb_image.html

Base class for RGB image decoders that handle resizing and cropping operations.

```python
    def __init__(self, output_size):
        super().__init__()
        self.output_size = output_size

    def declare_state_and_memory(self, previous_state: State) -> Tuple[State, AllocationQuery]:
        widths = self.metadata['width']
        heights = self.metadata['height']
        # We convert to uint64 to avoid overflows
        self.max_width = np.uint64(widths.max())
        self.max_height = np.uint64(heights.max())
        output_shape = (self.output_size[0], self.output_size[1], 3)
        my_dtype = np.dtype('<u1')

        return (
            replace(previous_state, jit_mode=True,
                    shape=output_shape, dtype=my_dtype),
            (AllocationQuery(output_shape, my_dtype),
            AllocationQuery((self.max_height * self.max_width * np.uint64(3),), my_dtype),
            )
        )

    def generate_code(self) -> Callable:

        jpg = IMAGE_MODES['jpg']

        mem_read = self.memory_read
        my_range = Compiler.get_iterator()
        imdecode_c = Compiler.compile(imdecode)
        resize_crop_c = Compiler.compile(resize_crop)
        get_crop_c = Compiler.compile(self.get_crop_generator)

        scale = self.scale
        ratio = self.ratio
        if isinstance(scale, tuple):
            scale = np.array(scale)
        if isinstance(ratio, tuple):
            ratio = np.array(ratio)

        def decode(batch_indices, my_storage, metadata, storage_state):
            destination, temp_storage = my_storage
            for dst_ix in my_range(len(batch_indices)):
                source_ix = batch_indices[dst_ix]
                field = metadata[source_ix]
                image_data = mem_read(field['data_ptr'], storage_state)
                height = np.uint32(field['height'])
                width = np.uint32(field['width'])

                if field['mode'] == jpg:
                    temp_buffer = temp_storage[dst_ix]
                    imdecode_c(image_data, temp_buffer,
                               height, width, height, width, 0, 0, 1, 1, False, False)
                    selected_size = 3 * height * width
                    temp_buffer = temp_buffer.reshape(-1)[:selected_size]
                    temp_buffer = temp_buffer.reshape(height, width, 3)

                else:
                    temp_buffer = image_data.reshape(height, width, 3)

                i, j, h, w = get_crop_c(height, width, scale, ratio)

                resize_crop_c(temp_buffer, i, i + h, j, j + w,
                              destination[dst_ix])

            return destination[:len(batch_indices)]
        decode.is_parallel = True
        return decode

    @property
    @abstractmethod
    def get_crop_generator():
        raise NotImplementedError
```
Tokens:
Raw