Search code examples
swiftlstmcoreml

swift coreML: prediction function without "options" parameter


In the swift documentation, MLModel have two prediction functions

  1. func prediction(from: MLFeatureProvider) -> MLFeatureProvider. Predicts output feature values from the given input feature values.
  2. func prediction(from: MLFeatureProvider, options: MLPredictionOptions) -> MLFeatureProvider. Predicts output feature values from the given input feature values.

However, in my auto-generated MLModel class, the function with options parameter was not been generated. The following code are my auto-generated predict functions.

func prediction(input: coreML_1denses_80iters_k213_2Input) throws -> coreML_1denses_80iters_k213_2Output {
    let outFeatures = try model.prediction(from: input)
    let result = coreML_1denses_80iters_k213_2Output(output1: outFeatures.featureValue(for: "output1")!.multiArrayValue!, lstm_1_h_out: outFeatures.featureValue(for: "lstm_1_h_out")!.multiArrayValue!, lstm_1_c_out: outFeatures.featureValue(for: "lstm_1_c_out")!.multiArrayValue!)
    return result
}

func prediction(input1: MLMultiArray, input2: MLMultiArray, lstm_1_h_in: MLMultiArray?, lstm_1_c_in: MLMultiArray?) throws -> coreML_1denses_80iters_k213_2Output {
    let input_ = coreML_1denses_80iters_k213_2Input(input1: input1, input2: input2, lstm_1_h_in: lstm_1_h_in, lstm_1_c_in: lstm_1_c_in)
    return try self.prediction(input: input_)
}

NOTE: btw, the reason why I want to find the prediction function with "options" parameter is this error message:

[coreml] Cannot evaluate a sequence of length 600, which is longer than maximum of 400.

And I have found a solution, which is adding a forceCPU Flag in prediction function. The option can be found in MLPredictionOptions called "usesCPUOnly". However, I cannot find a place to put the options.


Solution

  • One way to do this is to add your own prediction method in an extension of the auto-generated class (in a different source file).