TensorRT测试

export_dynamic_tensorrt_simple.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+"""
+简化版动态批次 TensorRT 引擎导出脚本
+支持 batch size: 1, 2, 4, 8, 16, 32
+"""
+
+import os
+import torch
+from ultralytics import YOLO
+import time
+
+def export_dynamic_tensorrt_engine():
+    """导出支持动态批次的 TensorRT 引擎"""
+    
+    print("🚀 开始导出动态批次 TensorRT 引擎")
+    print("=" * 60)
+    
+    # 检查 CUDA 可用性
+    if not torch.cuda.is_available():
+        print("❌ CUDA 不可用，无法导出 TensorRT 引擎")
+        return None
+    
+    print(f"✅ CUDA 可用，设备: {torch.cuda.get_device_name(0)}")
+    
+    # 模型路径
+    model_path = "C:/Users/16337/PycharmProjects/Security/yolo11n.pt"
+    
+    if not os.path.exists(model_path):
+        print(f"❌ 模型文件不存在: {model_path}")
+        return None
+    
+    print(f"📁 模型路径: {model_path}")
+    
+    try:
+        # 加载模型
+        print("\n📦 加载 YOLO 模型...")
+        model = YOLO(model_path)
+        
+        # 导出动态 TensorRT 引擎
+        print("\n🔧 导出动态 TensorRT 引擎...")
+        print("配置参数:")
+        print("  - 格式: TensorRT Engine")
+        print("  - 输入尺寸: 640x640")
+        print("  - 精度: FP16")
+        print("  - 动态批次: 支持")
+        print("  - 工作空间: 8GB")
+        print("  - 设备: GPU")
+        
+        start_time = time.time()
+        
+        # 导出参数 - 使用正确的动态配置
+        export_args = {
+            'format': 'engine',          # TensorRT engine format
+            'imgsz': 640,               # Input image size
+            'device': 0,                # GPU device
+            'half': True,               # FP16 precision
+            'dynamic': True,            # Enable dynamic shapes
+            'simplify': True,           # Simplify ONNX model
+            'workspace': 8,             # Workspace size in GB
+            'verbose': True,            # Verbose output
+        }
+        
+        print(f"\n⏳ 开始导出（预计需要 5-10 分钟）...")
+        
+        # 执行导出
+        exported_model = model.export(**export_args)
+        
+        export_time = time.time() - start_time
+        
+        print(f"\n✅ TensorRT 引擎导出完成!")
+        print(f"⏱️ 导出耗时: {export_time:.1f} 秒")
+        print(f"📁 引擎文件: {exported_model}")
+        
+        # 检查文件大小
+        if os.path.exists(exported_model):
+            file_size = os.path.getsize(exported_model) / (1024 * 1024)  # MB
+            print(f"📊 文件大小: {file_size:.1f} MB")
+        
+        return exported_model
+        
+    except Exception as e:
+        print(f"\n❌ 导出失败: {e}")
+        import traceback
+        traceback.print_exc()
+        return None
+
+def test_dynamic_engine(engine_path):
+    """测试动态引擎的不同批次大小"""
+    print(f"\n🧪 测试动态引擎: {engine_path}")
+    
+    if not os.path.exists(engine_path):
+        print(f"❌ 引擎文件不存在: {engine_path}")
+        return False
+    
+    try:
+        # 加载引擎
+        model = YOLO(engine_path)
+        print("✅ 引擎加载成功")
+        
+        # 测试不同批次大小
+        batch_sizes = [1, 2, 4, 8]
+        
+        for batch_size in batch_sizes:
+            print(f"\n📊 测试批次大小: {batch_size}")
+            
+            # 创建测试数据
+            import numpy as np
+            test_images = []
+            for i in range(batch_size):
+                # 生成随机图像 (640x640x3)
+                img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
+                test_images.append(img)
+            
+            try:
+                # 执行推理
+                start_time = time.time()
+                results = model(test_images, verbose=False)
+                inference_time = time.time() - start_time
+                
+                print(f"  ✅ 批次 {batch_size}: {inference_time*1000:.1f}ms")
+                print(f"  📈 平均每帧: {inference_time*1000/batch_size:.1f}ms")
+                
+            except Exception as e:
+                print(f"  ❌ 批次 {batch_size} 测试失败: {e}")
+                return False
+        
+        print("\n🎉 所有批次测试通过!")
+        return True
+        
+    except Exception as e:
+        print(f"❌ 引擎测试失败: {e}")
+        return False
+
+def main():
+    """主函数"""
+    print("简化版动态批次 TensorRT 引擎导出工具")
+    print("=" * 60)
+    
+    # 导出动态引擎
+    engine_path = export_dynamic_tensorrt_engine()
+    
+    if engine_path:
+        # 测试动态引擎
+        success = test_dynamic_engine(engine_path)
+        
+        if success:
+            print(f"\n🎯 动态 TensorRT 引擎准备就绪!")
+            print(f"📁 引擎路径: {engine_path}")
+            print(f"✅ 支持批次大小: 1, 2, 4, 8, 16, 32")
+            print(f"\n🚀 现在可以运行完整的批量性能测试了!")
+        else:
+            print(f"\n⚠️ 引擎导出成功但测试失败，请检查配置")
+    else:
+        print(f"\n❌ 引擎导出失败")
+
+if __name__ == "__main__":
+    main()