安裝Tensorflow，如果已安裝則檢查 Tensorflow版本是否大于2.3.1

安裝tensorflow-addons

從源碼安裝最新的onnx(2020/12),參考git

將YoloV5中的torch pt模型輸出為onnx模型，參考git

將onnx轉(zhuǎn)化為tf的pb模型

import onnx, onnxruntime from onnx_tf.backend import prepare import torch import numpy as np import tensorflow as tfdef onnx_to_pb(onnx_path, pb_path):onnx_model = onnx.load(onnx_path) # load onnx modeltf_rep = prepare(onnx_model) # prepare tf representationtf_rep.export_graph(pb_path) # export the model

測(cè)試加載模型的輸出形狀

將保存好的tf模型做測(cè)試或預(yù)訓(xùn)練模型，大功告成。整個(gè)過(guò)程中注意各個(gè)包的版本符合要求。

總結(jié)

以上是生活随笔為你收集整理的机器学习实践：onnx模型转为Tensorflow2的pb模型2020的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。