fix hard coded n_state in predictWith and rename variables

Author: wangchou

convert_whisper_encoder_to_coreml.py

@@ -12,13 +12,13 @@
 encoder = model.encoder
 encoder.eval()
 
-mel_segment = torch.ones((1, 80, 3000))
-traced_encoder = torch.jit.trace(encoder, mel_segment)
+melSegment = torch.ones((1, 80, 3000))
+traced_encoder = torch.jit.trace(encoder, melSegment)
 
 # convert to coreml model
 encoder = ct.convert(
     traced_encoder,
-    inputs=[ct.TensorType(name="mel_segment", shape=mel_segment.shape)],
+    inputs=[ct.TensorType(name="melSegment", shape=melSegment.shape)],
     outputs=[ct.TensorType(name="output")],
     compute_units=ct.ComputeUnit.ALL,
 )
@@ -27,11 +27,11 @@
 encoder_fp16.save(f"encoder_{modelSize}_fp16.mlmodel")
 
 # test accuracy
-torch_output = traced_encoder.forward(mel_segment)
+torch_output = traced_encoder.forward(melSegment)
 print("torch model output:", torch_output)
-mel_segment = mel_segment.cpu().detach().numpy()
+melSegment = melSegment.cpu().detach().numpy()
 coreml_output = torch.from_numpy(
-  list(encoder_fp16.predict({'mel_segment': mel_segment}).values())[0]
+  list(encoder_fp16.predict({'melSegment': melSegment}).values())[0]
 )
 print(f"coreml {modelSize} model output:", coreml_output)
 diff = torch.abs(torch_output - coreml_output).detach()

coreml/CoremlEncoder.h

@@ -16,10 +16,10 @@ NS_ASSUME_NONNULL_BEGIN
 API_AVAILABLE(macos(10.15), ios(13.0), watchos(6.0), tvos(13.0)) __attribute__((visibility("hidden")))
 @interface CoremlEncoderInput : NSObject<MLFeatureProvider>
 
-/// mel_segment as 1 × 80 × 3000 3-dimensional array of floats
-@property (readwrite, nonatomic, strong) MLMultiArray * mel_segment;
+/// melSegment as 1 × 80 × 3000 3-dimensional array of floats
+@property (readwrite, nonatomic, strong) MLMultiArray * melSegment;
 - (instancetype)init NS_UNAVAILABLE;
-- (instancetype)initWithMel_segment:(MLMultiArray *)mel_segment NS_DESIGNATED_INITIALIZER;
+- (instancetype)initWithMelSegment:(MLMultiArray *)melSegment NS_DESIGNATED_INITIALIZER;
 
 @end
 
@@ -123,11 +123,11 @@ API_AVAILABLE(macos(10.15), ios(13.0), watchos(6.0), tvos(13.0)) __attribute__((
 
 /**
     Make a prediction using the convenience interface
-    @param mel_segment as 1 × 80 × 3000 3-dimensional array of floats:
+    @param melSegment as 1 × 80 × 3000 3-dimensional array of floats:
     @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 CoremlEncoderOutput
 */
-- (nullable CoremlEncoderOutput *)predictionFromMel_segment:(MLMultiArray *)mel_segment error:(NSError * _Nullable __autoreleasing * _Nullable)error;
+- (nullable CoremlEncoderOutput *)predictionFromMelSegment:(MLMultiArray *)melSegment error:(NSError * _Nullable __autoreleasing * _Nullable)error;
 
 /**
     Batch prediction

coreml/CoremlEncoder.m

@@ -12,21 +12,21 @@
 
 @implementation CoremlEncoderInput
 
-- (instancetype)initWithMel_segment:(MLMultiArray *)mel_segment {
+- (instancetype)initWithMelSegment:(MLMultiArray *)melSegment {
     self = [super init];
     if (self) {
-        _mel_segment = mel_segment;
+        _melSegment = melSegment;
     }
     return self;
 }
 
 - (NSSet<NSString *> *)featureNames {
-    return [NSSet setWithArray:@[@"mel_segment"]];
+    return [NSSet setWithArray:@[@"melSegment"]];
 }
 
 - (nullable MLFeatureValue *)featureValueForName:(NSString *)featureName {
-    if ([featureName isEqualToString:@"mel_segment"]) {
-        return [MLFeatureValue featureValueWithMultiArray:self.mel_segment];
+    if ([featureName isEqualToString:@"melSegment"]) {
+        return [MLFeatureValue featureValueWithMultiArray:self.melSegment];
     }
     return nil;
 }
@@ -176,8 +176,8 @@ - (nullable CoremlEncoderOutput *)predictionFromFeatures:(CoremlEncoderInput *)i
     return [[CoremlEncoderOutput alloc] initWithOutput:(MLMultiArray *)[outFeatures featureValueForName:@"output"].multiArrayValue];
 }
 
-- (nullable CoremlEncoderOutput *)predictionFromMel_segment:(MLMultiArray *)mel_segment error:(NSError * _Nullable __autoreleasing * _Nullable)error {
-    CoremlEncoderInput *input_ = [[CoremlEncoderInput alloc] initWithMel_segment:mel_segment];
+- (nullable CoremlEncoderOutput *)predictionFromMelSegment:(MLMultiArray *)melSegment error:(NSError * _Nullable __autoreleasing * _Nullable)error {
+    CoremlEncoderInput *input_ = [[CoremlEncoderInput alloc] initWithMelSegment:melSegment];
     return [self predictionFromFeatures:input_ error:error];
 }
 

coreml/CoremlEncoder.mlmodel

@@ -59,7 +59,7 @@
         "formattedType" : "MultiArray (Float32 1 × 80 × 3000)",
         "shortDescription" : "",
         "shape" : "[1, 80, 3000]",
-        "name" : "mel_segment",
+        "name" : "melSegment",
         "type" : "MultiArray"
       }
     ],

coreml/CoremlEncoder.mlmodelc/analytics/coremldata.bin

@@ -13,7 +13,7 @@
   "layers" : [
     {
       "size_of_axes" : 1,
-      "bottom" : "mel_segment",
+      "bottom" : "melSegment",
       "axes_0" : -2,
       "weights" : {
 

coreml/CoremlEncoder.mlmodelc/coremldata.bin

@@ -284,13 +284,6 @@
       "_rank" : 4,
       "h" : 1
     },
-    "mel_segment" : {
-      "k" : 1,
-      "w" : 3000,
-      "n" : 1,
-      "_rank" : 3,
-      "h" : 80
-    },
     "input.17" : {
       "k" : 1,
       "w" : 1536,
@@ -720,6 +713,13 @@
       "_rank" : 4,
       "h" : 1500
     },
+    "melSegment" : {
+      "k" : 1,
+      "w" : 3000,
+      "n" : 1,
+      "_rank" : 3,
+      "h" : 80
+    },
     "q.3" : {
       "k" : 1500,
       "w" : 64,

coreml/CoremlEncoder.mlmodelc/metadata.json

@@ -7,15 +7,19 @@ int main() {
     const void* coremlEncoder = loadModel("coreml/CoremlEncoder.mlmodelc");
 
     // prepare easy test input like torch.ones
-    float* inFloats = (float *)malloc(sizeof(float) * 80 * 3000);
+    float* mel = (float *)malloc(sizeof(float) * 80 * 3000);
     for(int i=0;i<80*3000; i++) {
-        inFloats[i] = 1.0;
+        mel[i] = 1.0;
     }
 
-    float* outFloats = predictWith(coremlEncoder, inFloats);
+    // alloc output buffer
+    int n_state = 384; // tiny=384
+    float* outFloats = (float *)malloc(sizeof(float) * 1500 * n_state);
+
+    predictWith(coremlEncoder, mel, outFloats);
 
     // it should match
     // pytorch output: {'output': array([[[-0.28637695, -0.25561523, ..., -0.10253906]]], dtype=float32)
-    cout << outFloats[0] << " " << outFloats[1] << " " << outFloats[384*1500-1];
+    cout << outFloats[0] << " " << outFloats[1] << " " << outFloats[n_state*1500-1];
     closeModel(coremlEncoder);
 }