带有自定义层的 CoreML 在带有 Apple 神经引擎的设备上存在错误

Question

看起来带有自定义层的 CoreML 在带有 Apple 神经引擎的设备上存在错误。

错误症状：在具有 ANE（如 iPhone XS）的设备上用于自定义层，函数“ouputShapes”在“setWeightData”之前调用。由于形状取决于输入权重数据的自定义层的结果可能会崩溃。在 iPad Air 2 等旧设备上一切正常。 正常情况下，函数 'setWeightData' 必须在 'ouputShapes' 之前调用。

相关讨论：https://forums.developer.apple.com/thread/113861

Answer 1

解决方案是阻止在 ANE 上运行带有自定义层的 CoreML 模型。为此使用https://developer.apple.com/documentation/coreml/mlcomputeunits

let config = MLModelConfiguration()
config.computeUnits = MLComputeUnits.cpuAndGPU

但是如果你有大模型，你可以使用 CoreML 的黑魔法来使用 ANE。 需要将您的模型划分为两个 CoreML 部分，其中一个模型没有自定义层并且可以在 ANE 上运行，而其他部分可以在 CPU 或 GPU 上运行。并将第一个模型的输出连接到 SWIFT 应用端的第二个模型的输入。

示例，我有为图像生成标题的模型。它由两部分组成：图像特征提取器和标题生成器。

为了将此模型转换为 CoreML，字幕生成器需要一些自定义层，因此我将模型分为两个 CoreML 部分：

// Standart initialize CoreML model
let model_features = CaptionMobile_features()

// Initialize CoreML model with options
// Prevent run model on ANE but alow run on CPU and GPU
let config = MLModelConfiguration()
config.computeUnits = MLComputeUnits.cpuAndGPU
guard let model_caption = try? CaptionMobile_caption(configuration: config)
else { fatalError("Can't intitalize Caption CoreML model") }

结果重特征模型在 ANE 上运行，速度可能高达 10 倍。而小型模型在 CPU 或 GPU 上运行。

Answer 2

根据 Matthijs Hollemans 的建议，还有另一种可能的解决方案：如果 outShape 依赖于小数据，我们可能会将此类数据存储到权重而不是传递到自定义层的 init 的参数中。

Python 端（Coremltools）：

# MatMul - matrix multiplication of two matrix A * B
def  _convert_matmul(**kwargs):
    tf_op = kwargs["op"]
    coreml_nn_builder = kwargs["nn_builder"]
    constant_inputs = kwargs["constant_inputs"]

    params = NeuralNetwork_pb2.CustomLayerParams()
    # The name of the Swift or Obj-C class that implements this layer.
    params.className = 'MatMul'
    params.description = 'Custom layer that corresponds to the MatMul TF op'

    # Get specific parameters (constant inputs) of operation
    ################
    # Store matrix (B) as weight parameter, in weights by index [0]
    w = constant_inputs.values()[0]

    ########
    # We Store B matrix shape for ability calculate out results matrix shape during matrix multiplication in Swift code,
    # Store shape of B matrix, in parameters which passed to function init in SWIFT app side
    params.parameters["b_shape_0"].intValue = w.shape[0]
    params.parameters["b_shape_1"].intValue = w.shape[1]
    ########

    # Store constant input as weights because this array/matrix
    w_as_weights = params.weights.add()
    # Suppoerted types for WeightParams see in:
    # https://github.com/apple/coremltools/blob/5bcd4f8aa55df82792deb9a4491c3f37d5b89149/mlmodel/format/NeuralNetwork.proto#L658
    w_as_weights.floatValue.extend(map(float, w.flatten()))
    ################

    # This operation receive first input (A) as standard tensor and second input as const which we resend via 'weights', see above
    input_names = [tf_op.inputs[0].name]
    # Get list of out tensors names
    output_names = [out_tensor.name for out_tensor in tf_op.outputs]

    coreml_nn_builder.add_custom(name=tf_op.name,
                                    input_names=input_names,
                                    output_names=output_names,
                                    custom_proto_spec=params)

SWIFT 应用端：

@objc(MatMul) class MatMul: NSObject, MLCustomLayer {
    private var b_shape = [Int]()

    required init(parameters: [String : Any]) throws {
        //print(String(describing: self), #function, parameters)
        super.init()

        // Parameters came from _convert_matmul() 
        b_shape.append(parameters["b_shape_0"] as? Int ?? 0)
        b_shape.append(parameters["b_shape_1"] as? Int ?? 0)
    }
}