Changes fix model transfer bugs

This commit is contained in:
fly6516 2025-06-04 02:20:09 +08:00
parent 7e38b3acfd
commit 3ba3097b6f
2 changed files with 12 additions and 27 deletions

View File

@ -1,31 +1,13 @@
import onnx
import tensorflow as tf
from onnx_tf.backend import prepare
from ultralytics import YOLO
# Load the YOLO11 model
model = YOLO("runs/detect/train/weights/best.pt")
def onnx_to_tflite(onnx_model_path, tflite_model_path):
# 加载 ONNX 模型
onnx_model = onnx.load(onnx_model_path)
# Export the model to TFLite format
model.export(format="tflite") # creates 'yolo11n_float32.tflite'
# 使用 onnx-tf 将 ONNX 模型转换为 TensorFlow 模型
tf_rep = prepare(onnx_model)
# Load the exported TFLite model
tflite_model = YOLO("runs/detect/train/weights/best.tflite")
# 保存 TensorFlow 模型
saved_model_dir = 'runs/detect/train/weights/best.pt'
tf_rep.export_graph(saved_model_dir)
# 加载 TensorFlow 模型并转换为 TFLite 模型
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()
# 保存 TFLite 模型
with open(tflite_model_path, 'wb') as f:
f.write(tflite_model)
print(f"模型已成功转换并保存为 {tflite_model_path}")
# 使用示例
onnx_model_path = 'model.onnx' # 替换为你的 ONNX 模型文件路径
tflite_model_path = 'model.tflite' # 替换为你希望保存的 TFLite 模型路径
onnx_to_tflite(onnx_model_path, tflite_model_path)
# Run inference
results = tflite_model("datasets/valid/images/2007_000042_jpg.rf.dadefef10ec4b9422a1689f7ce40588e.jpg")

3
test_tf.py Normal file
View File

@ -0,0 +1,3 @@
import tensorflow as tf
print("TF Version:", tf.__version__)
print("Is GPU available:", tf.test.is_gpu_available())