API Reference

This section contains the complete API documentation for DeepAugment, automatically generated from the source code.

---

Main API

Core Functions

deepaugment.optimize(X_train, y_train, X_val=None, y_val=None, iterations=None, epochs=None, verbose=True, **kwargs)

Quick optimization function for minimal code.

Convention: if no val set, auto-split 80/20.

Parameters:

• X_train – Training images

• y_train – Training labels

• X_val – Validation images (optional, will split if None)

• y_val – Validation labels (optional)

• iterations – Optimization iterations (default from config)

• epochs – Training epochs per policy (default from config)

• verbose – Show progress

• **kwargs – Additional args for DeepAugment

Returns:

Best augmentation policy

Example

>>> from deepaugment import optimize
>>> best = optimize(X, y, iterations=50)

DeepAugment Class

class deepaugment.DeepAugment

Bases: object

Find optimal image augmentation policies using Bayesian optimization.

Beautiful API inspired by Rails: Convention over Configuration, but with sharp knives when you need them.

Example

>>> aug = DeepAugment(X_train, y_train, X_val, y_val)
>>> best = aug.optimize(iterations=50)
>>> aug.show_best()

__init__(X_train, y_train, X_val, y_val, model='simple', device='auto', random_state=None, experiment_name=None, method='bayesian', save_history=None, transform_categories=None, custom_reward_fn=None, resume_from=None, n_operations=None, train_size=None, val_size=None)

Initialize DeepAugment with beautiful defaults and configuration freedom.

All args are optional - sensible defaults from config.

Parameters:

• X_train – Training data (N, H, W, C) and labels (N,)

• y_train – Training data (N, H, W, C) and labels (N,)

• X_val – Validation data and labels

• y_val – Validation data and labels

• 1 (# PHASE) – Essential

• model – Model architecture (default: ‘simple’)

• device – ‘cuda’, ‘mps’, ‘cpu’, or ‘auto’ (default: ‘auto’)

• random_state – Seed for reproducibility (default from config)

• experiment_name – Name for tracking (default: timestamp)

• 2 (# PHASE) – Useful

• method – ‘bayesian’ or ‘random’ search (default: ‘bayesian’)

• save_history – Save optimization history (default from config)

• 3 (# PHASE) – Advanced

• transform_categories – Restrict transforms, e.g., [‘geometric’, ‘color’]

• custom_reward_fn – Custom reward function(entry) -> float

• resume_from – Path to checkpoint to continue from

• Core (#)

• n_operations – Transforms per policy (default from config)

• train_size – Training subset size (default from config)

• val_size – Validation subset size (default from config)

optimize(iterations=None, epochs=None, samples=None, batch_size=None, learning_rate=None, early_stopping=None, patience=None, verbose=True)

Search for best augmentation policy.

All args optional - sensible defaults from config.

Parameters:

• iterations – Number of policies to try (default from config)

• epochs – Training epochs per evaluation (default from config)

• samples – Training runs per policy, for averaging (default from config)

• batch_size – Training batch size (default from config)

• learning_rate – Learning rate (default from config)

• early_stopping – Stop if no improvement (default from config)

• patience – Early stopping patience (default from config)

• verbose – Show progress bar

Returns:

Best policy as list of (transform, magnitude) tuples

best_policy()

Get best policy found (human-readable).

best_score()

Get best validation accuracy found.

show_best(n=5)

Print top N policies. Beautiful output for programmer happiness.

---

Transforms

Image transformations - elegant, minimal, powerful.

Uses torchvision v2 transforms. Convention: magnitude in [0, 1].

deepaugment.transforms.make_transform(name, magnitude)

Create single transform. Fail fast if invalid.

Convention: magnitude clamped to [0, 1].

deepaugment.transforms.apply_policy(image, policy)

Apply augmentation policy to image.

Pure function: image → policy → augmented image.

Parameters:

• image – RGB image (H, W, C) numpy array or PIL Image

• policy – List of (transform_name, magnitude) tuples

Returns:

Augmented image as numpy array

deepaugment.transforms.create_augmenter(policy)

Create reusable augmenter from policy.

Returns: Callable that augments images.

deepaugment.transforms.get_transform_names(categories=None)

Get transform names, optionally filtered by category.

Convention: None means all transforms.

Transform Constants

deepaugment.TRANSFORMS

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object’s

(key, value) pairs

dict(iterable) -> new dictionary initialized as if via:

d = {} for k, v in iterable:

d[k] = v

dict(**kwargs) -> new dictionary initialized with the name=value pairs

in the keyword argument list. For example: dict(one=1, two=2)

---

Models

Model architectures - minimal, elegant, effective.

Just the models. Training logic lives in trainer.py (separation of concerns).

class deepaugment.models.SimpleCNN

Bases: Module

Minimal CNN for fast policy evaluation.

Compact but effective: ~1.2M parameters for 32x32 images.

__init__(num_classes=10, in_channels=3)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Forward pass. Clean and simple.

deepaugment.models.create_model(model_name='simple', num_classes=10, in_channels=3)

Create model by name. Convention: ‘simple’ is default.

Extensible: add more models here as needed.

Model Creation

deepaugment.create_model(model_name='simple', num_classes=10, in_channels=3)

Create model by name. Convention: ‘simple’ is default.

Extensible: add more models here as needed.

SimpleCNN Model

class deepaugment.SimpleCNN

Bases: Module

Minimal CNN for fast policy evaluation.

Compact but effective: ~1.2M parameters for 32x32 images.

__init__(num_classes=10, in_channels=3)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Forward pass. Clean and simple.

---

Policy

Policy operations - elegant, composable, functional.

A policy is a list of (transform_name, magnitude) tuples. Everything related to policy representation lives here (SSOT).

class deepaugment.policy.PolicySpace

Bases: object

Policy search space. Knows everything about what policies can be.

Follows SSOT: all policy space logic lives here.

n_operations: int

transform_categories: list

random_state: int

transform_names: list

n_transforms: int

n_dimensions: int

__attrs_post_init__()

Initialize derived attributes. Convention: computed, not stored.

random_policy()

Sample random policy from space.

format_policy(raw_policy)

Convert raw policy (optimizer format) to human-readable.

Raw: [idx1, mag1, idx2, mag2, …] Human: [(name1, mag1), (name2, mag2), …]

dimensions()

Get search space dimensions for optimizer.

Returns list of (min, max) or categorical values for each dimension.

__init__(n_operations=4, transform_categories=None, random_state=42)

Method generated by attrs for class PolicySpace.

Parameters:

• n_operations (int)

• transform_categories (list)

• random_state (int)

Return type:

None

class deepaugment.policy.PolicyHistory

Bases: object

Track evaluated policies and scores.

Minimal, elegant storage. Unix: do one thing well.

policies: list

scores: list

add(policy, score)

Add evaluation result.

best()

Get best policy and score. Simple.

top_k(k=5)

Get top K policies.

to_dict()

Export as dict for serialization.

classmethod from_dict(data)

Load from dict.

__init__(policies=NOTHING, scores=NOTHING)

Method generated by attrs for class PolicyHistory.

Parameters:

• policies (list)

• scores (list)

Return type:

None

deepaugment.policy.display_policy(policy, indent=4)

Pretty print policy. For programmer happiness.

deepaugment.policy.policy_summary(policy)

One-line policy summary.

PolicySpace

class deepaugment.PolicySpace

Bases: object

Policy search space. Knows everything about what policies can be.

Follows SSOT: all policy space logic lives here.

n_operations: int

transform_categories: list

random_state: int

transform_names: list

n_transforms: int

n_dimensions: int

__attrs_post_init__()

Initialize derived attributes. Convention: computed, not stored.

random_policy()

Sample random policy from space.

format_policy(raw_policy)

Convert raw policy (optimizer format) to human-readable.

Raw: [idx1, mag1, idx2, mag2, …] Human: [(name1, mag1), (name2, mag2), …]

dimensions()

Get search space dimensions for optimizer.

Returns list of (min, max) or categorical values for each dimension.

__init__(n_operations=4, transform_categories=None, random_state=42)

Method generated by attrs for class PolicySpace.

Parameters:

• n_operations (int)

• transform_categories (list)

• random_state (int)

Return type:

None

PolicyHistory

class deepaugment.policy.PolicyHistory

Bases: object

Track evaluated policies and scores.

Minimal, elegant storage. Unix: do one thing well.

policies: list

scores: list

add(policy, score)

Add evaluation result.

best()

Get best policy and score. Simple.

top_k(k=5)

Get top K policies.

to_dict()

Export as dict for serialization.

classmethod from_dict(data)

Load from dict.

__init__(policies=NOTHING, scores=NOTHING)

Method generated by attrs for class PolicyHistory.

Parameters:

• policies (list)

• scores (list)

Return type:

None

---

Search

Search strategies - Bayesian optimization and random search.

Minimal, composable, elegant. Strategy pattern without the boilerplate.

class deepaugment.search.SearchStrategy

Bases: object

Base for search strategies. Minimal interface.

ask() → get next policy to try tell() → report result

policy_space: PolicySpace

history: PolicyHistory

__attrs_post_init__()

Initialize history if not provided.

ask()

Get next policy to try. Subclasses implement.

tell(policy, score)

Report evaluation result. Updates history.

best()

Get best policy found so far.

__init__(policy_space, history=None)

Method generated by attrs for class SearchStrategy.

Parameters:

• policy_space (PolicySpace)

• history (PolicyHistory)

Return type:

None

class deepaugment.search.RandomSearch

Bases: SearchStrategy

Random search baseline. Simple, fast, surprisingly effective.

No learning, just exploration. Perfect for sanity checks.

ask()

Sample random policy.

__init__(policy_space, history=None)

Method generated by attrs for class RandomSearch.

Parameters:

• policy_space (PolicySpace)

• history (PolicyHistory)

Return type:

None

class deepaugment.search.BayesianSearch

Bases: SearchStrategy

Bayesian optimization using scikit-optimize.

Learn from past evaluations to suggest better policies. Uses Random Forest + Expected Improvement.

n_initial: int

random_state: int

__attrs_post_init__()

Initialize Bayesian optimizer.

ask()

Ask optimizer for next policy.

Convention: random initially, smart after.

tell(policy, score)

Report result to optimizer.

scikit-optimize minimizes, so we negate the score.

load_history(history_data)

Resume from saved history. Progress over stability!

__init__(policy_space, history=None, n_initial=None, random_state=None, optimizer=None)

Method generated by attrs for class BayesianSearch.

Parameters:

• policy_space (PolicySpace)

• history (PolicyHistory)

• n_initial (int)

• random_state (int)

• optimizer (Optimizer)

Return type:

None

deepaugment.search.create_search(method='bayesian', policy_space=None, n_initial=None, random_state=None)

Create search strategy by name.

Convention: ‘bayesian’ is default, ‘random’ is baseline.

---

Training

Training logic - functional, composable, minimal.

Pure functions for training models. No classes, just functions. Unix philosophy: do one thing well.

deepaugment.trainer.prepare_data(X, y)

Convert numpy to torch tensors.

Convention: assumes (N, H, W, C) format, converts to (N, C, H, W).

deepaugment.trainer.create_loaders(train_data, val_data, batch_size=None)

Create data loaders. Convention: sensible batch size from config.

deepaugment.trainer.train_epoch(model, loader, optimizer, criterion, device)

Train for one epoch. Pure function, returns metrics.

deepaugment.trainer.validate_epoch(model, loader, criterion, device)

Validate for one epoch. Pure function, returns metrics.

deepaugment.trainer.train_model(model, train_data, val_data, epochs=None, batch_size=None, learning_rate=None, device=None)

Train model to completion. Functional composition.

Parameters:

• model – PyTorch model

• train_data – (X_train, y_train) as numpy arrays

• val_data – (X_val, y_val) as numpy arrays

• epochs – Training epochs (default from config)

• batch_size – Batch size (default from config)

• learning_rate – Learning rate (default from config)

• device – Device string (default from config)

Returns:

Training history dict

deepaugment.trainer.evaluate_policy(policy, train_data, val_data, num_classes, augmenter, model_factory, epochs=None, samples=None, **kwargs)

Evaluate augmentation policy by training model.

Functional approach: augment → train → evaluate → score.

Parameters:

• policy – List of (transform, magnitude) tuples

• train_data – (X_train, y_train)

• val_data – (X_val, y_val)

• num_classes – Number of classes

• augmenter – Function to augment images

• model_factory – Function that creates model

• epochs – Training epochs

• samples – Number of training runs (for averaging)

• **kwargs – Additional args for train_model

Returns:

Average validation accuracy

---

Configuration

Single Source of Truth for all configuration.

Following Rails doctrine: Convention over Configuration. All magic numbers live here. Change once, apply everywhere.

class deepaugment.config.Defaults

Bases: object

Beautiful defaults that just work.

epochs: int

batch_size: int

learning_rate: float

iterations: int

samples: int

n_operations: int

n_initial_points: int

train_size: int

val_size: int

image_size: int

model_name: str

dropout_rate: float

device: str

method: str

random_state: int

save_history: bool

checkpoint_every: int

early_stopping: bool

patience: int

__init__(epochs=10, batch_size=64, learning_rate=0.001, iterations=50, samples=1, n_operations=4, n_initial_points=10, train_size=2000, val_size=500, image_size=32, model_name='simple', dropout_rate=0.3, device='auto', method='bayesian', random_state=42, save_history=True, checkpoint_every=10, early_stopping=False, patience=10)

Method generated by attrs for class Defaults.

Parameters:

• epochs (int)

• batch_size (int)

• learning_rate (float)

• iterations (int)

• samples (int)

• n_operations (int)

• n_initial_points (int)

• train_size (int)

• val_size (int)

• image_size (int)

• model_name (str)

• dropout_rate (float)

• device (str)

• method (str)

• random_state (int)

• save_history (bool)

• checkpoint_every (int)

• early_stopping (bool)

• patience (int)

Return type:

None

deepaugment.config.auto_device()

Auto-detect best available device. CUDA > MPS > CPU.

deepaugment.config.resolve_device(device)

Resolve device string. Convention: ‘auto’ means smart detection.

---

Utilities

Elegant utilities inspired by Unix philosophy and functional programming.

Each function does one thing well. Compose them for power.

deepaugment.utils.numpy_to_pil(image)

numpy → PIL. Pure function, no side effects.

deepaugment.utils.pil_to_numpy(image)

PIL → numpy. Pure function, no side effects.

deepaugment.utils.sample_indices(total='__no__default__', size='__no__default__', seed=42)

Sample indices without replacement. Functional and reproducible.

deepaugment.utils.split_data(X, y, ratio=0.8, seed=42)

Split data into train/val. Convention: 80/20 split.

deepaugment.utils.to_native(obj)

Convert numpy types to native Python for JSON serialization.

Recursive, handles nested structures. Unix philosophy: do one thing well.

deepaugment.utils.save_checkpoint(data, filename, directory=None)

Save checkpoint as JSON. Convention: pretty print for human readability.

deepaugment.utils.load_checkpoint(filepath)

Load checkpoint from JSON. Simple, no magic.

deepaugment.utils.generate_experiment_name(prefix='exp')

Generate unique experiment name. Convention: ISO format timestamp.

deepaugment.utils.format_policy(policy, indent=4)

Format policy as beautiful string for display.

deepaugment.utils.format_policy_summary(policy)

One-line policy summary.

deepaugment.utils.validate_images(X)

Validate image array format. Fail fast with helpful error.

deepaugment.utils.validate_labels(y, X)

Validate labels match images. Fail fast.

deepaugment.utils.clamp(x, min_val=0.0, max_val=1.0)

Clamp value to range. Simple, pure.

deepaugment.utils.normalize(arr, low=0.0, high=1.0)

Normalize array to range [low, high].