机器学习实践：onnx模型转为Tensorflow2的pb模型2020

1. 安装Tensorflow，如果已安装则检查 Tensorflow版本是否大于2.3.1
   

2. 安装tensorflow-addons
   

3. 从源码安装最新的onnx(2020/12),参考git

4. 将YoloV5中的torch pt模型输出为onnx模型，参考git
   

5. 将onnx转化为tf的pb模型

import onnx, onnxruntime
from onnx_tf.backend import prepare
import torch
import numpy as np
import tensorflow as tf

def onnx_to_pb(onnx_path, pb_path):
    onnx_model = onnx.load(onnx_path)  # load onnx model
    tf_rep = prepare(onnx_model)  # prepare tf representation
    tf_rep.export_graph(pb_path)  # export the model

测试加载模型的输出形状

6. 将保存好的tf模型做测试或预训练模型，大功告成。整个过程中注意各个包的版本符合要求。