Class BatchNormOps

executor/operator/common/batchnorm.cpp:39–166

namespace BatchNormImpl {

struct BatchNormOps : public NodeOps
{
    bool Prerun(Node* node)
    {
        const Tensor* input_tensor = node->GetInputTensor(0);
        const TShape& shape = input_tensor->GetShape();

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

        int channel_num = dims[1];

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

        const Tensor* mean_tensor = node->GetInputTensor(3);
        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++)
        {
            scale_var_inv[c] = 1.f / sqrt(var[c] * rescale_factor + eps);
            scale_mean[c] = -mean[c] * rescale_factor * scale_var_inv[c];
        }

        node->SetAttr("scale_mean", scale_mean);
        node->SetAttr("scale_var_inv", scale_var_inv);

        return true;
    }

    bool Run(Node* node)
    {
        const Tensor* input_tensor = node->GetInputTensor(0);
        Tensor* output_tensor = node->GetOutputTensor(0);
        const TShape& shape = input_tensor->GetShape();
        const std::vector<int> dims = shape.GetDim();

        int batch_number = dims[0];
        int channel_num = dims[1];
        int channel_size = dims[2] * dims[3];
        int img_size = channel_num * channel_size;

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

        const float* input = ( const float* )get_tensor_mem(input_tensor);
        float* output = ( float* )get_tensor_mem(output_tensor);

        if(param->caffe_flavor)
        {