-
安装Tensorflow,如果已安装则检查 Tensorflow版本是否大于2.3.1
-
安装tensorflow-addons
-
从源码安装最新的onnx(2020/12),参考git
-
将YoloV5中的torch pt模型输出为onnx模型,参考git
-
将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
测试加载模型的输出形状
- 将保存好的tf模型做测试或预训练模型,大功告成。整个过程中注意各个包的版本符合要求。