Changes fix model transfer bugs
This commit is contained in:
parent
7e38b3acfd
commit
3ba3097b6f
@ -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
3
test_tf.py
Normal file
@ -0,0 +1,3 @@
|
||||
import tensorflow as tf
|
||||
print("TF Version:", tf.__version__)
|
||||
print("Is GPU available:", tf.test.is_gpu_available())
|
Loading…
Reference in New Issue
Block a user