我想实现实例规范化

Question

我正在写一个金属cnn代码。 Metal提供MPSCNNLocalContrastNormalization， 由于实例规范化的概念略有不同，我打算将其实现为内核函数。

然而，问题是当从核函数中的输入接收到纹理中的特征是R，G，B时，应当获得每个R，G，B的均值和方差。 我想得到一些关于如何实现这一点的提示。

kernel void instance_normalization_2darray(texture2d_array<float, access::sample> src [[ texture(0) ]],
                                        texture2d_array<float, access::write> dst [[ texture(1) ]],
                                        uint3 tid [[thread_position_in_grid]]) {

}


    kernel void calculate_avgA(texture2d_array<float, access::read> texture_in [[texture(0)]],
                          texture2d_array<float, access::write> texture_out [[texture(1)]],
                          uint3 tid [[thread_position_in_grid]])
{
    int width = texture_in.get_width();
    int height = texture_in.get_height();
    int depth = texture_in.get_array_size();
    float4 outColor;


    uint3 kernelIndex(0,0,0);
    uint3 textureIndex(0,0,0);

    for(int k = 0; k < depth; k++) {
        outColor = (0.0, 0.0, 0.0, 0.0);
        for (int i=0; i < width; i++)
        {
            for (int j=0; j < height; j++)
            {
                kernelIndex = uint3(i, j, k);
                textureIndex = uint3(tid.x + i, tid.y + j, tid.z + k);
                float4 color = texture_in.read(textureIndex.xy, textureIndex.z).rgba;
                outColor += color;
            }
        }
        outColor = outColor / (width * height);
        texture_out.write(float4(outColor.rgba), tid.xy, textureIndex.z);
    }
}