Реализация пользовательского слоя CoreML с двумя входами

Question

У меня есть график тензорного потока, который я хочу преобразовать в CoreML, но он использует некоторые пропущенные операции, которые мне придется реализовать как пользовательские слои.

Две операции, на которых я сейчас сосредотачиваюсь: Sin а также FloorDiv,

Sin было довольно просто, я мог следовать этому уроку, и у меня есть рабочий класс Swift и Metal ядро, которое выполняет работу, которую я протестировал с помощью игрушечного файла coreml:

import Foundation
import CoreML
import Accelerate

@objc(Sin) class Sin: NSObject, MLCustomLayer {

    let sinPipeline: MTLComputePipelineState

    required init(parameters: [String : Any]) throws {
        print(#function, parameters)

        let sinFunction = GPUDispatch.sharedInstance.library.makeFunction(name: "sin")!
        sinPipeline = try! GPUDispatch.sharedInstance.device.makeComputePipelineState(
            function: sinFunction)


        super.init()
    }

    func setWeightData(_ weights: [Data]) throws {
        print(#function, weights)
    }


    func outputShapes(forInputShapes inputShapes: [[NSNumber]]) throws
        -> [[NSNumber]] {
            print(#function, inputShapes)
            return inputShapes
    }

    func evaluate(inputs: [MLMultiArray], outputs: [MLMultiArray]) throws {

        for i in 0..<inputs.count {
            let input = inputs[i]
            let output = outputs[i]

            var count = Int32(input.count)
            let iptr = UnsafeMutablePointer<Float>(OpaquePointer(input.dataPointer))
            let optr = UnsafeMutablePointer<Float>(OpaquePointer(output.dataPointer))

            vvsinf(optr, iptr, &count)
        }

    }

    func encode(commandBuffer: MTLCommandBuffer,
                inputs: [MTLTexture], outputs: [MTLTexture]) throws {
        if let encoder = commandBuffer.makeComputeCommandEncoder() {
            for i in 0..<inputs.count {
                encoder.setTexture(inputs[i], index: 0)
                encoder.setTexture(outputs[i], index: 1)
                encoder.dispatch(pipeline: sinPipeline, texture: inputs[i])
                encoder.endEncoding()
            }
        }
    }

}

И в Sin.metal:

kernel void sin(
                  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 = sin(x);
    outTexture.write(half4(y), gid.xy, gid.z);
}

То, что я не понимаю, как это будет работать, если пользовательский слой имеет два входа, такие как мне нужно для FloorDiv, который возвращает floor(x / y),

Как бы я адаптировал Sin класс, который я предоставил, чтобы произвести что-то вроде sin(x*y)даже если это только на процессоре? Есть ли другие хорошие учебники для такого рода вещей?

Answer 1

Шаблон отличается от того, что я ожидал, но теперь он совершенно очевиден, я немного поиграл с кодом.

Это класс, который реализует FloorDiv:

import Foundation
import CoreML
import Accelerate

@objc(FloorDiv) class FloorDiv: NSObject, MLCustomLayer {

    let floorDivPipeline: MTLComputePipelineState

    required init(parameters: [String : Any]) throws {
        print(#function, parameters)

        let floorDivFunction = GPUDispatch.sharedInstance.library.makeFunction(name: "floordiv")!
        floorDivPipeline = try! GPUDispatch.sharedInstance.device.makeComputePipelineState(
            function: floorDivFunction)

        super.init()
    }

    func setWeightData(_ weights: [Data]) throws {
        print(#function, weights)
    }

    func outputShapes(forInputShapes inputShapes: [[NSNumber]]) throws
        -> [[NSNumber]] {
            print(#function, inputShapes)
            return inputShapes
    }

    func evaluate(inputs: [MLMultiArray], outputs: [MLMultiArray]) throws {

        let numerator = inputs[0]
        let denominator = inputs[1]
        var output = outputs[0]


        assert(numerator.count == denominator.count)

        var count = Int32(numerator.count)
        let numerator_ptr = UnsafeMutablePointer<Float>(OpaquePointer(numerator.dataPointer))
        let denominator_ptr = UnsafeMutablePointer<Float>(OpaquePointer(denominator.dataPointer))

        let output_ptr = UnsafeMutablePointer<Float>(OpaquePointer(output.dataPointer))

        vvdivf(output_ptr, numerator_ptr, denominator_ptr, &count)
        vvfloorf(output_ptr, output_ptr, &count)

    }


    func encode(commandBuffer: MTLCommandBuffer,
                inputs: [MTLTexture], outputs: [MTLTexture]) throws {

        if let encoder = commandBuffer.makeComputeCommandEncoder() {

                encoder.setTexture(inputs[0], index: 0)
                encoder.setTexture(inputs[1], index: 1)
                encoder.setTexture(outputs[0], index: 2)

                encoder.dispatch(pipeline: floorDivPipeline, texture: inputs[0])
                encoder.endEncoding()

        }
    }

}

А вот ядро ​​Металл:

#include <metal_stdlib>
using namespace metal;

kernel void floordiv(
                 texture2d_array<half, access::read> inTexture [[texture(0)]],
                 texture2d_array<half, access::read> inTexture2 [[texture(1)]],
                 texture2d_array<half, access::write> outTexture [[texture(2)]],
                 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 x2 = float4(inTexture2.read(gid.xy, gid.z));
    const float4 y = floor(x / x2);
    outTexture.write(half4(y), gid.xy, gid.z);
}