Class BNOps

executor/operator/arm32/batch_norm.cpp:71–201

}

struct BNOps : public NodeOps
{
    BNOps()
    {
        name_ = "arm_batchnorm_fp32";
    }

    bool OnBind(Node* node)
    {
        // set the inplace feature
        inplace_t io_map;

        io_map[0] = 0;

        node->SetAttr(ATTR_INPLACE, io_map);

        return true;
    }

    bool Prerun(Node* node)
    {
        const Tensor* mean_tensor = node->GetInputTensor(3);
        const TShape& shape = mean_tensor->GetShape();

        const std::vector<int> dims = shape.GetDim();

        int channel_num = dims[0];

        float* scale_mean = ( float* )mem_alloc(channel_num * sizeof(float));
        float* scale_var_inv = ( float* )mem_alloc(channel_num * sizeof(float));

        const Tensor* var_tensor = node->GetInputTensor(4);
        const float* mean = ( const float* )get_tensor_mem(mean_tensor);
        const float* var = ( const float* )get_tensor_mem(var_tensor);

        BatchNorm* bn_op = dynamic_cast<BatchNorm*>(node->GetOp());
        BatchNormParam* param = bn_op->GetParam();

        float rescale_factor;
        float eps = param->eps;

        rescale_factor = param->rescale_factor ? 1 / param->rescale_factor : 0;
        for(int c = 0; c < channel_num; c++)
        {
            float tmp = std::sqrt(var[c] * rescale_factor + eps);
            scale_var_inv[c] = (float)(1.f / tmp);
            tmp = rescale_factor * scale_var_inv[c];
            scale_mean[c] = (float)(-mean[c] * tmp);
        }
        if(!param->caffe_flavor)
        {
            const Tensor* gamma_tensor = node->GetInputTensor(1);
            const Tensor* beta_tensor = node->GetInputTensor(2);
            const float* gamma = ( const float* )get_tensor_mem(gamma_tensor);
            const float* beta = ( const float* )get_tensor_mem(beta_tensor);
            for(int c = 0; c < channel_num; c++)
            {
                scale_var_inv[c] *= gamma[c];