将其他模型文件转化成Core ML模型文件（PB）

我们在工作或学习中，往往并不能够直接得到Core ML模型。这就需要将其他模型转化成Core ML。这篇文章将会介绍如何借助一个名为tf-coreml的开源项目，将一个手写字识别的PB模型，转化为Core ML的mlmodel模型。并将其部署到我们的iOS项目中。

本文依赖：

Python
tensorflow >= 1.5.0
coremltools >= 0.8
numpy >= 1.6.2
protobuf >= 3.1.0
six >= 1.10.0

注意：
1.阅读本文，你最好亲自动手使用TensorFlow生成并训练过一个神经网络模型，并使用模型做过预测。当然，这样做的目的，是为了让你自己对TensorFlow和神经网络模型有一个大致的概念，因此你大可不必关心你的模型的准确率。

2.由于本文从网络上获取了一些Python代码，他们有些是Python2，有些是Python3，因此本文将严格区分Python2和Python3的命令，包括pip2和pip3。不会出现类似：Python file.py 的命令。你需要根据自己的Python环境，自行决定如何调用这些Python脚本。

一、准备一个PB文件：

首先，我们需要准备一个PB格式的模型文件。如果你手上没有PB文件，你可以使用我已经训练好的模型文件：

首先，你需要Clone我的项目：

 

git clone [email protected]:yangchenlarkin/CoreML.git

在files目录中找到pb_model.pb文件

 

cd CoreML/files/
ls pb_model.pb

二、安装tf-coreml

首先你需要安装好文章开头描述的相关依赖项。然后安装tf-coreml

 

pip3 install -U tfcoreml

三、观察tfcoreml.convert函数

我们需要使用tfcoreml.convert函数来进行转换，传入PB文件路径和一些参数，这个函数会帮助你生成一个.mlmodel文件存到你指定的路径：

 

def convert(tf_model_path,
            mlmodel_path,
            output_feature_names,
            input_name_shape_dict=None,
            image_input_names=None,
            is_bgr=False,
            red_bias=0.0,
            green_bias=0.0,
            blue_bias=0.0,
            gray_bias=0.0,
            image_scale=1.0,
            class_labels=None,
            predicted_feature_name=None,
            predicted_probabilities_output='',
            add_custom_layers=False,  # type: bool
            custom_conversion_functions={},  # type: Dict[Text, Any]
            )

去掉我们不需要关注的字段，剩下的，我们一一解释一下。然后我会给出一些手段，帮助你了解如何获得这些字段的数据。

 

def convert(tf_model_path,
            mlmodel_path,
            output_feature_names,
            input_name_shape_dict=None,
            image_input_names=None,
            class_labels=None,
            )

• tf_model_path：
  • PB文件路径，即我们第一步中准备的文件
• mlmodel_path：
  • 我们最终需要的Core ML文件的存储路径
• output_feature_names：
  • 需要转化的输出层名字列表。
  • TensorFlow中，一个模型可能有好几个输出，每个输出都会有一个名字，这里需要定义我们需要将哪些输出显示的暴露出来。我们在上一篇文章中曾经介绍到，我们将CoreML模型文件，在Xcode中打开后，可以看到他的输出，这个输出，就是在这里定义的。
• input_name_shape_dict：
  • 输入层的shape字典。
  • TensorFlow中，一个模型也会有好几个输入，每个输入都会有一个名字，输入的shape可能不尽相同，因此需要在这里进行描述。
  • 对于图片类型的输入而言，一般都会是形如[?, sizeW, sizeH, n]的shape。解释一下：第一个'？'是因为模型支持批量传入数据集，而数据集的大小（有几条数据，或者说一批传入的图片数）是可变的，因此这里是一个问号。sizeW和sizeH表示图片宽高；n表示图片通道数，比如说，1表示灰度图，每个像素只有灰度，3表示彩色，可能是RGB，也可能是其他，4表示彩色加一个透明度通道。这个需要和构建这个神经网络模型的工程师沟通一下。
  • 另外，由于Core ML只用于单个样本的预测，不会批量传样本到模型，因此在转化模型的时候，这里需要把'?'修改成1；
• image_input_names：
  • 用于描述哪些输入是图片格式的，这个列表应该是input_name_shape_dict所有key的子集，
  • 如果这里不加描述，那么生成的Core ML模型将会认为这些输入只是简单的多维数组。加了这个描述之后，Core ML会将其描述为图片，这样，你才可以在Xcode中，以CVPixelBufferRef对象作为输入。
• class_labels：
  我们这里是一个图片分类模型，这里可以设置一个字符串数组，用于对应输出层的每一个节点，本文中因为是手写字体识别，所以这里恰好应该是:

 

[str(x) for x in range(10)]

四、获取信息

由上一小节我们可以知道，我们主要需要获取的信息有：

 

input_name_shape_dict = None
image_input_names = None
output_feature_names = None
class_labels = None

1.利用tf-coreml内置的工具打印出PB文件的信息

首先我们将tf-coreml项目clone下来

 

git clone [email protected]:tf-coreml/tf-coreml.git

在utils文件加下找到一个名叫inspect_pb.py的文件：

 

cd tf-coreml/utils/
ls inspect_pb.py

这个脚本，需要两个参数：

参数1：PB模型文件路径
参数2：模型信息输出的文件路径

 

python3 inspect_pb.py  

其中 model.pb就是第一步中获得的pb模型，pb_info.txt，你可以按照自己的喜好随意指定，比如:

 

~/Desktop/pb_info.txt

执行完毕后让我们打开这个文件，我们看到很多内容，你可以和创建这个神经网络的工程师沟通，了解一下他的输入和输入到底是哪个。在本文中，输入和输出分别是：

 

输入

输出

这里输入和输出的最后一行就是name和shape，值得注意的是：

• 首先，我们不要前面的import/
• 第二，如上文中所描述的，这里我们需要把输入的'?'改成1
  因此，我们可以得到参数：

 

input_name_shape_dict={'conv2d_1_input:0': [1, 28, 28, 1]},
image_input_names='conv2d_1_input:0',
output_feature_names=['dense_2/Softmax:0'],
class_labels=[str(x) for x in range(10)],

意思是分别是：

• 输入'conv2d_1_input:0'的shape是[1, 28, 28, 1]
• 输入'conv2d_1_input:0'是一个图片
• 输出是'dense_2/Softmax:0'
• 输出的标签是['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

五、编写代码

以上准备工作都做完之后，代码其实很简单，创建一个名为convert.py的python脚本：

 

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# file name: convert.py

import tfcoreml
import sys


def main():
    if len(sys.argv) == 1:
        print('please input the file path of a pb model')
        return

    tfcoreml.convert(tf_model_path=sys.argv[1],
                     mlmodel_path='mlmodel.mlmodel',
                     output_feature_names=['dense_2/Softmax:0'],
                     input_name_shape_dict={'conv2d_1_input:0': [1, 28, 28, 1]},
                     image_input_names='conv2d_1_input:0',
                     class_labels=[str(x) for x in range(10)])


if __name__ == '__main__':
    main()

然后直接执行就OK了：

 

python3 convert.py 

你会在同级目录下，得到一个mlmodel.mlmodel的Core ML文件。

有时候，执行可能会报错：

 

NotImplementedError: Unsupported Ops of type: Shape,Pack

不用理会，再跑一次就可以了。

六、部署到iOS

你可以参考《二、使用Core ML加载.mlmodel模型文件》来完成模型的部署，你可以直接从我的GitHub获取一个已经构建好的项目半成品：

 

git clone [email protected]:yangchenlarkin/CoreML.git

运行项目，点击首页的“手写数字识别”，可以进入到识别界面。该ViewController位于项目中的DigitRecognition文件夹。打开DRViewController.m，你需要完成最后的方法：

 

- (NSString *)predict:(UIImage *)image {
    return @"TODO";
}

而如下两个方法已经为你写好了：

 

#pragma mark - predict

- (CVPixelBufferRef)pixelBufferFromCGImage:(CGImageRef)image {
    //...
}
- (UIImage*)scaleImage:(UIImage *)image size:(CGFloat)size {
    //...
}

七、分析与coding

在动手之前，我们看一下这个方法：

 

#pragma mark - predict

- (CVPixelBufferRef)pixelBufferFromCGImage:(CGImageRef)image {
    //...
    CVReturn status = CVPixelBufferCreate(kCFAllocatorDefault,
                                          frameWidth,
                                          frameHeight,
                                          kCVPixelFormatType_OneComponent8,
                                          (__bridge CFDictionaryRef) options,
                                          &pxbuffer);
    //...
    CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceGray();
    CGContextRef context = CGBitmapContextCreate(pxdata,
                                                 frameWidth,
                                                 frameHeight,
                                                 8,
                                                 CVPixelBufferGetBytesPerRow(pxbuffer),
                                                 colorSpace,
                                                 (CGBitmapInfo)kCGImageAlphaNone);
    //...
}

对比一下《二、使用Core ML加载.mlmodel模型文件》中的同名方法：

 

#pragma mark - predict

- (CVPixelBufferRef)pixelBufferFromCGImage:(CGImageRef)image {
    //...
    CVReturn status = CVPixelBufferCreate(kCFAllocatorDefault,
                                          frameWidth,
                                          frameHeight,
                                          kCVPixelFormatType_32ARGB,
                                          (__bridge CFDictionaryRef) options,
                                          &pxbuffer);
    //...
    CGColorSpaceRef rgbColorSpace = CGColorSpaceCreateDeviceRGB();
    CGContextRef context = CGBitmapContextCreate(pxdata,
                                                 frameWidth,
                                                 frameHeight,
                                                 8,
                                                 CVPixelBufferGetBytesPerRow(pxbuffer),
                                                 rgbColorSpace,
                                                (CGBitmapInfo)kCGImageAlphaNoneSkipFirst);
    //...
}

我们可以发现，我们在两端代码中使用的图片格式是不同的，原因是，《二、使用Core ML加载.mlmodel模型文件》中我们使用的是彩色图片，模型输入的shape是[1, size, size, 3]，而这里我们使用的是灰度图，模型输入的shape是[1, size, size, 1]。这是值得注意的一个地方。

下面直接给出实现代码：

 

- (NSString *)predict:(UIImage *)image {
    UIImage *scaleImage = [self scaleImage:image size:28];
    CVPixelBufferRef buffer = [self pixelBufferFromCGImage:scaleImage.CGImage];
    mlmodel *m = [[mlmodel alloc] init];
    NSError *error = nil;
    mlmodelOutput *o = [m predictionFromConv2d_1_input__0:buffer error:&error];
    if (error) {
        NSLog(@"%@", error);
        return nil;
    }
    return o.classLabel;
}

运行之后，点击首页的“手写数字识别”，在黑色区域内书写数字，就可以在下方看到分类结果了。

为了缩短本文中脚本的运行时间，这个模型使用的层数较少，同时也没有使用非常大量的训练集，所以并不能很好的识别手写数字，你大概需要写的规范一点、把数字写满黑色区域的中间位置才能得到比较好的结果。

七、成果

最后，还是来看几张成果图吧：

 

链接：https://www.jianshu.com/p/0b2ef1d76b68