Automatic Augmentation Methods#

Augmentation Policy#

This module contains common data augmentation policies that can improve the accuracy of image classification models.

class kornia.augmentation.auto.AutoAugment(policy='imagenet', transformation_matrix_mode='silent')[source]#

Apply AutoAugment [CZM+18] searched strategies.

See the Convention block on PolicyAugmentBase.

Convention:

  • selects one sub-policy uniformly for each forward call. A sub-policy is an ordered sequence of (name, probability, magnitude_bin) entries, so every listed operation is considered in order.

  • a magnitude bin selects the adjacent interval of that operation’s eleven-point scale. Bins 0 through 9 are valid for operations with a magnitude; operations that ignore magnitude accept any value. The wrapped operation samples a value inside its selected interval for every batch row; posterize then rounds to the nearest integer number of bits (so bin 1, spanning 4.4 to 4.8, produces 4 or 5), translate_x / translate_y multiply it by the image size in pixels, and shear_x / shear_y, whose scale is a shear fraction from -0.3 to 0.3, multiply it by 180 to give degrees.

Parameters:

  • policy (Union[str, List[List[Tuple[str, Union[float, int], Union[float, int, None]]]]], optional) – a customized policy config or presets of “imagenet”, “cifar10”, and “svhn”. Default: "imagenet"

  • transformation_matrix_mode (str, optional) – computation mode for the chained transformation matrix, via .transform_matrix attribute. If silent, transformation matrix will be computed silently and the non-rigid modules will be ignored as identity transformations. If rigid, the result is the same for a policy: every operation wrapper contributes a matrix (an intensity operation the identity), so none is rejected. If skip, transformation matrix will be totally ignored. Default: "silent"

Examples

>>> import torch
>>> import kornia.augmentation as K
>>> from kornia.augmentation.auto import AutoAugment
>>> in_tensor = torch.rand(5, 3, 30, 30)
>>> aug = K.AugmentationSequential(AutoAugment())
>>> aug(in_tensor).shape
torch.Size([5, 3, 30, 30])
get_transformation_matrix(input, params=None, recompute=False, extra_args=None)[source]#

Compute the transformation matrix according to the provided parameters.

Parameters:

  • input (Tensor) – the input torch.Tensor.

  • params (Optional[List[ParamItem]], optional) – params for the sequence. Default: None

  • recompute (bool, optional) – if to recompute the transformation matrix according to the params. default: False. Default: False

  • extra_args (Optional[Dict[str, Any]], optional) – Optional dictionary of extra arguments with specific options for different input types. Default: None

Return type:

Optional[Tensor]

forward_parameters(batch_shape)[source]#

Generate per-module parameters for one policy forward pass.

Parameters:

batch_shape (Size) – Input shape used to sample operation parameters.

Return type:

List[ParamItem]

Returns:

Parameters for each selected child module, in execution order.

forward(input, params=None, extra_args=None)[source]#

Apply the policy to input and cache the parameters used.

Parameters:

  • input (Tensor) – Input tensor.

  • params (Optional[List[ParamItem]], optional) – Optional precomputed parameters. If omitted, parameters are sampled on-the-fly. Default: None

  • extra_args (Optional[Dict[str, Any]], optional) – Optional overrides forwarded to child transforms. Default: None

Return type:

Tensor

Returns:

Augmented tensor.

inverse(input, params=None, extra_args=None)[source]#

Undo the drawn sub-policy, or refuse when it cannot be undone.

Only geometric operations are invertible. A policy draw that contains an intensity operation is not round-trippable, and inverting it would return a tensor that still carries that operation – for a draw with no geometry at all, the input unchanged. This raises instead, as inverse on MixAugmentationBaseV2 already does for the mix classes.

Parameters:

  • input (Tensor) – Tensor produced by a forward pass.

  • params (Optional[List[ParamItem]], optional) – Parameters used during that forward pass. Defaults to the cached ones. Default: None

  • extra_args (Optional[Dict[str, Any]], optional) – Optional per-input-type overrides. Default: None

Return type:

Tensor

Returns:

The inverse-transformed tensor.

Raises:

RuntimeError – The drawn sub-policy contains a non-invertible operation.

class kornia.augmentation.auto.RandAugment(n, m, policy=None, transformation_matrix_mode='silent')[source]#

Apply RandAugment [CZSL20] augmentation strategies.

See the Convention block on PolicyAugmentBase.

Convention:

  • selects n distinct one-operation candidate sub-policies uniformly without replacement, then applies them in the order drawn. n must be in [1, len(policy)].

  • m must be strictly between 0 and 30 and sets the magnitude of every selected operation that has one to low + (high - low) * m / 30 over its magnitude range. A symmetric operation’s range is (0, max) with a random sign per row: m=15 on ("rotate", -30, 30) gives +15 or -15. A translate_x / translate_y value is a fraction of the image width / height, applied in pixels. The operation’s magnitude mapping is the identity for every default entry except shear_x / shear_y, which multiply by 180 (27 degrees at m=15 for the default range), and posterize, which truncates to integer bits. posterize also runs its range backwards, high - (high - low) * m / 30, so a larger m keeps fewer bits (7, 6 and 4 at m=7, 15 and 29).

Parameters:

  • n (int) – the number of augmentations to apply sequentially. Must be at least 1 and at most the number of sub-policies in policy, since they are sampled without replacement.

  • m (int) – magnitude for all the augmentations, ranged from (0, 30) exclusive.

  • policy (Optional[List[List[Tuple[str, Union[float, int], Union[float, int, None]]]]], optional) – candidate transformations. If None, a default candidate list will be used. Default: None

  • transformation_matrix_mode (str, optional) – computation mode for the chained transformation matrix, via .transform_matrix attribute. If silent, transformation matrix will be computed silently and the non-rigid modules will be ignored as identity transformations. If rigid, the result is the same for a policy: every operation wrapper contributes a matrix (an intensity operation the identity), so none is rejected. If skip, transformation matrix will be totally ignored. Default: "silent"

Examples

>>> import torch
>>> import kornia.augmentation as K
>>> from kornia.augmentation.auto import RandAugment
>>> in_tensor = torch.rand(5, 3, 30, 30)
>>> aug = K.AugmentationSequential(RandAugment(n=2, m=10))
>>> aug(in_tensor).shape
torch.Size([5, 3, 30, 30])
get_transformation_matrix(input, params=None, recompute=False, extra_args=None)[source]#

Compute the transformation matrix according to the provided parameters.

Parameters:

  • input (Tensor) – the input torch.Tensor.

  • params (Optional[List[ParamItem]], optional) – params for the sequence. Default: None

  • recompute (bool, optional) – if to recompute the transformation matrix according to the params. default: False. Default: False

  • extra_args (Optional[Dict[str, Any]], optional) – Optional dictionary of extra arguments with specific options for different input types. Default: None

Return type:

Optional[Tensor]

forward_parameters(batch_shape)[source]#

Sample parameters for the selected RandAugment operations.

Parameters:

batch_shape (Size) – Input shape used for parameter sampling.

Return type:

List[ParamItem]

Returns:

Parameters for the sampled operations, including magnitude scaled from the global m value.

forward(input, params=None, extra_args=None)[source]#

Apply the policy to input and cache the parameters used.

Parameters:

  • input (Tensor) – Input tensor.

  • params (Optional[List[ParamItem]], optional) – Optional precomputed parameters. If omitted, parameters are sampled on-the-fly. Default: None

  • extra_args (Optional[Dict[str, Any]], optional) – Optional overrides forwarded to child transforms. Default: None

Return type:

Tensor

Returns:

Augmented tensor.

inverse(input, params=None, extra_args=None)[source]#

Undo the drawn sub-policy, or refuse when it cannot be undone.

Only geometric operations are invertible. A policy draw that contains an intensity operation is not round-trippable, and inverting it would return a tensor that still carries that operation – for a draw with no geometry at all, the input unchanged. This raises instead, as inverse on MixAugmentationBaseV2 already does for the mix classes.

Parameters:

  • input (Tensor) – Tensor produced by a forward pass.

  • params (Optional[List[ParamItem]], optional) – Parameters used during that forward pass. Defaults to the cached ones. Default: None

  • extra_args (Optional[Dict[str, Any]], optional) – Optional per-input-type overrides. Default: None

Return type:

Tensor

Returns:

The inverse-transformed tensor.

Raises:

RuntimeError – The drawn sub-policy contains a non-invertible operation.

class kornia.augmentation.auto.TrivialAugment(policy=None, transformation_matrix_mode='silent')[source]#

Apply TrivialAugment [MullerH21] augmentation strategies.

See the Convention block on PolicyAugmentBase.

Convention:

  • selects one candidate sub-policy uniformly for each forward call. Each candidate holds one operation whose probability is fixed at 1, so the selected operation is applied to every row and no probability is sampled; a candidate’s second and third entries are its magnitude range, not a probability, unlike AutoAugment. Only the magnitude is drawn independently for every batch row.

  • the operation wrapper’s magnitude mapping applies, as under RandAugment. A ("rotate", -30, 30) candidate draws angles in [-30, 30] with a random sign, and ("shear_x", -0.3, 0.3) shears by up to 0.3 * 180 degrees.

Parameters:

  • policy (Optional[List[List[Tuple[str, Union[float, int], Union[float, int, None]]]]], optional) – candidate transformations. If None, a default candidate list will be used. Default: None

  • transformation_matrix_mode (str, optional) – computation mode for the chained transformation matrix, via .transform_matrix attribute. If silent, transformation matrix will be computed silently and the non-rigid modules will be ignored as identity transformations. If rigid, the result is the same for a policy: every operation wrapper contributes a matrix (an intensity operation the identity), so none is rejected. If skip, transformation matrix will be totally ignored. Default: "silent"

Examples

>>> import torch
>>> import kornia.augmentation as K
>>> from kornia.augmentation.auto import TrivialAugment
>>> in_tensor = torch.rand(5, 3, 30, 30)
>>> aug = K.AugmentationSequential(TrivialAugment())
>>> aug(in_tensor).shape
torch.Size([5, 3, 30, 30])
get_transformation_matrix(input, params=None, recompute=False, extra_args=None)[source]#

Compute the transformation matrix according to the provided parameters.

Parameters:

  • input (Tensor) – the input torch.Tensor.

  • params (Optional[List[ParamItem]], optional) – params for the sequence. Default: None

  • recompute (bool, optional) – if to recompute the transformation matrix according to the params. default: False. Default: False

  • extra_args (Optional[Dict[str, Any]], optional) – Optional dictionary of extra arguments with specific options for different input types. Default: None

Return type:

Optional[Tensor]

forward_parameters(batch_shape)[source]#

Generate per-module parameters for one policy forward pass.

Parameters:

batch_shape (Size) – Input shape used to sample operation parameters.

Return type:

List[ParamItem]

Returns:

Parameters for each selected child module, in execution order.

forward(input, params=None, extra_args=None)[source]#

Apply the policy to input and cache the parameters used.

Parameters:

  • input (Tensor) – Input tensor.

  • params (Optional[List[ParamItem]], optional) – Optional precomputed parameters. If omitted, parameters are sampled on-the-fly. Default: None

  • extra_args (Optional[Dict[str, Any]], optional) – Optional overrides forwarded to child transforms. Default: None

Return type:

Tensor

Returns:

Augmented tensor.

inverse(input, params=None, extra_args=None)[source]#

Undo the drawn sub-policy, or refuse when it cannot be undone.

Only geometric operations are invertible. A policy draw that contains an intensity operation is not round-trippable, and inverting it would return a tensor that still carries that operation – for a draw with no geometry at all, the input unchanged. This raises instead, as inverse on MixAugmentationBaseV2 already does for the mix classes.

Parameters:

  • input (Tensor) – Tensor produced by a forward pass.

  • params (Optional[List[ParamItem]], optional) – Parameters used during that forward pass. Defaults to the cached ones. Default: None

  • extra_args (Optional[Dict[str, Any]], optional) – Optional per-input-type overrides. Default: None

Return type:

Tensor

Returns:

The inverse-transformed tensor.

Raises:

RuntimeError – The drawn sub-policy contains a non-invertible operation.

class kornia.augmentation.auto.PolicyAugmentBase(policy, transformation_matrix_mode='silence')[source]#

Base class for policy-based image augmentations.

Convention:

  • a concrete policy selects one or more PolicySequential children for each forward call and records that selected path, including every operation parameter dictionary, in _params. Passing that list to forward(input, params=...) selects the recorded children rather than drawing a new path and reproduces their output.

  • the selected operations run in the order recorded in _params: the listed order inside a sub-policy, and for RandAugment the order in which its sub-policies were drawn. When matrix computation is enabled, their geometric transformation matrices compose in that same execution order; a nonempty policy containing only intensity operations has the identity matrix. An empty selected sub-policy has no matrix. inverse reverses a geometry-only selected path and raises RuntimeError when an applied intensity operation cannot be undone.

  • input normalization, random-number generation, parameter placement, replay, and serialization follow the canonical augmentation contract in Conventions & Pitfalls. This base does not expose a per-instance generator.

Concrete policies define how they select children and interpret policy magnitudes.

class kornia.augmentation.auto.PolicySequential(*operations)[source]#

Policy tuple for applying multiple operations.

Convention:

  • accepts only OperationBase children. Generated parameters apply every child in construction order. Supplied params define the execution path instead: their order and membership select the referenced children, so they may replay a reordered or partial path. forward stores the ParamItem objects for the path it executed in _params; those parameters can replay the same operations through forward(input, params=...).

  • its transformation matrix is the ordered product of its geometric wrapped operations. Intensity operations do not contribute to that matrix.

  • forward_parameters samples each child through forward_parameters(), so the operation’s magnitude, magnitude mapping and symmetric sign apply; the gate comes from the wrapped augmentation’s p.

Parameters:

operations (OperationBase) – the operations to perform, passed as positional arguments rather than as one list.

class kornia.augmentation.auto.operations.OperationBase(operation, initial_magnitude=None, temperature=0.1, is_batch_operation=False, magnitude_fn=None, symmetric_megnitude=False)[source]#

Base class of differentiable augmentation operations.

Parameters:

  • operation (_AugmentationBase) – Kornia augmentation module.

  • initial_magnitude (Optional[List[Tuple[str, Optional[float]]]], optional) – targeted magnitude parameter name and its initial magnitude value. The magnitude parameter name shall align with the attribute inside the random_generator in each augmentation. If None, the augmentation will be randomly applied according to the augmentation sampling range. Default: None

  • temperature (float, optional) – retained for API compatibility. Default: 0.1

  • is_batch_operation (bool, optional) – determine if to obtain the probability from p or p_batch. Set to True for most non-shape-persistent operations (e.g. cropping). Default: False

Convention:

  • probability is initialized from the wrapped augmentation’s p (p_batch for a batch operation), clamped to [1e-7, 1 - 1e-7] and kept in state_dict() for checkpoint compatibility, but it takes no part in sampling and receives no gradient: the gate is drawn from the wrapped augmentation’s own p and p_batch as a hard 0 or 1. magnitude, where the operation has one, is clamped to the wrapped generator’s range and receives a gradient where the wrapped augmentation is differentiable in it; forward_parameters substitutes it into the wrapped augmentation’s draw.

  • forward linearly blends the wrapped output with the input using batch_prob. The wrapped augmentation transforms every row whose gate is above 0, so the same supplied fractional batch_prob produces the same soft blend whatever the wrapped p or p_batch. Below p = 1 the wrapped augmentation skips a row with gate 0, so that row’s matrix stays the identity.

  • a symmetric magnitude applies the magnitude mapping first and then a random sign per row, so a mapping that quantizes to zero stays zero (Posterize maps 0.5 to 0 bits with magnitude_range=(0, 8)).

  • the concrete classes in kornia.augmentation.auto.operations.ops wrap public 2D augmentations and inherit their input, dtype, RNG and replay contracts. A wrapper pickles, with either sign setting, when its magnitude mapping is picklable.

forward_parameters(batch_shape, mag=None)[source]#

Sample parameters for the wrapped augmentation op.

Parameters:

  • batch_shape (Size) – Input batch shape used for sampling.

  • mag (Optional[Tensor], optional) – Optional magnitude override. When set, it replaces the sampled factor before the magnitude mapping function is applied. Default: None

Return type:

Dict[str, Tensor]

Returns:

Parameter dictionary consumed by forward().

forward(input, params=None)[source]#

Apply the operation with soft probabilistic blending.

Parameters:

  • input (Tensor) – Input tensor.

  • params (Optional[Dict[str, Tensor]], optional) – Optional precomputed parameters. If omitted, parameters are sampled from input.shape. Default: None

Return type:

Tensor

Returns:

Tensor blended with the wrapped augmentation’s output according to batch_prob.

Augmentation Search Methods#

WIP. This module contains common data augmentation search methods.