CODEWITHSUNDEEP
coremlcoremltools
Jarikus
Jarikus

Reputation: 834

CoreML with Custom Layers have bug on devices with Apple Neural Engine

Looks like CoreML with Custom Layers have bug on devices with Apple Neural Engine.

Bug symptoms: On devices with ANE like iPhone XS for custom layers the function 'ouputShapes' call before 'setWeightData'. As results for custom layers which of out shape depend from input weights data may crash. While on older devices like iPad Air 2 all works fine. In normal the function 'setWeightData' must call before 'ouputShapes'.

Related discussion: https://forums.developer.apple.com/thread/113861

Upvotes: 0

Views: 866

Answers (2)

Jarikus
Jarikus

Reputation: 834

By advise Matthijs Hollemans, there is another possible solution: if outShape depended from small data we may store such data not to Weights while to Parameters that gets passed into the init of the custom layer.

Python side (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 app side:

@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)
    }
}

Upvotes: 1

Jarikus
Jarikus

Reputation: 834

Solution is prevent run CoreML model with customs layers on ANE. For this use https://developer.apple.com/documentation/coreml/mlcomputeunits

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

But if you have big model you may use black magic of CoreML for use ANE. Need divide your model on two CoreML parts where one model not have custom layers and may run on ANE while other part may run on CPU or GPU. And connect output of first model to input of second in SWIFT app side.

Sample, I have model which generate caption for image. And it consist from two parts: image features extractor and caption generator.

For convert this model to CoreML, the caption generator need some custom layers so I divided model on two CoreML parts:

// 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") }

As results heavy features model run on ANE which may speed up to 10x. While small model run on CPU or GPU.

Upvotes: 1

Related Questions