pb graph鼠标移上显示数据_pytorch模型转tensorflow的pb文件

pytorch训练操作方便，tensorflow部署上线效果好，所以就有模型转化这个需求。

1.首先新建一个py脚本

# 导入torch训练时的模型类
from ** import YourTorchModel

2.导入工具包

# onnx-tf要求tensorflow==2.2.0
# onnx-tf==1.6

import torch
import torch.nn as nn
import torch.onnx
import onnx
from onnx_tf.backend import prepare
import argparse
import os

3.加载torch模型文件

# 需要转化的模型文件路径
model_path = "model.pth" 
# 加载模型
model = torch.load(model_path)
print(model)

4.把pth文件转为onnx文件

model = torch.nn.DataParallel(model)

# 设置模型输入维度
input = torch.randn([1,2,3])

# 设置输入张量名，多个输入就是多个名
input_names = ["input"]

# 设置输出张量名
output_names = ["output"]

# 自定义onnx文件名和路径
onnx_filename = "model.onnx"

# 执行转化和保存
torch.onnx.export(model.module, input, onnx_filename, verbose=True, input_names=input_names,
                  output_names=output_names)

5. onnx文件转为pb文件

# 
onnx_model = onnx.load("model.onnx")  # load onnx model
tf_exp = prepare(onnx_model)  # prepare tf representation
tf_exp.export_graph("model.pb")  # export the model

6. tf1.x加载pb文件

with tf.Graph().as_default():
    output_graph_def = tf.GraphDef()
    output_graph_path = "model.pb"

    with open(output_graph_path, 'rb') as f:
        output_graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(output_graph_def, name="")
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        input = sess.graph.get_tensor_by_name("input:0")  # "input" 是在pth文件转为onnx文件时定义好的，名字要一致
        output = sess.graph.get_tensor_by_name("output:0") # ”output“ 也是
        input_data = torch.randn(2, 10) # 输入要测试的数据，格式要一致
        predictions = sess.run(output, feed_dict={input: input_data})
        print("predictions:", predictions)

7. tf2.x加载pb文件件

with tf.Graph().as_default():
    output_graph_def = tf.compat.v1.GraphDef()
    output_graph_path = "ConvNet.pb"

    with open(output_graph_path, 'rb') as f:
        output_graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(output_graph_def, name="")
    with tf.compat.v1.Session() as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        input = sess.graph.get_tensor_by_name("input:0")  # "input" 是在pth文件转为onnx文件时定义好的，名字要一致
        output = sess.graph.get_tensor_by_name("output:0")  # ”output“ 也是
        input_data = torch.randn(2, 10)  # 输入要测试的数据，格式要一致
        predictions = sess.run(output, feed_dict={input: input_data})
        print("predictions:", predictions)

参考资料：

[1]pytorch模型转tensorflow pb模型