ttnn.batch_norm

ttnn.batch_norm = Operation(python_fully_qualified_name='ttnn.batch_norm', function=<ttnn._ttnn.operations.normalization.batch_norm_t object>, preprocess_golden_function_inputs=<function default_preprocess_golden_function_inputs>, golden_function=None, postprocess_golden_function_outputs=<function default_postprocess_golden_function_outputs>, is_cpp_operation=True, is_experimental=False)

ttnn.batch_norm(input: ttnn.Tensor, running_mean: Optional[ttnn.Tensor] = None, running_var: Optional[ttnn.Tensor] = None, training: bool = False, eps: float = 1e-05, momentum: float = 0.1, weight: Optional[ttnn.Tensor] = None, bias: Optional[ttnn.Tensor] = None, output: Optional[ttnn.Tensor] = None, memory_config: Optional[ttnn.MemoryConfig] = None, compute_kernel_config: Optional[ttnn.DeviceComputeKernelConfig] = None) -> ttnn.Tensor

Applies batch norm over each channel on input_tensor. See Spatial Batch Normalization for more details.

\[\text{batch_norm}(x, \gamma, \beta, \epsilon) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta\]

Where:

\(\mu\) and \(\sigma^2\) are the mean and variance of the input tensor, respectively
\(\gamma\) and \(\beta\) are the learnable scale and shift parameters, respectively
\(\epsilon\) is a small constant.

Parameters:

input_tensor (ttnn.Tensor) – the input tensor of shape [N, C, H, W].

Keyword Arguments:

eps (float, optional) – Epsilon value. Defaults to 1e-05.
momentum (float, optional) – Momentum value. Defaults to 0.1.
running_mean (ttnn.Tensor, optional) – the running_mean of shape [1, C, 1, 1], required in inference mode. When in training mode, this tensor is optional and the updated running mean value is stored in-place based on the inputs provided. Defaults to None.
running_var (ttnn.Tensor, optional) – the running_var of shape [1, C, 1, 1], required in inference mode. When in training mode, this tensor is optional and the updated running variance value is stored in-place based on the inputs provided. Defaults to None.
weight (ttnn.Tensor, optional) – the weight or gamma value of shape [1, C, 1, 1]. Defaults to None.
bias (ttnn.Tensor, optional) – the bias or beta value of shape [1, C, 1, 1]. Defaults to None.
training (bool, optional) – Selection between training mode and inference (evaluation) mode. Defaults to False (Inference mode).
output (ttnn.Tensor, optional) – Preallocated output tensor to store batch norm result of shape [N, C, H, W]. Defaults to None.
memory_config (ttnn.MemoryConfig, optional) – memory configuration for the operation. Defaults to None.
compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) – device compute kernel configuration for the operation. Defaults to None.

Returns:

ttnn.Tensor – the output tensor.

Note

Supported dtypes, layouts, and ranks:

Dtypes	Layouts	Ranks
BFLOAT16, FLOAT32	TILE	4

These apply for all the tensor inputs to this operation, including the optional output tensor.

Memory Support:

Interleaved: DRAM and L1

Limitations:

All input tensors must be tilized, interleaved, rank 4, and on-device.

Example

N, C, H, W = 2, 3, 4, 5

input_tensor = ttnn.rand([N, C, H, W], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
running_mean = ttnn.rand([1, C, 1, 1], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
running_var = ttnn.rand([1, C, 1, 1], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
weight = ttnn.rand([1, C, 1, 1], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
bias = ttnn.from_torch(torch.rand([1, C, 1, 1], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device)

output = ttnn.batch_norm(
    input_tensor,
    running_mean = running_mean,
    running_var = running_var,
    weight = weight,
    bias = bias,
    eps = 1e-05,
    momentum = 0.1,
    training = True
)