TensorRT测试

export_480_tensorrt.py

@@ -0,0 +1,53 @@
+#!/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())