美文网首页iOS学习开发iOS学习笔记
如何把tensorflow的pb模型转换为CoreML的mlmo

如何把tensorflow的pb模型转换为CoreML的mlmo

作者: QYiZHong | 来源:发表于2018-04-23 11:19 被阅读1283次

    前言

    鉴于发布这篇博客以来已经有不少人私信问我如何去转换,我想可能是我的博客并没有写明白,于是我决定把这篇文章完善一下。

    私信我的人有下面这两种情况:

    1. 只是自己写iOS,想尝试如何在iOS上跑神经网络模型,偏向于玩的性质。
    2. 确实有需求需要转换模型,偏向于工作的性质。

    如果你是第一种的话,我推荐直接使用iOS12的CreateML或者我的另一篇文章Apple机器学习库turicreate实战——使用cifar-10数据集做图像分类。如果你是第二种的话,你可能需要与交付你模型的人沟通一下。因为涉及到模型的输入输出的shape,如果你对tensorflow了解不多的话这个转化过程是十分不友好的。

    我的需求

    最近有需求需要把tensorflow训练的模型在iOS上使用,然后我在GitHub上发现了一个叫tf-coreml的库,他可以把pb模型转化为mlmodel模型。

    转换

    获得模型

    你的目标一定是通过训练来保存模型,最后放在xcode里可以用来调用。使用tensorflow训练保存的模型有两种格式:

    1. checkpoint格式
    2. pb格式

    第一种格式需要你通过代码把checkpont固化为pb模型,但是这一步的坑特别多,我尝试之后觉得不太友好。第二种是训练完用tensorflow直接保存为pb格式的模型,就已经满足转化的条件了。

    转换前需要了解的知识

    如果你对tensorflow并不了解,我觉得在看下面的转换过程之前,你需要了解一些基本的知识以便你能顺利的转换。先放上我的demo方便讲解。

    一般来说,转换的过程中你需要关注输入和输出,我随意举例:

    input = tf.placeholder(tf.float32, shape=[None, 36], name="Input")
    # 可能是这样,我记不太清了,输出可以是千奇百怪的
    output = tf.layers.dense(inputs=self.layer2, units=1, activation=tf.nn.relu, name="Prediction")
    

    如果这个神经网络模型并不是你一手操办的,那你肯定是不知道你的输入输出到底是啥了。这个时候你打开我demo里的一个叫network-info.txt的文件,里面有两段内容:

    ---------------------------------------------------------------------------------------------------------------------------------------------
    0: op name = import/Input, op type = ( Placeholder ), inputs = , outputs = import/Input:0
    @input shapes:
    @output shapes:
    name = import/Input:0 : (?, 36)
    ---------------------------------------------------------------------------------------------------------------------------------------------
    
    ---------------------------------------------------------------------------------------------------------------------------------------------
    19: op name = import/Prediction, op type = ( Relu ), inputs = import/add_5:0, outputs = import/Prediction:0
    @input shapes:
    name = import/add_5:0 : (?, 1)
    @output shapes:
    name = import/Prediction:0 : (?, 1)
    ---------------------------------------------------------------------------------------------------------------------------------------------
    

    这个是根据我demo里的inspect_pb.py生成的,这两个就是输入输出的op。“outputs = import/Input:0",最后那个Input:0就是你需要关注的内容。你会发现他们俩刚好一个在最前面一个在最后面。这不是一定的,但经常是这样的。实际情况下你可能需要与交付你模型的同事沟通一下到底输入和输出是哪一个。

    转换过程

    下载inspect_pb.py文件

    进入tf-coreml的github,然后下载他们那个库里utils/下的一个inspect_pb.py文件,如图1、2。

    图1 图2

    把这个py改一下,里面的方法可以把pb模型图里的所有信息写在一个txt格式的文件里,如图3。

    图3

    你可以在图3显示的txt里找到你需要的输入和输出信息的全名,这里你需要找的就与上文所述的一样。转换代码如下:

    import tfcoreml
    tfcoreml.convert(tf_model_path="./model_0.pb",
                         mlmodel_path="./model.mlmodel",
                         output_feature_names=['Prediction:0'],
                         input_name_shape_dict={'Input:0': [1, 36]})
    

    你可能很纳闷为啥名字得是'Prediction:0','Input:0',这我也不知道,但是我知道tensorflow里面通过名字获得变量的时候也得这么写。

    参数解释一下,第一个参数是pb模型的路径,第二个参数是生成mlmodel模型的路径,第三个是输出的名字,这个输出的名字需要在上文所述生成的txt文件里有,且是你需要的,第四个是可选的,像我的Input的shape是[?, 36],因为我的输入是下面这样的

    input = tf.placeholder(tf.float32, shape=[None, 36], name="Input")
    

    这里要保证输入固定的,所以改为 [1, 36]。

    demo: GitHub

    相关文章

      网友评论

        本文标题:如何把tensorflow的pb模型转换为CoreML的mlmo

        本文链接:https://www.haomeiwen.com/subject/wvvklftx.html