Pytorch模型转Tensorflow模型的那些事

2024-02-21 19:44:28

Pytorch模型转Tensorflow模型部署上线

pytorch模型转tensorflow流程

最近在研究Query2Title模型，学术界上快速实验一般都用pytorch，但是业界部署模型上大多都还是tensorflow模型部署。也可能是自己太懒了，哈哈，pytorch用久了tensorflow有点生疏，懒的去用啦！结果就…，线上部署必须得tensorflow模型，所以就把torch的.pt文件，转成tensorflow的pb文件把，中间遇到一些问题记录下来，仅供借鉴参考！<ps!立下flag，坚持写一些博客>

pytorch模型转tensorflow流程

总体上pytorch转tensorflow遵循2个步骤：

torch模型文件.pt保存成onnx模型文件.onnx
.onnx模型文件转换成tensorflow .pb文件

torch模型文件转onnx文件

这里直接利用torch的onnx库将.pt文件load后导出为.onnx文件即可，基本上没有什么坑，如下：

def torch2Onnx():
    """
    pytorch转onnx
    """
    model_pytorch = YourModel()
    model_pytorch.load_state_dict(torch.load('xxx.pt',
                                       map_location='cpu'))
    # 输入placeholder
    dummy_input = torch.randint(0, 10000, (1, 20))
    dummy_output = model_pytorch(dummy_input)
    print(dummy_output.shape)

    # Export to ONNX format
    torch.onnx.export(model_pytorch, 
                      dummy_input, 
                      'model.onnx', 
                      input_names=['inputs'], 
                      output_names=['outputs'])

需要注意的是：tensorflow1.x模型是需要先定义计算图的，计算图需要输入输出节点，所以最好指定计算图的输入输出节点名称，以便后续调用tensorflow模型

.onnx文件转tensorflow .pb文件

这一步利用开源工具onnx-tf (github) 将.onnx文件转换成.pb文件。可以直接源码安装，也可以直接 pip安装，转换过程如下：

def onnx2Tensorflow(onnx_model="model.onnx", tf_model="model.tf.pb"):
    """
    onnx转tensorflow
    """
    """
    # ---------------------报错-------------------------
    # Load ONNX model and convert to TensorFlow format
    model_onnx = onnx.load(onnx_model)
    tf_rep = prepare(model_onnx)
    # Export model as .pb file
    tf_rep.export_graph(tf_model)
    """
    os.system("onnx-tf convert -i %s -o %s"%(onnx_model, tf_model))

官方教程是按照注释中的内容做的，但是我自己执行报错，但是命令行方式转换可以成功。

导入计算图，测试样例

转换成功后，我们加载计算图测试一下样例输出是否正确，如下：

def load_pb(path_to_pb):
    """
    加载pb文件
    """
    with tf.gfile.GFile(path_to_pb, 'rb') as f:
        # 定义计算图
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')
        return graph

def loadTFModel():
    """
    加载tensorflow模型测试
    """
    graph = load_pb('model.tf.pb')
    with tf.Session(graph=graph) as sess:
        # Show tensor names in graph
        for op in graph.get_operations():
            # 获取节点名称
            print(op.values())
        # 获取输入输出节点
        output_tensor = graph.get_tensor_by_name('div_3:0')
        input_tensor = graph.get_tensor_by_name('inputs:0')

        # dummy_input = np.random.randint(0, 1000, (1, 20), dtype=np.int64)
        query1 = np.array([[4158, 7811, 6653,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]], dtype=np.int64)
        query2 = np.array([[9914, 10859, 6907, 8719, 7098, 
                            8861, 4158, 10785, 6299, 1264, 
                            1612, 10285, 6973, 7811, 0,
                            0, 0, 0, 0, 0]],
                           dtype=np.int64)
        output1 = sess.run(output_tensor, 
                          feed_dict={input_tensor: query1})
        output2 = sess.run(output_tensor,
                          feed_dict={input_tensor: query2})
        
        simi = np.dot(output1[0], output2[0])
        print(simi)

很明显在测试过程中，需要给输入张量喂入输出，获取输出节点数据，这个时候我们在第一步指定的输入、输出节点变量名就派上了用场。input_tensor获取输入节点的张量值就使用了变量"inputs:0"，":0"表示该节点的第一个输出，如果该节点只有一个输出就是:0;可以看到，在第一步中我们设定的输出变量名是outputs，但是这里并不是。（个人觉得是有点bug的，所以具体的输出节点名称要从计算图上的节点中去查看。当然如果使用最新版本的onnx-tf这个问题是不存在的，由于我需要的pb文件是tf 1.12版本的，亲测1.13版本以上，没有变量名映射的bug）

注意版本差异和环境

不同的tensorflow版本和onnx-tf版本需要适配：个人总结tf 1.13以上的版本可以使用最新的onnx-tf 1.7版本；tf 1.12版本可以使用onnx-tf 1.3版本。
注意模型推理是在cpu还是gpu上：tensorflow模型在cpu上推理和gpu上推理存在差异！torch文件默认的输入数据格式是NxCxHxW，而TensorFlow中的函数在gpu环境上都是支持NxCxHxW格式的。但是部分函数，例如max_pool在cpu 上只支持NxHxHxC的格式。解决办法就是，如果要在cpu上推理，需要禁用gpu: os.environ[“CUDA_VISIBLE_DEVICES”]="-1" ，否则转换会默认按照gpu环境导出；当然你如果没有gpu环境，那就不存在这个问题，自动按照cpu运行环境导出。

各位看客小伙伴，有什么问题或者好的解决办法可以留言评论哦！

码农公寓