54 lines
1.3 KiB
Python
54 lines
1.3 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
导出 480x480 分辨率的 TensorRT 引擎
|
||
|
|
"""
|
||
|
|
|
||
|
|
from ultralytics import YOLO
|
||
|
|
import torch
|
||
|
|
|
||
|
|
def main():
|
||
|
|
print("="*60)
|
||
|
|
print("导出 480x480 TensorRT 引擎")
|
||
|
|
print("="*60)
|
||
|
|
|
||
|
|
# 加载 YOLOv11n 模型
|
||
|
|
model_path = "C:/Users/16337/PycharmProjects/Security/yolo11n.pt"
|
||
|
|
print(f"\n加载模型: {model_path}")
|
||
|
|
model = YOLO(model_path)
|
||
|
|
|
||
|
|
# 导出为 TensorRT 引擎
|
||
|
|
print("\n开始导出 TensorRT 引擎...")
|
||
|
|
print("配置:")
|
||
|
|
print(" - 输入尺寸: 480x480")
|
||
|
|
print(" - 精度: FP16")
|
||
|
|
print(" - 批次大小: 动态 (1-32)")
|
||
|
|
print()
|
||
|
|
|
||
|
|
try:
|
||
|
|
# 导出 TensorRT 引擎
|
||
|
|
model.export(
|
||
|
|
format='engine',
|
||
|
|
imgsz=480, # 480x480 分辨率
|
||
|
|
half=True, # FP16 精度
|
||
|
|
dynamic=True, # 动态批次
|
||
|
|
batch=8, # 优化批次大小
|
||
|
|
workspace=4, # 4GB workspace
|
||
|
|
verbose=True
|
||
|
|
)
|
||
|
|
|
||
|
|
print("\n✅ TensorRT 引擎导出成功！")
|
||
|
|
print(f"引擎文件: yolo11n.engine (480x480)")
|
||
|
|
print("\n注意: 引擎文件会保存在当前目录")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print(f"\n❌ 导出失败: {e}")
|
||
|
|
import traceback
|
||
|
|
traceback.print_exc()
|
||
|
|
return 1
|
||
|
|
|
||
|
|
return 0
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
import sys
|
||
|
|
sys.exit(main())
|