机器学习实践:onnx模型转为Tensorflow2的pb模型2020
安装Tensorflow,如果已安装则检查 Tensorflow版本是否大于2.3.1安装tensorflow-addons从源码安装最新的onnx(2020/12),参考git将YoloV5中的torch pt模型输出为onnx模型,参考git将onnx转化为tf的pb模型import onnx, onnxruntimefrom onnx_tf.backend import prepareimp
·
-
安装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模型做测试或预训练模型,大功告成。整个过程中注意各个包的版本符合要求。
魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。
更多推荐



所有评论(0)