文章目录
- 1、安装TensorFlow-GPU版本
- 1.2、验证是否安装正常
- 2、将训练好的pt文件转换成onnx文件
- 2.2、什么是Onnx模型和Tensorflow模型
- 2.1、将onnx文件转换成pb文件
1、安装TensorFlow-GPU版本
1、创建虚拟环境python=3.8
conda create -n TF2.4 python=3.82、进入虚拟环境
conda activate TF2.43、去官网查看tensorflow对应cuda的版本
https://tensorflow.google.cn/install/source_windows?hl=zh-cn
conda search找找当前源下的CUDA与cuDNN有没有我们要的版本：
conda search cuda
conda search cudnnconda install cudatoolkit=11.0.221
conda install cudnn=8.2.1# 安装tensorflow-gpu版本
pip install tensorflow-gpu==2.4.0
1.2、验证是否安装正常
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 export.py --weights weights/best.pt --include onnx engine --img 640 --device 0
pb文件是tensorflow中可以使用的文件
使用代码进行转换
pip install --user tensorflow_probability==0.7.0
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
之后再终端安装：
pip install -e.
pip install tensorflow-addons
都安装好之后就可以直接在终端进行转换了
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’
# 点击进入报错的文件中，修改这行
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)
# 解决方法直接将这行注释掉
from keras.utils import tf_utils
