Pytorch模型转Tensorflow模型的那些事

Pytorch模型转Tensorflow模型部署上线


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

pytorch模型转tensorflow流程

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

  1. torch模型文件.pt保存成onnx模型文件.onnx
  2. .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)

注意版本差异和环境

  1. 不同的tensorflow版本和onnx-tf版本需要适配:个人总结tf 1.13以上的版本可以使用最新的onnx-tf 1.7版本;tf 1.12版本可以使用onnx-tf 1.3版本。
  2. 注意模型推理是在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运行环境导出。

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

上一篇:VBS删除文件


下一篇:Windows还能这么玩❔开机自动念情书、DOS窗口变透明加上炫酷的命令特效,撩妹变得如此简单❕❕❕