112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
"""
|
|
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()
|