Method Run

executor/operator/arm64/batch_norm.cpp:140–189

    }

    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;

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

        float* scale_mean = any_cast<float*>(node->GetAttr("scale_mean"));
        float* scale_var_inv = any_cast<float*>(node->GetAttr("scale_var_inv"));


        int cpu_number = cpu_info->GetCPUNumber();
        int block = channel_num;
        block = block > 0 ? block : 1;
        int num_task = cpu_number < block ? cpu_number : block;
        int step = channel_num / num_task;
        /* only use mean and var */
        for(int i = 0; i < batch_number; i++)
        {
            const float* cur_input = input + i * img_size;
            float* cur_output = output + i * img_size;

            if(num_task == 1)
                batchnorm_kernel( 0, 0, &step, cur_input, cur_output, scale_mean, scale_var_inv, channel_size);
            else
            {
                MULTI_THREAD_START(num_task, step, p_id, p_param)
                    batchnorm_kernel( 0, p_id, p_param, cur_input, cur_output, scale_mean, scale_var_inv, channel_size);
                MULTI_THREAD_END();
            }
            if(num_task * step != channel_num)
            {
                int offset = num_task * step;
                int remain_num = channel_num - offset;
                cur_input += offset * channel_size;
                cur_output += offset * channel_size;
                batchnorm_kernel( 0, 0, &remain_num, cur_input, cur_output, scale_mean + offset, scale_var_inv + offset, channel_size);
            }
        }

        return true;
    }

    bool Postrun(Node* node)
    {