chore: initial commit for Infrared Equipment Detection System
- Add YOLO-based infrared thermal image detection system - Add dataset configuration (dataset.yaml) with 12 equipment classes - Add training, validation and detection scripts - Add temperature feature extraction and fault diagnosis modules - Update README and CLAUDE.md with correct dataset paths - Optimize .gitignore to track dataset.yaml while ignoring large data files
This commit is contained in:
SkyLi
@@ -0,0 +1,230 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
CUDA和cuDNN环境测试脚本
|
||||
验证PyTorch GPU支持是否正常
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
|
||||
def test_cuda():
|
||||
"""测试CUDA支持"""
|
||||
print("=" * 50)
|
||||
print("CUDA 环境测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
import torch
|
||||
print(f"PyTorch版本: {torch.__version__}")
|
||||
print(f"CUDA可用: {torch.cuda.is_available()}")
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用")
|
||||
return False
|
||||
|
||||
print(f"CUDA版本: {torch.version.cuda}")
|
||||
|
||||
# 检查cuDNN版本
|
||||
try:
|
||||
cudnn_ver = torch.backends.cudnn.version()
|
||||
print(f"cuDNN版本: {cudnn_ver}")
|
||||
except RuntimeError as e:
|
||||
print(f"cuDNN版本检测失败: {e}")
|
||||
print("警告: cuDNN可能无法正常使用，但GPU张量运算可能仍然正常")
|
||||
return False
|
||||
|
||||
print(f"GPU数量: {torch.cuda.device_count()}")
|
||||
|
||||
for i in range(torch.cuda.device_count()):
|
||||
print(f"\nGPU {i}:")
|
||||
print(f" 名称: {torch.cuda.get_device_name(i)}")
|
||||
props = torch.cuda.get_device_properties(i)
|
||||
print(f" 计算能力: {props.major}.{props.minor}")
|
||||
print(f" 总内存: {props.total_memory / 1024**3:.2f} GB")
|
||||
|
||||
return True
|
||||
|
||||
except ImportError:
|
||||
print("错误: PyTorch未安装")
|
||||
return False
|
||||
|
||||
|
||||
def test_gpu_tensor_ops():
|
||||
"""测试GPU张量运算"""
|
||||
print("\n" + "=" * 50)
|
||||
print("GPU 张量运算测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用，跳过GPU测试")
|
||||
return False
|
||||
|
||||
# 创建张量并进行运算
|
||||
print("创建1000x1000随机张量...")
|
||||
x = torch.randn(1000, 1000).cuda()
|
||||
y = torch.randn(1000, 1000).cuda()
|
||||
|
||||
print("执行矩阵乘法...")
|
||||
z = torch.matmul(x, y)
|
||||
|
||||
print("执行归约操作...")
|
||||
result = z.sum().item()
|
||||
|
||||
print(f"结果: {result:.2f}")
|
||||
print("GPU张量运算测试通过!")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"错误: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_cudnn():
|
||||
"""测试cuDNN加速"""
|
||||
print("\n" + "=" * 50)
|
||||
print("cuDNN 加速测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用")
|
||||
return False
|
||||
|
||||
# 检查cuDNN是否启用
|
||||
print(f"cuDNN可用: {torch.backends.cudnn.is_available()}")
|
||||
try:
|
||||
cudnn_ver = torch.backends.cudnn.version()
|
||||
print(f"cuDNN版本: {cudnn_ver}")
|
||||
except RuntimeError as e:
|
||||
print(f"cuDNN版本检测跳过: {e}")
|
||||
print("注意: PyTorch绑定的cuDNN与系统版本可能不匹配，但不影响GPU使用")
|
||||
print(f"cuDNN基准测试启用: {torch.backends.cudnn.benchmark}")
|
||||
|
||||
# cuDNN卷积测试
|
||||
print("\n执行cuDNN卷积测试...")
|
||||
x = torch.randn(1, 3, 224, 224).cuda()
|
||||
conv = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1).cuda()
|
||||
|
||||
# 预热
|
||||
_ = conv(x)
|
||||
|
||||
# 计时测试
|
||||
import time
|
||||
start = time.time()
|
||||
for _ in range(10):
|
||||
_ = conv(x)
|
||||
torch.cuda.synchronize()
|
||||
elapsed = time.time() - start
|
||||
|
||||
print(f"10次卷积耗时: {elapsed*1000:.2f} ms")
|
||||
print("cuDNN测试通过!")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"错误: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def test_ultralytics_gpu():
|
||||
"""测试ultralytics YOLO GPU支持"""
|
||||
print("\n" + "=" * 50)
|
||||
print("Ultralytics YOLO GPU 测试")
|
||||
print("=" * 50)
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
import torch
|
||||
import os
|
||||
|
||||
# 确保使用GPU加载模型
|
||||
if not torch.cuda.is_available():
|
||||
print("错误: CUDA不可用")
|
||||
return False
|
||||
|
||||
print("加载YOLOv8n模型到GPU...")
|
||||
# 关键：加载模型时指定device=0，确保模型在GPU上
|
||||
model = YOLO("yolov8n.pt", task="detect")
|
||||
|
||||
# 将模型移到GPU
|
||||
model.to("cuda:0")
|
||||
|
||||
print(f"模型设备: {next(model.model.parameters()).device}")
|
||||
|
||||
# 使用本地测试图像进行推理
|
||||
test_image = "bus.jpg"
|
||||
if not os.path.exists(test_image):
|
||||
print("下载测试图像...")
|
||||
import urllib.request
|
||||
urllib.request.urlretrieve(
|
||||
"https://ultralytics.com/images/bus.jpg",
|
||||
test_image
|
||||
)
|
||||
|
||||
print("执行GPU推理测试...")
|
||||
# 关键：device=0指定使用GPU
|
||||
results = model(test_image, device=0, verbose=False)
|
||||
|
||||
# 验证结果
|
||||
num_detections = len(results[0].boxes)
|
||||
print(f"检测到 {num_detections} 个目标")
|
||||
|
||||
# 检查模型是否真的在GPU上运行
|
||||
model_device = next(model.model.parameters()).device
|
||||
print(f"模型实际设备: {model_device}")
|
||||
assert model_device.type == "cuda", "模型未在GPU上运行!"
|
||||
|
||||
print("YOLO GPU测试通过!")
|
||||
return True
|
||||
|
||||
except ImportError:
|
||||
print("错误: ultralytics未安装")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"错误: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
print("\n" + "#" * 50)
|
||||
print("# 红外故障检测 - CUDA/cuDNN 环境测试")
|
||||
print("#" * 50 + "\n")
|
||||
|
||||
results = {
|
||||
"CUDA": test_cuda(),
|
||||
"GPU张量运算": test_gpu_tensor_ops(),
|
||||
"cuDNN": test_cudnn(),
|
||||
"YOLO GPU": test_ultralytics_gpu(),
|
||||
}
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("测试结果汇总")
|
||||
print("=" * 50)
|
||||
|
||||
all_passed = True
|
||||
for name, passed in results.items():
|
||||
status = "✅ 通过" if passed else "❌ 失败"
|
||||
print(f"{name}: {status}")
|
||||
if not passed:
|
||||
all_passed = False
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
if all_passed:
|
||||
print("所有测试通过! 环境配置正确。")
|
||||
else:
|
||||
print("部分测试失败，请检查环境配置。")
|
||||
print("=" * 50)
|
||||
|
||||
return 0 if all_passed else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user