BatchNorm

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product/Atlas A3 inference product	√
Atlas A2 training product/Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	x
Atlas inference product AI Core	√
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

BatchNorm normalizes the input of each layer to make the distribution of each layer as similar as possible, thereby expediting training and improving the model's generalization capability (effectively reducing vanishing and exploding gradients). The basic idea is to normalize each input feature of samples in each batch along the batch dimension. Specifically, for input feature x, the calculation process of BatchNorm can be expressed as:

For input feature x, calculate the mean μ and variance σ along the batch dimension:
$\text{[math]}$
For each feature i, normalize input feature x:
$\text{[math]}$
Scale and translate the normalized feature:
$\text{[math]}$

Prototype

Allocate the temporary space through the API framework.

template <typename T, bool isReuseSource = false, bool isBasicBlock = false>
__aicore__ inline void BatchNorm(const LocalTensor<T>& output, const LocalTensor<T>& outputMean, const LocalTensor<T>& outputVariance, const LocalTensor<T>& inputX, const LocalTensor<T>& gamm, const LocalTensor<T>& beta, const T epsilon, BatchNormTiling& tiling)

Pass the temporary space through the sharedTmpBuffer input parameter.

template <typename T, bool isReuseSource = false, bool isBasicBlock = false>
__aicore__ inline void BatchNorm(const LocalTensor<T>& output, const LocalTensor<T>& outputMean, const LocalTensor<T>& outputVariance, const LocalTensor<T>& inputX, const LocalTensor<T>& gamm, const LocalTensor<T>& beta, const LocalTensor<uint8_t>& sharedTmpBuffer, const T epsilon, BatchNormTiling& tiling)

Parameters

**Table 1** Template parameters
Parameter	Description
T	Data type of the operand. For the Atlas 350 Accelerator Card, the supported data types are half and float. For the Atlas A3 training product/Atlas A3 inference product, the supported data types are half and float. For the Atlas A2 training product/Atlas A2 inference product, the supported data types are half and float. For the Atlas inference product AI Core, the supported data types are half and float.
isReuseSource	Whether the source operand can be modified. This parameter is reserved. Pass the default value false.
isBasicBlock	If the shape information and the tiling policy of inputX and output meet the basic block requirements, this parameter can be enabled to improve performance. By default, this parameter is disabled. The basic block requirements are as follows: The value of originB is a multiple of 8. The value of S × H is a multiple of 64 but less than 2048.

**Table 2** API parameters
Parameter	Input/Output	Description
output	Output	Destination operand, with a shape of [B, S, H]. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
outputMean	Output	Mean, destination operand, with a shape of [S, H]. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
outputVariance	Output	Variance, destination operand, with a shape of [S, H]. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
inputX	Input	Source operand, with a shape of [B, S, H]. The data type of inputX must be the same as that of the destination operand, and the value of *SH must be 32-byte aligned. The address of inputX can overlap with that of output. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT**.
gamm	Input	Source operand, with a shape of [B]. The data type of gamm must be the same as that of the destination operand, and the length must be 32-byte aligned. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
beta	Input	Source operand, with a shape of [B]. The data type of beta must be the same as that of the destination operand, and the length must be 32-byte aligned. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
sharedTmpBuffer	Input	This parameter is used to store intermediate variables during complex internal API computation and is provided by developers. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. For details about how to obtain the temporary space size (BufferSize), see BatchNorm Tiling.
epsilon	Input	Weight coefficient for preventing division by zero. The data type must be the same as that of inputX or output.
tiling	Input	Tiling information of input data. For details about how to obtain the tiling information, see BatchNorm Tiling.

Returns

None

Constraints

For details about the operand address alignment requirements, see General Address Alignment Restrictions.
Currently, only the ND format is supported.
The S*H of input data must be 32-byte aligned.

Examples

For a complete call example, see BatchNorm operator sample.

// outputLocal: tensor for storing the BatchNorm computation result
// meanLocal: tensor for storing the mean value of the computation result
// varianceLocal: tensor for storing the variance of the computation result
// inputXLocal: input tensor involved in computation
// gammaLocal: input tensor, scaling coefficient γ of the normalized data
// betaLocal: input tensor, translation coefficient β of the normalized data
// epsilon: weight coefficient ε for preventing division by zero
// batchNormTiling: tiling data, obtained from the host

AscendC::BatchNorm<dataType, isReuseSource, isBasicBlock>(outputLocal, meanLocal,varianceLocal, 
                                                          inputXLocal, gammaLocal, betaLocal, 
                                                          (dataType)epsilon, batchNormTiling); 

The following is an example:

Input (inputXLocal, shape: [8, 4, 2]):
[  0  1  2  3  4  5  6  7  
   8  9 10 11 12 13 14 15
  16 17 18 19 20 21 22 23 
  24 25 26 27 28 29 30 31 
  32 33 34 35 36 37 38 39 
  40 41 42 43 44 45 46 47 
  48 49 50 51 52 53 54 55 
  56 57 58 59 60 61 62 63 ]
Input (gammaLocal, shape: [4]):
[ 0 1 2 3 4 5 6 7 ]
Input (betaLocal, shape: [4]):
[ 0 1 2 3 4 5 6 7 ]
Output (dstLocal):
[ 0. 0. 0. 0. 0. 0. 0. 0. 
  -0.091073155 -0.091073155 -0.091073155 -0.091073155 -0.091073155 -0.091073155 -0.091073155 -0.091073155 
  0.6907122 0.6907122 0.6907122 0.6907122 0.6907122 0.6907122 0.6907122 0.6907122 
  2.345356 2.345356 2.345356 2.345356 2.345356 2.345356 2.345356 2.345356 
  4.8728585 4.8728585 4.8728585 4.8728585 4.8728585 4.8728585 4.8728585 4.8728585 
  8.27322 8.27322 8.27322 8.27322 8.27322 8.27322 8.27322 8.27322 
  12.546439 12.546439 12.546439 12.546439 12.546439 12.546439 12.546439 12.546439 
  17.692516 17.692516 17.692516 17.692516 17.692516 17.692516 17.692516 17.692516 ]
Output (meanLocal):
[ 28 29 30 31 32 33 34 35 ]
Output (varianceLocal):
[ 336 336 336 336 336 336 336 336 ]

Parent topic: Normalization