首页 文章

CoreML自定义图层:使用金属着色器进行像素标准化

提问于
浏览
0

我正在将Nvidia的GAN发电机的逐步发展转换为coreML . 我已经设法将所有内容转移到coreML,但Pixelwise规范化(Lambda)层除外,我计划将其作为Swift / Metal中的自定义coreML层实现 .

在TensorFlow.Keras中,我已经将像素规范实现为

def pixelwise_norm(a):
    return a / tf.sqrt(tf.reduce_mean(a * a, axis=3, keep_dims=True) + 1e-8)

现在,我几乎没有使用着色器/金属,但按照这里的说明:http://machinethink.net/blog/coreml-custom-layers/,我有一个自定义图层设置为使用Metal进行前馈操作 . 我正在使用MTLComputePipelineState(调用?编码?)以下着色器进行图层操作:

#include <metal_stdlib>
using namespace metal;


kernel void pixelwise_norm(
              texture2d_array<half, access::read> inTexture [[texture(0)]],
              texture2d_array<half, access::write> outTexture [[texture(1)]],
              ushort3 gid [[thread_position_in_grid]])
{
    if (gid.x >= outTexture.get_width() ||
        gid.y >= outTexture.get_height()) {
        return;
    }

    const float4 x = float4(inTexture.read(gid.xy, gid.z));
    const float4 y = 0.0000001f + (x / sqrt(pow(x,2)));
    outTexture.write(half4(y), gid.xy, gid.z);
}

我无法搞清楚“reduce_mean”的金属等价物,现在这个着色器实现了一个〜tensorflow~操作就像

return a / tf.sqrt((a * a) + 1e-8)

有没有人有任何指针?谢谢

1 回答

  • 0

    如果我正确读取这个,对于要素图中的每个像素,这个像素除了那个像素的通道上的L2范数?

    在这种情况下,您需要使用for循环来读取该像素的通道,将这些数字相加,然后除以通道数 . (如果通道数超过4,则只需要执行此循环 . )

    另请注意,您的1e-8需要位于sqrt()内或至少在分母内 .

相关问题