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

在这里插入图片描述
测试加载模型的输出形状
在这里插入图片描述

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

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐