coreML单文件部署多个模型

image.png
image.png

如上图几个模型会对应生成相应的m文件，采用下面的方法可以只用一个m文件来加载多个模型

#import <Foundation/Foundation.h>

#import <CoreML/CoreML.h>
#import <stdint.h>

NS_ASSUME_NONNULL_BEGIN

/// Model Prediction Input Type
API_AVAILABLE(macos(10.13), ios(11.0), watchos(4.0), tvos(11.0))
@interface MLCoreModelInput : NSObject<MLFeatureProvider>

//the input name,default is image
@property (nonatomic, strong) NSString *inputName;

//data as color (kCVPixelFormatType_32BGRA) image buffer, 224 pixels wide by 224 pixels high
@property (readwrite, nonatomic) CVPixelBufferRef data;

- (instancetype)init NS_UNAVAILABLE;

- (instancetype)initWithData:(CVPixelBufferRef)data;

@end

API_AVAILABLE(macos(10.13), ios(11.0), watchos(4.0), tvos(11.0))
@interface MLCoreModelOutput : NSObject<MLFeatureProvider>

//the output name, defalut is prob
@property (nonatomic, strong) NSString *outputName;

// prob as multidimensional array of doubles
@property (readwrite, nonatomic) MLMultiArray *prob;

- (instancetype)init NS_UNAVAILABLE;
- (instancetype)initWithProb:(MLMultiArray *)prob;
@end

/// Model Prediction Output Type
API_AVAILABLE(macos(10.13), ios(11.0), watchos(4.0), tvos(11.0))
@interface MLCoreModelMapOutput : NSObject<MLFeatureProvider>

//the output value name, defalut is prob
@property (nonatomic, strong) NSString *outputValueName;
//the output label name, defalut is classLabel
@property (nonatomic, strong) NSString *outputLabelName;

/// prob as dictionary of strings to doubles
@property (readwrite, nonatomic) NSDictionary<NSString *, NSNumber *> * prob;

/// classLabel as string value
@property (readwrite, nonatomic) NSString * classLabel;
- (instancetype)init NS_UNAVAILABLE;
- (instancetype)initWithProb:(NSDictionary<NSString *, NSNumber *> *)prob classLabel:(NSString *)classLabel;
@end


// Class for model loading and prediction
API_AVAILABLE(macos(10.13), ios(11.0), watchos(4.0), tvos(11.0))
@interface MLCoreModel : NSObject

@property (readonly, nonatomic, nullable) MLModel * model;

//the input name,default is image
@property (nonatomic, strong) NSString *inputNodeName;
//the output value name, defalut is prob
@property (nonatomic, strong) NSString *outputValueName;
//the output label name, defalut is classLabel
@property (nonatomic, strong) NSString *outputLabelName;

- (nullable instancetype)initWithContentsOfURL:(NSURL *)url error:(NSError * _Nullable * _Nullable)error;

/**
 Make a prediction using the standard interface
 @param input an instance of ResnetNSFWInput to predict from
 @param error If an error occurs, upon return contains an NSError object that describes the problem. If you are not interested in possible errors, pass in NULL.
 @return the prediction as ResnetNSFWOutput
 */
- (nullable MLCoreModelOutput *)predictionFromFeatures:(MLCoreModelInput *)input error:(NSError * _Nullable * _Nullable)error;

/**
 Make a prediction using the standard interface
 @param input an instance of ResnetNSFWInput to predict from
 @param error If an error occurs, upon return contains an NSError object that describes the problem. If you are not interested in possible errors, pass in NULL.
 @return the prediction as MLCoreModelMapOutput
 */
- (nullable MLCoreModelMapOutput *)predictionMapFromFeatures:(MLCoreModelInput *)input error:(NSError * _Nullable * _Nullable)error;
/// All models can predict on a specific set of input features.
- (nullable id<MLFeatureProvider>)prediction:(MLCoreModelInput *)input
                                                   error:(NSError **)error;
/**
 Make a prediction using the convenience interface
 @param data as color (kCVPixelFormatType_32BGRA) image buffer, 224 pixels wide by 224 pixels high:
 @param error If an error occurs, upon return contains an NSError object that describes the problem. If you are not interested in possible errors, pass in NULL.
 @return the prediction as ResnetNSFWOutput
 */
- (nullable MLCoreModelOutput *)predictionFromData:(CVPixelBufferRef)data error:(NSError * _Nullable * _Nullable)error;

/**
 Make a prediction using the convenience interface
 @param data as color (kCVPixelFormatType_32BGRA) image buffer, 224 pixels wide by 224 pixels high:
 @param error If an error occurs, upon return contains an NSError object that describes the problem. If you are not interested in possible errors, pass in NULL.
 @return the prediction as MLCoreModelMapOutput
 */
- (nullable MLCoreModelMapOutput *)predictionMapFromData:(CVPixelBufferRef)data error:(NSError * _Nullable * _Nullable)error;

@end

NS_ASSUME_NONNULL_END

#import "MLCoreModel.h"

#define DefalutInputName            @"image"
#define DefalutOutputValueName      @"prob"
#define DefalutOutputLabelName      @"classLabel"

@implementation MLCoreModelInput

- (instancetype)initWithData:(CVPixelBufferRef)data {
    if (self) {
        _data = data;
        _inputName = DefalutInputName;
    }
    return self;
}

- (NSSet<NSString *> *)featureNames {
    return [NSSet setWithArray:@[self.inputName]];
}

- (nullable MLFeatureValue *)featureValueForName:(nonnull NSString *)featureName {
    if ([featureName isEqualToString:self.inputName]) {
        return [MLFeatureValue featureValueWithPixelBuffer:_data];
    }
    
    return nil;
}

@end

@implementation MLCoreModelOutput

- (instancetype)initWithProb:(MLMultiArray *)prob{
    if (self) {
        _prob = prob;
        _outputName = DefalutOutputValueName;
    }
    return self;
}

- (NSSet<NSString *> *)featureNames{
    return [NSSet setWithArray:@[self.outputName]];
}

- (nullable MLFeatureValue *)featureValueForName:(nonnull NSString *)featureName {
    if ([featureName isEqualToString:self.outputName]) {
        return [MLFeatureValue featureValueWithMultiArray:_prob];
    }
    
    return nil;
}

@end

@implementation MLCoreModelMapOutput

- (instancetype)initWithProb:(NSDictionary<NSString *, NSNumber *> *)prob classLabel:(NSString *)classLabel {
    if (self) {
        _prob = prob;
        _classLabel = classLabel;
        _outputValueName = DefalutOutputValueName;
        _outputLabelName = DefalutOutputLabelName;
    }
    return self;
}

- (NSSet<NSString *> *)featureNames {
    return [NSSet setWithArray:@[self.outputValueName, self.outputLabelName]];
}

- (nullable MLFeatureValue *)featureValueForName:(NSString *)featureName {
    if ([featureName isEqualToString:self.outputValueName]) {
        return [MLFeatureValue featureValueWithDictionary:_prob error:nil];
    }
    if ([featureName isEqualToString:self.outputLabelName]) {
        return [MLFeatureValue featureValueWithString:_classLabel];
    }
    return nil;
}

@end

@implementation MLCoreModel

- (nullable instancetype)initWithContentsOfURL:(NSURL *)url error:(NSError * _Nullable * _Nullable)error{
    self = [super init];
    if (!self) { return nil; }
    
    _model = [MLModel modelWithContentsOfURL:url error:error];
    if (_model == nil) {
        return nil;
    }
    
    _outputValueName = DefalutOutputValueName;
    _outputLabelName = DefalutOutputLabelName;
    _inputNodeName = DefalutInputName;
    return self;
}

- (nullable MLCoreModelOutput *)predictionFromFeatures:(MLCoreModelInput *)input error:(NSError * _Nullable * _Nullable)error{
    id<MLFeatureProvider> outFeatures = [_model predictionFromFeatures:input error:error];
    MLCoreModelOutput * result = [[MLCoreModelOutput alloc] initWithProb:[outFeatures featureValueForName:self.outputValueName].multiArrayValue];
    return result;
}
- (nullable id<MLFeatureProvider>)prediction:(MLCoreModelInput *)input
                                       error:(NSError **)error
{
   id<MLFeatureProvider> outFeatures = [_model predictionFromFeatures:input error:error];
    return outFeatures;
}
- (nullable MLCoreModelMapOutput *)predictionMapFromFeatures:(MLCoreModelInput *)input error:(NSError * _Nullable * _Nullable)error{
    id<MLFeatureProvider> outFeatures = [_model predictionFromFeatures:input error:error];
    MLCoreModelMapOutput * result = [[MLCoreModelMapOutput alloc] initWithProb:(NSDictionary<NSString *, NSNumber *> *)[outFeatures featureValueForName:self.outputValueName].dictionaryValue classLabel:[outFeatures featureValueForName:self.outputLabelName].stringValue];
    return result;
}


- (nullable MLCoreModelOutput *)predictionFromData:(CVPixelBufferRef)data error:(NSError * _Nullable * _Nullable)error{
    MLCoreModelInput *input_ = [[MLCoreModelInput alloc] initWithData:data];
    input_.inputName = self.inputNodeName;
    return [self predictionFromFeatures:input_ error:error];
}

- (nullable MLCoreModelMapOutput *)predictionMapFromData:(CVPixelBufferRef)data error:(NSError * _Nullable * _Nullable)error{
    MLCoreModelInput *input_ = [[MLCoreModelInput alloc] initWithData:data];
    input_.inputName = self.inputNodeName;
    return [self predictionMapFromFeatures:input_ error:error];
}

@end

使用模型的时候设置一下inputNodeName，outputValueName

        self.coreModel = [[MLCoreModel alloc] initWithContentsOfURL:[NSURL URLWithString:saveURL] error:nil];
        if (self.outputLabelsName)  self.coreModel.outputLabelName = self.outputLabelsName;
        if (self.outputValueName)   self.coreModel.outputValueName = self.outputValueName;

模型具有多个输出示例如下：

    NSError *error;
    _model_pnet =[[MLCoreModel alloc]initWithContentsOfURL:[NSURL URLWithString:saveURL] error:&error];
    NSError *error;
    MLCoreModelInput *input = [[MLCoreModelInput alloc] initWithData:pixelBuffer];
    input.inputName = @"data";
    id<MLFeatureProvider> outFeatures = [_model_pnet prediction:input error:&error];
    MLMultiArray *conv4_2 = [outFeatures featureValueForName:@"conv4-2"].multiArrayValue;
    MLMultiArray * prob1 = [outFeatures featureValueForName:@"prob1"].multiArrayValue;