BatchNorm

功能说明

BatchNorm是对于每一层的输入做规范化处理，使得每一层的分布尽可能的相同，从而加速训练过程和提高模型的泛化能力（有效减少梯度消失和梯度爆炸问题）。基本思想是对于每个 batch中的样本，对其输入的每个特征在batch的维度上进行归一化。具体来说，对于输入特征x，BatchNorm的计算过程可以表示为：

1. 对输入特征x，在batch维度上计算均值μ和方差σ：

   

2. 对于每个特征i，对输入特征x进行归一化：

   

3、对归一化后的特征进行缩放和平移：

函数原型

• 接口框架申请临时空间

  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)

• 通过sharedTmpBuffer入参传入临时空间

  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)

参数说明

返回值

无

支持的型号

Atlas推理系列产品AI Core

Atlas A2训练系列产品/Atlas 800I A2推理产品

注意事项

• 操作数地址偏移对齐要求请参见通用约束。
• 当前仅支持ND格式的输入，不支持其他格式。
• 输入数据的S*H必须满足32B对齐的要求。
• 支持源操作数与目的操作数地址重叠。

调用示例

#include "kernel_operator.h"
using namespace AscendC;
template <typename dataType, bool isReuseSource = false, bool isBasicBlock = false> class KernelBatchnorm {
public:
    __aicore__ inline KernelBatchnorm() {}
    __aicore__ inline void Init(GM_ADDR inputX_gm, GM_ADDR gamm_gm, GM_ADDR beta_gm, GM_ADDR output_gm,
        GM_ADDR outputMean_gm, GM_ADDR outputVariance_gm, const BatchNormTiling& tiling)
    {
        this->bLength = tiling.bLength;
        this->sLength = tiling.sLength;
        this->hLength = tiling.hLength;
        this->originalBLength = tiling.originalBLength;

        bshLength = originalBLength * sLength * hLength;
        shLength = sLength * hLength;

        inputX_global.SetGlobalBuffer(reinterpret_cast<__gm__ dataType*>(inputX_gm), bshLength);
        gamm_global.SetGlobalBuffer(reinterpret_cast<__gm__ dataType*>(gamm_gm), bLength);
        beta_global.SetGlobalBuffer(reinterpret_cast<__gm__ dataType*>(beta_gm), bLength);

        output_global.SetGlobalBuffer(reinterpret_cast<__gm__ dataType*>(output_gm), bshLength);
        outputMean_global.SetGlobalBuffer(reinterpret_cast<__gm__ dataType*>(outputMean_gm), shLength);
        outputVariance_global.SetGlobalBuffer(reinterpret_cast<__gm__ dataType*>(outputVariance_gm), shLength);

        pipe.InitBuffer(inQueueX, 1, sizeof(dataType) * bshLength);
        pipe.InitBuffer(inQueueGamma, 1, sizeof(dataType) * bLength);
        pipe.InitBuffer(inQueueBeta, 1, sizeof(dataType) * bLength);
        pipe.InitBuffer(outQueue, 1, sizeof(dataType) * bshLength);
        pipe.InitBuffer(outQueueMean, 1, sizeof(dataType) * shLength);
        pipe.InitBuffer(outQueueVariance, 1, sizeof(dataType) * shLength);

    }
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn()
    {
        LocalTensor<dataType> inputXLocal = inQueueX.AllocTensor<dataType>();
        LocalTensor<dataType> gammaLocal = inQueueGamma.AllocTensor<dataType>();
        LocalTensor<dataType> betaLocal = inQueueBeta.AllocTensor<dataType>();

        DataCopy(inputXLocal, inputX_global, bshLength);
        DataCopy(gammaLocal, gamm_global, bLength);
        DataCopy(betaLocal, beta_global, bLength);

        inQueueX.EnQue(inputXLocal);
        inQueueGamma.EnQue(gammaLocal);
        inQueueBeta.EnQue(betaLocal);
    }
    __aicore__ inline void Compute()
    {
        LocalTensor<dataType> inputXLocal = inQueueX.DeQue<dataType>();
        LocalTensor<dataType> gammaLocal = inQueueGamma.DeQue<dataType>();
        LocalTensor<dataType> betaLocal = inQueueBeta.DeQue<dataType>();

        LocalTensor<dataType> outputLocal = outQueue.AllocTensor<dataType>();
        LocalTensor<dataType> meanLocal = outQueueMean.AllocTensor<dataType>();
        LocalTensor<dataType> varianceLocal = outQueueVariance.AllocTensor<dataType>();
        BatchNorm<dataType, isReuseSource, isBasicBlock>(outputLocal, meanLocal, varianceLocal,inputXLocal, gammaLocal, betaLocal, (dataType)epsilon, tilling);
        outQueue.EnQue<dataType>(outputLocal);
        outQueueMean.EnQue<dataType>(meanLocal);
        outQueueVariance.EnQue<dataType>(varianceLocal);

        inQueueX.FreeTensor(inputXLocal);
        inQueueGamma.FreeTensor(gammaLocal);
        inQueueBeta.FreeTensor(betaLocal);
    }
    __aicore__ inline void CopyOut()
    {
        LocalTensor<dataType> outputLocal = outQueue.DeQue<dataType>();
        LocalTensor<dataType> meanLocal = outQueueMean.DeQue<dataType>();
        LocalTensor<dataType> varianceLocal = outQueueVariance.DeQue<dataType>();

        DataCopy(output_global, outputLocal, bshLength);
        DataCopy(outputMean_global, meanLocal, shLength);
        DataCopy(outputVariance_global, varianceLocal, shLength);

        outQueue.FreeTensor(outputLocal);
        outQueueMean.FreeTensor(meanLocal);
        outQueueVariance.FreeTensor(varianceLocal);
    }

private:
    GlobalTensor<dataType> inputX_global;
    GlobalTensor<dataType> gamm_global;
    GlobalTensor<dataType> beta_global;
    GlobalTensor<dataType> output_global;
    GlobalTensor<dataType> outputMean_global;
    GlobalTensor<dataType> outputVariance_global;

    TPipe pipe;
    TQue<QuePosition::VECIN, 1> inQueueX;
    TQue<QuePosition::VECIN, 1> inQueueGamma;
    TQue<QuePosition::VECIN, 1> inQueueBeta;
    TQue<QuePosition::VECOUT, 1> outQueue;
    TQue<QuePosition::VECOUT, 1> outQueueMean;
    TQue<QuePosition::VECOUT, 1> outQueueVariance;

    uint32_t bLength;
    uint32_t sLength;
    uint32_t hLength;
    uint32_t originalBLength;
    dataType epsilon = 0.001;

    uint32_t bshLength;
    uint32_t shLength;

    BatchNormTiling tilling;
};

extern "C" __global__ __aicore__ void kernel_batchnorm_operator(GM_ADDR inputX_gm, GM_ADDR gamm_gm, GM_ADDR beta_gm, GM_ADDR output_gm, GM_ADDR outputMean_gm, GM_ADDR outputVariance_gm, GM_ADDR tiling)
{
    GET_TILING_DATA(tilingData, tiling);
    KernelBatchnorm<half, false, false> op;
    op.Init(inputX_gm, gamm_gm, beta_gm, output_gm, outputMean_gm, outputVariance_gm, tilingData.batchnormTilingData);
    op.Process();
}