DocumentationDocumentation

Serialization of Augmentation Pipelines

On this page

An augmentation pipeline is part of an experiment, not just helper code. The transforms you choose, their probabilities, target-processing rules, and seed settings define the data distribution that the model sees during training. Saving that pipeline gives you a durable artifact that can be reviewed in a pull request, attached to a training run, referenced from a model card, and loaded again for debugging or deployment.

Serialization stores the declared augmentation policy. It does not store the image, mask, bounding box, or keypoint values from one specific call, and it does not store the random choices sampled for one specific image. For per-sample debugging, use save_applied_params=True or ReplayCompose; those are covered later in this guide.

Core Functions

The main serialization APIs are:

Use A.save and A.load for artifact files. Use A.to_dict and A.from_dict when your experiment tracker, training framework, or configuration system already manages file storage.

When to Use Serialization

Use serialization when the augmentation policy must survive outside the Python file that created it:

  • Experiment tracking: Save the exact training policy next to the model checkpoint, metrics, dataset version, and code revision.
  • Deployment preprocessing: Load the same validation or preprocessing policy in training, evaluation, and serving code instead of reimplementing it by hand.
  • Target-aware review: Make bounding-box formats, keypoint settings, label fields, and extra synchronized targets visible in one artifact.
  • Model cards and audits: Reference the saved policy when you need to answer what data distribution the model was trained or evaluated on.

Without serialization, those details live in scattered Python code, default arguments, and memory. With serialization, the augmentation contract becomes a file that can be diffed, archived, loaded, and reviewed.

What Gets Serialized

A saved pipeline contains the parts needed to reconstruct the augmentation policy:

  • The composition object, such as Compose or ReplayCompose.
  • The ordered transform tree, including nested compositions.
  • Transform constructor arguments, such as sizes, limits, interpolation settings, and probability p.
  • Pipeline-level options such as seed, is_check_shapes, and the replay save key when relevant.
  • Target-processing configuration: bbox_params, keypoint_params, and additional_targets.
  • The Albumentations package version in the serialized output.

It does not contain your training data, model weights, dataloader code, environment lockfile, or one sample's random draw. Store those separately when you need a complete experiment record.

JSON vs YAML

Albumentations supports two human-readable formats:

  • JSON is the default. Use it as the canonical experiment artifact when configs are generated by code, checked by CI, compared in exact diffs, or consumed by other tools. JSON has no extra dependency beyond the Python standard library.
  • YAML is useful when humans author or review the policy directly. It supports comments and is easier to scan for long configs, but it requires PyYAML (pip install pyyaml) and can be less strict across tooling.

For teams, a practical rule is: keep the artifact that the training job actually used as JSON, and optionally keep a reviewed YAML source when that improves human review. The important part is that the committed or logged config is the one loaded by the run, not a hand-maintained description of it.

Basic Example

This example saves a training policy to JSON, loads it back, and applies it to an image. The loaded pipeline has the same structure and constructor parameters as the original pipeline.

from pathlib import Path

import albumentations as A
import numpy as np

transform = A.Compose(
    [
        A.Resize(height=256, width=256),
        A.Rotate(angle_range=(-45, 45), p=0.5),
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(p=0.3),
    ],
    seed=137,
)

artifact_path = Path("augmentation-policy.json")
A.save(transform, artifact_path)

loaded_transform = A.load(artifact_path)

rng = np.random.default_rng(137)
image = rng.integers(0, 256, size=(300, 300, 3), dtype=np.uint8)
result = loaded_transform(image=image)

print(result["image"].shape)
print(type(loaded_transform))

To use YAML instead, pass data_format="yaml":

A.save(transform, "augmentation-policy.yaml", data_format="yaml")
loaded_transform = A.load("augmentation-policy.yaml", data_format="yaml")

Saving Target-Aware Policies

For detection, pose, stereo, multi-modal, and segmentation pipelines, the transform list is not enough. The policy also needs the target-processing contract: box format, keypoint format, label fields, visibility filtering, and any extra named targets that must receive synchronized geometry.

A.save preserves those settings. This example saves a detection-style policy, loads it, and applies the loaded policy to every target it declares:

import albumentations as A
import numpy as np

train_transform = A.Compose(
    [
        A.Resize(height=512, width=512),
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(p=0.2),
    ],
    bbox_params=A.BboxParams(
        coord_format="pascal_voc",
        label_fields=["bbox_labels"],
        min_visibility=0.3,
    ),
    keypoint_params=A.KeypointParams(
        coord_format="xy",
        remove_invisible=False,
    ),
    additional_targets={
        "right_image": "image",
        "depth_mask": "mask",
    },
    seed=137,
)

A.save(train_transform, "detector-policy.json")
loaded_transform = A.load("detector-policy.json")

image = np.zeros((640, 640, 3), dtype=np.uint8)
right_image = np.zeros((640, 640, 3), dtype=np.uint8)
depth_mask = np.zeros((640, 640), dtype=np.uint8)

result = loaded_transform(
    image=image,
    right_image=right_image,
    depth_mask=depth_mask,
    bboxes=[[98, 120, 420, 510]],
    bbox_labels=["car"],
    keypoints=[[260, 300]],
)

print(result["image"].shape)
print(result["right_image"].shape)
print(result["depth_mask"].shape)
print(result["bboxes"])
print(result["bbox_labels"])
print(result["keypoints"])

This is what makes the file useful as an experiment artifact: a reviewer can see that bounding boxes are Pascal VOC, small boxes below the visibility threshold are filtered, keypoints outside the image are kept, and right_image receives the same geometric parameters as image.

If a keypoint policy uses semantic label swapping, include the label_mapping in A.KeypointParams. If a bounding-box or keypoint policy depends on label arrays, include those field names in label_fields; otherwise the loaded pipeline will not know which extra arrays must stay aligned with the filtered targets.

Experiment Artifacts and Review

A saved augmentation policy is most useful when it travels with the rest of the experiment metadata:

  • Commit or log the exact JSON file loaded by the training run.
  • Store it next to the model checkpoint, metrics, dataset version, code revision, and dependency lockfile.
  • Review augmentation changes like code changes. A probability change from 0.2 to 0.8, a new crop, or a changed bounding-box visibility threshold can change the training distribution as much as a model hyperparameter.
  • Reference the policy artifact from model cards and experiment reports. This lets another engineer answer "what data distribution was this model trained on?" without reverse-engineering the training script.
  • Keep validation and test-time augmentation policies as separate artifacts. They usually have different goals from the training policy.

The artifact should be the input to the run, not only a description written afterward. That keeps training, review, and later debugging tied to the same policy definition.

Custom Transform Caveats

Standard Albumentations transforms serialize cleanly because their constructor arguments are known and serializable. Custom transforms need the same discipline.

Prefer regular, importable custom transform classes with explicit constructor arguments:

import albumentations as A
from albumentations.core.transforms_interface import ImageOnlyTransform


class AddConstant(ImageOnlyTransform):
    def __init__(self, value: int, p: float = 1.0):
        super().__init__(p=p)
        self.value = value

    def apply(self, img, **params):
        return img + self.value


transform = A.Compose([AddConstant(value=10, p=1.0)], seed=137)
config = A.to_dict(transform)
restored = A.from_dict(config)

Avoid anonymous inline callables for durable experiment configs. A function object cannot be faithfully written into JSON or YAML. If you use A.Lambda, give it a stable name and provide the implementation when loading:

import albumentations as A


def invert_image(image, **kwargs):
    return 255 - image


invert = A.Lambda(image=invert_image, name="invert_image", p=1.0)
transform = A.Compose([invert], seed=137)

A.save(transform, "lambda-policy.json")

loaded_transform = A.load(
    "lambda-policy.json",
    nonserializable={"invert_image": invert},
)

The same pattern works with dictionaries:

config = A.to_dict(transform)
loaded_transform = A.from_dict(
    config,
    nonserializable={"invert_image": invert},
)

For custom policies, add a round-trip test: build the pipeline, serialize it, load it, apply it to a small synthetic sample, and assert that expected targets are present with the expected shapes. This catches missing constructor state, unstable imports, and nonserializable callables before a training run depends on them.

Policy Files vs Applied Parameters

Serialization answers "what policy did this run declare?" Sometimes you need a different answer: "what exactly happened to this sample?"

Use these tools for different jobs:

  • A.save / A.load: save and restore the declared policy.
  • save_applied_params=True: inspect the transforms and sampled parameters applied in one pipeline call.
  • ReplayCompose: capture the exact sampled decisions for one call and replay them on another compatible sample.

Example with save_applied_params=True:

import albumentations as A
import numpy as np

image = np.zeros((512, 512, 3), dtype=np.uint8)

transform = A.Compose(
    [
        A.RandomCrop(height=256, width=256),
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(p=0.5),
    ],
    save_applied_params=True,
    seed=137,
)

result = transform(image=image)
print(result["applied_transforms"])

Example with ReplayCompose:

import albumentations as A
import numpy as np

image = np.zeros((512, 512, 3), dtype=np.uint8)
another_image = np.zeros((512, 512, 3), dtype=np.uint8)

transform = A.ReplayCompose(
    [
        A.RandomCrop(height=256, width=256),
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(p=0.5),
    ],
    seed=137,
)

result = transform(image=image)
replayed = A.ReplayCompose.replay(result["replay"], image=another_image)

Use ReplayCompose when you need to apply the same sampled crop, flip, or color decision again. Use serialization when you need to save the policy that will sample future decisions.

Summary

  • Save augmentation policies with A.save and restore them with A.load.
  • Treat the saved config as an experiment artifact: commit it, diff it, log it, and reference it from reports or model cards.
  • Use JSON as the default canonical format; use YAML when human-authored review is worth the extra dependency.
  • Include target-processing settings by saving the whole Compose, not just the list of transforms.
  • Test custom transforms with a serialize-load-apply round trip.
  • Use save_applied_params=True or ReplayCompose for per-sample debugging; use serialization for the declared policy.

Where to Go Next

After understanding serialization:

  • Master reproducibility: Read the Reproducibility Guide for seed behavior, debugging techniques, and experiment practices.
  • Plan inference policies separately: Use the Test-Time Augmentation Guide when saving or reviewing inference-time augmentation policies.
  • Understand pipeline structure: Review Pipelines to see how Compose controls targets, probabilities, seeds, and target validation.
  • Handle extra inputs: Read Additional Targets when a pipeline must synchronize stereo pairs, multiple masks, or multi-modal arrays.
  • Write custom transforms carefully: Use the Creating Custom Transforms guide before placing custom code in a saved policy.