tensorflow2ncnn模型转换源码分析

由于ncnn作者nihui大佬说对tensorflow不是特别熟,所以ncnn的github里已经没有tensorflow2ncnn.cpp了,但是现在在tools文件夹里又出现了tensorflow文件夹,里面给出了第三方转换的地址 https://github.com/hanzy88/tensorflow2ncnn 。由于我平时还是用tensorflow较多,所以还是又这方面的需求,所以还是要了解一点tensorflow2ncnn的方法,这里就对这个地址里的tensorflow2ncnn.cpp做下源码分析,这样可以自己改这个文件满足转换需求。

其实转换的原理和《深度学习模型移植和转换》文集下的《深度学习模型转换(tensorflow2caffe)》其中阐述的原理是一样的。都是把一种框架的模型解析出算子和层,然后以另一种框架模型的格式写进文件。ncnn以层为最小的模型结构,所以对于同样以层为最小模型结构的模型框架的转换会比较简单,比如caffe。

tensorflow数据结构

tensorflow有一些特殊,它是以op为最小的模型结构。为了接下来源码分析的方便,先介绍一些tensorflow几个基本的数据结构,具体可参考 https://www.jianshu.com/p/236335897b30

  • Graph(图)
    把运算任务描述成一个直接的无环图形(DAG),图表中的节点(node)代表必须要实现的一些操作。图中的边代表数据或者可控的依赖。GratheDef 是系统中描述一个图表的协议(api),它由一个 NodeDefs 集合组成。一个GraphDef可以转化成一个更容易操作的图表对象。
  • Node(节点)
    图中的一个元素。
    把启动一个特定操作的方式称为特定运算图表的一个节点,包括任何用来配置这个操作的属性的值。对于那些多形态的操作,这些属性包括能完全决定这个节点(Node)签名的充分信息。详见graph.proto。
  • 操作(Op/operation)
    在 TensorFlow 的运行时中,它是一种类似 add 或 matmul 或 concat的运算。可以用 how to add an op 中的方法来向运行时添加新的操作。
    在 Python 的API中,它是图中的一个节点。在 tf.Operation 类中列举出了这些操作。一个操作(Operation)的 type 属性决定这个节点(node)的操作类型,比如add和matmul。
  • Tensor
    Tensor是一种特定的多维数组。比如,一个浮点型的四维数组表示一小批由[batch,height,width,channel]组成的图片。
    在一个运行的图(graph)中,它是一种流动在节点(node)之间的数据。
    在 Python 中,Tensor 类表示添加到图的操作中的输入和输出,见 tf.Tensor ,这样的类不持有数据。
  • tensorflow2ncnn

    image.png
    https://github.com/hanzy88/tensorflow2ncnn 项目的tools/tensorflow下就是tensorflow2ncnn的源码,见上图。主要的文件就是tensorflow2ncnn.cpp,其他的proto文件就是模型数据结构的protobuf文件。接下来我们就分1)解析tensorflow文件;2)解析出node和op对应ncnn的layer写进文件;来讲解tensorflow2ncnn.cpp。

  • 解析tensorflow文件
    从main函数进入
  • int main(int argc, char** argv)
        //传入参数1:tensorflow pb文件路径
        const char* tensorflowpb = argv[1];       
        //传入参数2:生成的ncnn param文件
        const char* ncnn_prototxt = argc >= 4 ? argv[2] : "ncnn.param"; 
        //传入参数3:生成的ncnn bin文件   
        const char* ncnn_modelbin = argc >= 4 ? argv[3] : "ncnn.bin";       
        //声明tensorflow的数据结构graph,又来接收解析出的tensorflow模型结构
        tensorflow::GraphDef graph;
        // read_proto_from_binary函数的作用是根据graph.proto文件定义好的数据结构解析tensorflow模型
        bool s1 = read_proto_from_binary(tensorflowpb, &graph);
        //如果解析失败
        if (!s1)
            fprintf(stderr, "read_proto_from_binary failed\n");
            return -1;
        //新建打开要生成的ncnn的模型文件
        FILE* pp = fopen(ncnn_prototxt, "wb");
        FILE* bp = fopen(ncnn_modelbin, "wb");
        //在param文件中写入魔法数字,如对ncnn模型文件的格式不熟悉,可以看看本文集下的ncnn源码笔记
        fprintf(pp, "7767517\n");
        //node_size()可以得到graph里node的数量
        int node_count = graph.node_size();
        // node的索引
        std::map<std::string, int> node_reference;
        //用于存储解析出的层的权重参数及对应层名
        std::map<std::string, tensorflow::TensorProto> weights;
        // 存储解析出的Dropout层
        std::set<std::string> dropouts;
        // 存储解析出的算子
        std::map<std::string, tensorflow::TensorProto> binaryop_consts;
    

    read_proto_from_binary函数

    static bool read_proto_from_binary(const char* filepath, google::protobuf::Message* message)
        std::ifstream fs(filepath, std::ifstream::in | std::ifstream::binary);
        if (!fs.is_open())
            fprintf(stderr, "open failed %s\n", filepath);
            return false;
        google::protobuf::io::IstreamInputStream input(&fs);
        google::protobuf::io::CodedInputStream codedstr(&input);
        codedstr.SetTotalBytesLimit(INT_MAX, INT_MAX / 2);
        bool success = message->ParseFromCodedStream(&codedstr);
        fs.close();
        return success;
    

    从for开始就是解析出node和op对应ncnn的layer写进文件

    for (int i=0; i<node_count; i++)