GPU测试

check_env.py

@@ -0,0 +1,111 @@
+"""
+TensorRT 环境诊断脚本
+"""
+import os
+import sys
+
+def check_env():
+    print("=" * 60)
+    print("TensorRT 环境诊断")
+    print("=" * 60)
+    
+    # 1. Python 版本
+    print(f"\n[1] Python: {sys.version}")
+    
+    # 2. CUDA 环境
+    print("\n[2] CUDA 环境:")
+    cuda_path = os.environ.get("CUDA_PATH", "未设置")
+    print(f"    CUDA_PATH: {cuda_path}")
+    
+    # 3. PyTorch CUDA
+    print("\n[3] PyTorch:")
+    try:
+        import torch
+        print(f"    版本: {torch.__version__}")
+        print(f"    CUDA 可用: {torch.cuda.is_available()}")
+        if torch.cuda.is_available():
+            print(f"    CUDA 版本: {torch.version.cuda}")
+            print(f"    GPU: {torch.cuda.get_device_name(0)}")
+    except ImportError:
+        print("    未安装")
+    
+    # 4. TensorRT Python 包
+    print("\n[4] TensorRT Python 包:")
+    try:
+        import tensorrt as trt
+        print(f"    版本: {trt.__version__}")
+        print(f"    文件位置: {trt.__file__}")
+    except ImportError as e:
+        print(f"    导入失败: {e}")
+    
+    # 5. PyCUDA
+    print("\n[5] PyCUDA:")
+    try:
+        import pycuda.driver as cuda
+        import pycuda.autoinit
+        print(f"    版本: {cuda.get_version()}")
+    except ImportError as e:
+        print(f"    导入失败: {e}")
+    except Exception as e:
+        print(f"    初始化失败: {e}")
+    
+    # 6. Ultralytics
+    print("\n[6] Ultralytics:")
+    try:
+        import ultralytics
+        print(f"    版本: {ultralytics.__version__}")
+    except ImportError:
+        print("    未安装")
+    
+    # 7. 测试 TensorRT Runtime 加载
+    print("\n[7] TensorRT Runtime 测试:")
+    try:
+        import tensorrt as trt
+        logger = trt.Logger(trt.Logger.WARNING)
+        runtime = trt.Runtime(logger)
+        print(f"    Runtime 创建成功")
+        print(f"    TRT 序列化版本: 检查 engine 文件...")
+    except Exception as e:
+        print(f"    失败: {e}")
+    
+    # 8. 检查 engine 文件
+    print("\n[8] Engine 文件检查:")
+    engine_paths = [
+        "C:/Users/16337/PycharmProjects/Security_project/yolov8n.engine",
+        "C:/Users/16337/PycharmProjects/Security_project/yolov8n_320.engine",
+    ]
+    for path in engine_paths:
+        if os.path.exists(path):
+            size = os.path.getsize(path) / (1024 * 1024)
+            print(f"    {path}")
+            print(f"      大小: {size:.2f} MB")
+            # 读取文件头
+            with open(path, 'rb') as f:
+                header = f.read(32)
+                print(f"      文件头 (hex): {header[:16].hex()}")
+        else:
+            print(f"    {path} - 不存在")
+    
+    # 9. 尝试加载 engine
+    print("\n[9] 尝试加载 Engine:")
+    for path in engine_paths:
+        if os.path.exists(path):
+            try:
+                import tensorrt as trt
+                logger = trt.Logger(trt.Logger.ERROR)
+                runtime = trt.Runtime(logger)
+                with open(path, 'rb') as f:
+                    engine = runtime.deserialize_cuda_engine(f.read())
+                if engine:
+                    print(f"    {os.path.basename(path)}: 加载成功 ✓")
+                else:
+                    print(f"    {os.path.basename(path)}: 加载失败 (engine=None)")
+            except Exception as e:
+                print(f"    {os.path.basename(path)}: 加载失败 - {e}")
+    
+    print("\n" + "=" * 60)
+    print("诊断完成")
+    print("=" * 60)
+
+if __name__ == "__main__":
+    check_env()