Tensorflow调用训练好的yolov5模型进行推理

文章目录

1、安装TensorFlow-GPU版本

bash 复制代码
1、创建虚拟环境python=3.8
conda create -n TF2.4 python=3.8

2、进入虚拟环境
conda activate TF2.4

3、去官网查看tensorflow对应cuda的版本
https://tensorflow.google.cn/install/source_windows?hl=zh-cn

conda search找找当前源下的CUDA与cuDNN有没有我们要的版本:

bash 复制代码
conda search cuda
conda search cudnn

conda install cudatoolkit=11.0.221
conda install cudnn=8.2.1

# 安装tensorflow-gpu版本
pip install tensorflow-gpu==2.4.0


1.2、验证是否安装正常

python 复制代码
import tensorflow as tf
print(tf.__version__)
print(tf.test.is_gpu_available())

2、将训练好的pt文件转换成onnx文件

2.2、什么是Onnx模型和Tensorflow模型

Onnx(Open Neural Network Exchange)是一种开放的深度学习模型交换格式,用于在不同的深度学习框架之间共享模型。它提供了一个中间格式,可以将模型从一个框架转换为另一个框架。

Tensorflow是一个广泛使用的深度学习框架,提供了强大的模型构建和训练工具。Tensorflow模型通常以.pb文件格式保存,它包含了模型的结构和参数。

2.1、将onnx文件转换成pb文件

可以通过yolov5中export.py文件进行转换

python 复制代码
python export.py --weights weights/best.pt --include onnx engine --img 640 --device 0

pb文件是tensorflow中可以使用的文件

使用代码进行转换

pip install --user tensorflow_probability==0.7.0

python 复制代码
import onnx
from onnx_tf.backend import prepare
import tensorflow.keras as keras



# 加载Onnx模型
onnx_model = onnx.load('best.onnx')

# 转换为Tensorflow模型
tf_model = prepare(onnx_model)

# 保存为.pb文件
tf_model.export_graph('best.pb')

源码进行转换

源码下载:https://github.com/onnx/onnx-tensorflow#installation

之后再终端安装:

python 复制代码
pip install -e.
pip install tensorflow-addons

都安装好之后就可以直接在终端进行转换了

python 复制代码
onnx-tf convert -i D:\\yolov5_back\\weights\\best.onnx -o D:\\yolov5_back\\weights\\best.onnx.pb

转换过程中可能会报两个错误:

错误1:from keras import backend

ModuleNotFoundError: No module named 'keras'

python 复制代码
# 点击进入报错的文件中,修改这行
from tensorflow.keras import backend

错误2:

from tensorflow.keras.utils import tf_utils

ImportError: cannot import name 'tf_utils' from 'tensorflow.keras.utils' (C:\Users\Administrator.conda\envs\Tensorflow_gpu_2.4\lib\site-packages\tensorflow\keras\utils_init_.py)

python 复制代码
# 解决方法直接将这行注释掉
from keras.utils import tf_utils
相关推荐
羊羊小栈5 天前
基于GraphRAG的医疗健康知识诊断系统(Neo4j_大语言模型)
人工智能·语言模型·毕业设计·知识图谱·创业创新·neo4j·大作业
羊羊小栈9 天前
基于GraphRAG的地质矿产知识管理系统(Neo4j_大语言模型)
人工智能·语言模型·自然语言处理·毕业设计·neo4j·大作业
夜郎king10 天前
SpringBoot 整合 Neo4j 实战:从零搭建经典小说知识图谱完整方案
spring boot·知识图谱·neo4j
羊羊小栈11 天前
基于知识图谱(Neo4j)和大语言模型(LLM)的图检索增强(GraphRAG)的数控车床主轴系统故障诊断智能问答系统
人工智能·语言模型·毕业设计·知识图谱·创业创新·neo4j·大作业
_Rookie._12 天前
neo4j图形数据库 -基础语法
neo4j
星川皆无恙12 天前
基于BERT+LSTM+CRF与知识图谱的医疗智能问答系统实战:Neo4j图数据库+实体识别+意图分析完整项目
数据库·人工智能·深度学习·bert·lstm·知识图谱·neo4j
飞行家贞贞14 天前
Windows 安装 Neo4j(2025最新·极简)
neo4j
ins_lizhiming15 天前
多模态+neo4j+langchain构建知识图谱问答
langchain·知识图谱·neo4j
qq_3564086616 天前
Neo4j RPM 安装与多实例配置说明
neo4j
codefan※17 天前
干掉幻觉实战:如何构建企业级知识图谱增强 RAG
人工智能·大模型·llm·知识图谱·neo4j·rag·graphrag