В CoreML с пользовательскими слоями есть ошибка на устройствах с Apple Neural Engine

Question

Похоже, что в CoreML с пользовательскими слоями есть ошибка на устройствах с Apple Neural Engine.

Симптомы ошибки: На устройствах с ANE, таких как iPhone XS для пользовательских слоев, функция 'ouputShapes' вызывается перед 'setWeightData'. Как результат для пользовательских слоев, чья внешняя форма зависит от входных весовых данных, может произойти сбой. В то время как на старых устройствах, таких как iPad Air 2, все работает нормально. В обычном случае функция setWeightData должна вызываться перед ouputShapes.

Похожие темы: https://forums.developer.apple.com/thread/113861

Answer 1

По совету Matthijs Hollemans, есть и другое возможное решение: если outShape зависит от небольших данных, мы можем хранить такие данные не в Weights, а в Parameters, которые передаются в 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)
    }
}

Answer 2

Решение - предотвратить запуск модели CoreML с таможенными слоями на ANE. Для этого используйте 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.