我们在工作或学习中，往往并不能够直接得到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 git@github.com: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 git@github.com: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_info.txt>
其中 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 <model.pb>
你会在同级目录下，得到一个mlmodel.mlmodel的Core ML文件。
有时候，执行可能会报错：
NotImplementedError: Unsupported Ops of type: Shape,Pack
不用理会，再跑一次就可以了。
六、部署到iOS
你可以参考《二、使用Core ML加载.mlmodel模型文件》来完成模型的部署，你可以直接从我的GitHub获取一个已经构建好的项目半成品：
git clone git@github.com: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;
}
运行之后，点击首页的“手写数字识别”，在黑色区域内书写数字，就可以在下方看到分类结果了。
为了缩短本文中脚本的运行时间，这个模型使用的层数较少，同时也没有使用非常大量的训练集，所以并不能很好的识别手写数字，你大概需要写的规范一点、把数字写满黑色区域的中间位置才能得到比较好的结果。
七、成果
最后，还是来看几张成果图吧：
网友评论