Как я могу рассчитать среднее значение и значение дисперсии изображения с 16 каналами, используя Metal Shader Lanuage

Question

как я могу рассчитать среднее значение и значение дисперсии изображения с 16 каналами, используя Metal?

Я хочу рассчитать среднее значение и значение дисперсии для каждого канала отдельно!

пример:

kernel void meanandvariance(texture2d_array<float, access::read> in[[texture(0)]],
                          texture2d_array<float, access::write> out[[texture(1)]],

                          ushort3 gid[[thread_position_in_grid]],
                          ushort tid[[thread_index_in_threadgroup]],
                          ushort3 tg_size[[threads_per_threadgroup]]) {

                          }

Answer 1

#include <metal_stdlib>
using namespace metal;

kernel void instance_norm(constant float4* scale[[buffer(0)]],
                          constant float4* shift[[buffer(1)]],
                          texture2d_array<float, access::read> in[[texture(0)]],
                          texture2d_array<float, access::write> out[[texture(1)]],

                          ushort3 gid[[thread_position_in_grid]],
                          ushort tid[[thread_index_in_threadgroup]],
                          ushort3 tg_size[[threads_per_threadgroup]]) {

    ushort width = in.get_width();
    ushort height = in.get_height();
    const ushort thread_count = tg_size.x * tg_size.y;

    threadgroup float4 shared_mem [256];

    float4 sum = 0;
    for(ushort xIndex = gid.x; xIndex < width; xIndex += tg_size.x) {
        for(ushort yIndex = gid.y; yIndex < height; yIndex += tg_size.y) {
            sum += in.read(ushort2(xIndex, yIndex), gid.z);
        }
    }
    shared_mem[tid] = sum;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Reduce to 32 values
    sum = 0;
    if (tid < 32) {
        for (ushort i = tid + 32; i < thread_count; i += 32) {
            sum += shared_mem[i];
        }
    }
    shared_mem[tid] += sum;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Calculate mean
    sum = 0;
    if (tid == 0) {
        ushort top = min(ushort(32), thread_count);
        for (ushort i = 0; i < top; i += 1) {
            sum += shared_mem[i];
        }
        shared_mem[0] = sum / (width * height);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    const float4 mean = shared_mem[0];

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Variance
    sum = 0;
    for(ushort xIndex = gid.x; xIndex < width; xIndex += tg_size.x) {
        for(ushort yIndex = gid.y; yIndex < height; yIndex += tg_size.y) {
            sum += pow(in.read(ushort2(xIndex, yIndex), gid.z) - mean, 2);
        }
    }

    shared_mem[tid] = sum;
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Reduce to 32 values
    sum = 0;
    if (tid < 32) {
        for (ushort i = tid + 32; i < thread_count; i += 32) {
            sum += shared_mem[i];
        }
    }
    shared_mem[tid] += sum;

    threadgroup_barrier(mem_flags::mem_threadgroup);

    // Calculate variance
    sum = 0;
    if (tid == 0) {
        ushort top = min(ushort(32), thread_count);
        for (ushort i = 0; i < top; i += 1) {
            sum += shared_mem[i];
        }
        shared_mem[0] = sum / (width * height);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    const float4 sigma = sqrt(shared_mem[0] + float4(1e-4));

    float4 multiplier = scale[gid.z] / sigma;
    for(ushort xIndex = gid.x; xIndex < width; xIndex += tg_size.x) {
        for(ushort yIndex = gid.y; yIndex < height; yIndex += tg_size.y) {
            float4 val = in.read(ushort2(xIndex, yIndex), gid.z);
            out.write(clamp((val - mean) * multiplier + shift[gid.z], -10.0, 10.0), ushort2(xIndex, yIndex), gid.z);
        }
    }

}

это то, как Blend реализует, но я не думаю, что это правда, кто-нибудь может решить это?

https://github.com/xmartlabs/Bender/blob/master/Sources/Metal/instanceNorm.metal

Answer 2

Вероятно, есть способ сделать это, создав последовательность представлений текстуры на входном массиве текстуры и выходном массиве текстуры, кодируя MPSImageStatisticsMeanAndVariance вызов ядра для каждого среза.

Но давайте посмотрим, как это сделать самим. Существует множество возможных подходов, поэтому я выбрал простой и использовал некоторые интересные результаты из статистики.

По существу, мы сделаем следующее:

1. Напишите ядро, которое может вывести среднее значение и дисперсию для одной строки изображения.
2. Запись ядро, которое может дать общее среднее значение и отклонение от частичных результатов, полученных на шаге 1.

Вот ядра:

kernel void compute_row_mean_variance_array(texture2d_array<float, access::read> inTexture [[texture(0)]],
                                            texture2d_array<float, access::write> outTexture [[texture(1)]],
                                            uint3 tpig [[thread_position_in_grid]])
{
    uint row = tpig.x;
    uint slice = tpig.y;
    uint width = inTexture.get_width();

    if (row >= inTexture.get_height() || slice >= inTexture.get_array_size()) { return; }

    float4 mean(0.0f);
    float4 var(0.0f);
    for (uint col = 0; col < width; ++col) {
        float4 rgba = inTexture.read(ushort2(col, row), slice);
        // http://datagenetics.com/blog/november22017/index.html
        float weight = 1.0f / (col + 1);
        float4 oldMean = mean;
        mean = mean + (rgba - mean) * weight;
        var = var + (rgba - oldMean) * (rgba - mean);
    }

    var = var / width;

    outTexture.write(mean, ushort2(row, 0), slice);
    outTexture.write(var, ushort2(row, 1), slice);
}

kernel void reduce_mean_variance_array(texture2d_array<float, access::read> inTexture [[texture(0)]],
                                       texture2d_array<float, access::write> outTexture [[texture(1)]],
                                       uint3 tpig [[thread_position_in_grid]])
{
    uint width = inTexture.get_width();
    uint slice = tpig.x;

    // https://arxiv.org/pdf/1007.1012.pdf
    float4 mean(0.0f);
    float4 meanOfVar(0.0f);
    float4 varOfMean(0.0f);
    for (uint col = 0; col < width; ++col) {
        float weight = 1.0f / (col + 1);

        float4 oldMean = mean;
        float4 submean = inTexture.read(ushort2(col, 0), slice);
        mean = mean + (submean - mean) * weight;

        float4 subvar = inTexture.read(ushort2(col, 1), slice);
        meanOfVar = meanOfVar +  (subvar - meanOfVar) * weight;

        varOfMean = varOfMean + (submean - oldMean) * (submean - mean);
    }
    float4 var = meanOfVar + varOfMean / width;

    outTexture.write(mean, ushort2(0, 0), slice);
    outTexture.write(var, ushort2(1, 0), slice);
}

В итоге, чтобы выполнить шаг 1, мы используйте «онлайн» (инкрементный) алгоритм для вычисления частичного среднего значения / дисперсии строки таким образом, чтобы она была более численно устойчивой, чем просто сложение всех значений пикселей и деление на ширину. Моя ссылка на написание этого ядра была в этом посте . Каждый поток в сетке записывает статистику своей строки в соответствующий столбец и фрагмент промежуточного массива текстур.

Чтобы выполнить шаг 2, нам нужно найти статистически обоснованный способ вычисления общей статистики на основе частичных результатов. , Это довольно просто в случае нахождения среднего значения: среднее из совокупности является средним из средних значений подмножеств (это имеет место, когда размер выборки каждого подмножества одинаков; в общем случае общее среднее значение представляет собой взвешенная сумма подмножества означает). Дисперсия сложнее, но превращает из в то, что дисперсия совокупности представляет собой сумму среднего значения дисперсий подмножеств и дисперсии средних подмножеств ( То же самое относится и к подмножествам одинакового размера. Это удобный факт, который мы можем объединить с нашим инкрементным подходом выше для получения окончательного среднего значения и дисперсии каждого среза, который записывается в соответствующий срез выходной текстуры.

Для полноты вот код Swift Я использовал для управления этими ядрами:

let library = device.makeDefaultLibrary()!

let meanVarKernelFunction = library.makeFunction(name: "compute_row_mean_variance_array")!
let meanVarComputePipelineState = try! device.makeComputePipelineState(function: meanVarKernelFunction)

let reduceKernelFunction = library.makeFunction(name: "reduce_mean_variance_array")!
let reduceComputePipelineState = try! device.makeComputePipelineState(function: reduceKernelFunction)

let width = sourceTexture.width
let height = sourceTexture.height
let arrayLength = sourceTexture.arrayLength

let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .rgba32Float, width: width, height: height, mipmapped: false)
textureDescriptor.textureType = .type2DArray
textureDescriptor.arrayLength = arrayLength
textureDescriptor.width = height
textureDescriptor.height = 2
textureDescriptor.usage = [.shaderRead, .shaderWrite]

let partialResultsTexture = device.makeTexture(descriptor: textureDescriptor)!

textureDescriptor.width = 2
textureDescriptor.height = 1
textureDescriptor.usage = .shaderWrite

let destTexture = device.makeTexture(descriptor: textureDescriptor)!

let commandBuffer = commandQueue.makeCommandBuffer()!

let computeCommandEncoder = commandBuffer.makeComputeCommandEncoder()!

computeCommandEncoder.setComputePipelineState(meanVarComputePipelineState)
computeCommandEncoder.setTexture(sourceTexture, index: 0)
computeCommandEncoder.setTexture(partialResultsTexture, index: 1)
let meanVarGridSize = MTLSize(width: sourceTexture.height, height: sourceTexture.arrayLength, depth: 1)
let meanVarThreadgroupSize = MTLSizeMake(meanVarComputePipelineState.threadExecutionWidth, 1, 1)
let meanVarThreadgroupCount = MTLSizeMake((meanVarGridSize.width + meanVarThreadgroupSize.width - 1) / meanVarThreadgroupSize.width,
                                          (meanVarGridSize.height + meanVarThreadgroupSize.height - 1) / meanVarThreadgroupSize.height,
                                          1)
computeCommandEncoder.dispatchThreadgroups(meanVarThreadgroupCount, threadsPerThreadgroup: meanVarThreadgroupSize)

computeCommandEncoder.setComputePipelineState(reduceComputePipelineState)
computeCommandEncoder.setTexture(partialResultsTexture, index: 0)
computeCommandEncoder.setTexture(destTexture, index: 1)
let reduceThreadgroupSize = MTLSizeMake(1, 1, 1)
let reduceThreadgroupCount = MTLSizeMake(arrayLength, 1, 1)
computeCommandEncoder.dispatchThreadgroups(reduceThreadgroupCount, threadsPerThreadgroup: reduceThreadgroupSize)

computeCommandEncoder.endEncoding()

let destTexture2DDesc = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .rgba32Float, width: 2, height: 1, mipmapped: false)
destTexture2DDesc.usage = .shaderWrite
let destTexture2D = device.makeTexture(descriptor: destTexture2DDesc)!

meanVarKernel.encode(commandBuffer: commandBuffer, sourceTexture: sourceTexture2D, destinationTexture: destTexture2D)

#if os(macOS)
let blitCommandEncoder = commandBuffer.makeBlitCommandEncoder()!
blitCommandEncoder.synchronize(resource: destTexture)
blitCommandEncoder.synchronize(resource: destTexture2D)
blitCommandEncoder.endEncoding()
#endif

commandBuffer.commit()

commandBuffer.waitUntilCompleted()

В моих экспериментах эта программа дала те же результаты, что и MPSImageStatisticsMeanAndVariance, давали или принимали некоторые различия порядка 1e-7. Он также был в 2,5 раза медленнее , чем MPS на моей Ma c, вероятно, частично из-за неспособности использовать скрытие с гранулярным параллелизмом.