pytorch 模型转onnx；onnx再转pb_onnx转pb

1、pytorch 模型转onnx

import torchvision
import torch.onnx
import torch.nn as nn


# Standard ImageNet input - 3 channels, 224x224,
# values don't matter as we care about network structure.
# But they can also be real inputs.
dummy_input = torch.randn(1, 3, 224, 224)
# Obtain your model, it can be also constructed in your script explicitly
model = torchvision.models.alexnet(pretrained=True)
# Invoke export
torch.onnx.export(model, dummy_input, "alexnet.onnx")
1
2
3
4
5
6
7
8
9
10
11
12
13

2、onnx再转pb

版本： tensorflow 2.5.0、tensorflow-probability 0.13.0
tensorflow与tensorflow-probability 版本要对应，不然报错，参考

import onnx
from onnx_tf.backend import prepare

def onnx2pb(onnx_input_path, pb_output_path):
    onnx_model = onnx.load(onnx_input_path)  # load onnx model
    tf_exp = prepare(onnx_model)  # prepare tf representation
    tf_exp.export_grap
1
2
3
4
5
6