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

Posted 2023-03-27

【中文标题】带有自定义层的 CoreML 在带有 Apple 神经引擎的设备上存在错误

【英文标题】：CoreML with Custom Layers have bug on devices with Apple Neural Engine

【发布时间】：2019-10-05 06:35:57

【问题描述】：

看起来。

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

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

【问题讨论】：

有趣的是，我根本无法让具有自定义层的模型在 ANE 上运行。您的自定义层是在神经网络的最末端，还是在中间的某个地方（即它后面是否还有其他层）？

@MatthijsHollemans 我有很多自定义层，从像 TensorFlow Gather 函数这样的小层到做一些后处理的绝对技术自定义层，自定义层在模型的开头和结尾。我描述的错误发生在模型初始化期间，我的应用程序在启动期间崩溃，我不知道它是否能够推断出 ANE 上的某些层。

你可以通过不使输出形状依赖于 weights 而是依赖于“参数”来解决这个问题，这是一个传递到自定义初始化的字典层。但我不确定这对您的自定义图层是否有意义。

@MatthijsHollemans hm 有趣的想法，考虑到矩阵只有二维因此我们可以将矩阵形状存储在两个整数参数中，我在几秒钟内发布可能的示例解决方案。

【参考方案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 上运行。

【参考方案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)