84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
GPU状态检查工具
|
|||
|
|
此脚本用于检查PyTorch是否能够识别和使用GPU
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import sys
|
|||
|
|
|
|||
|
|
def check_gpu():
|
|||
|
|
"""检查GPU可用性并打印详细信息"""
|
|||
|
|
print("=" * 40)
|
|||
|
|
print("📊 PyTorch GPU 支持检查")
|
|||
|
|
print("=" * 40)
|
|||
|
|
|
|||
|
|
# 检查PyTorch版本
|
|||
|
|
print(f"PyTorch版本: {torch.__version__}")
|
|||
|
|
|
|||
|
|
# 检查CUDA是否可用
|
|||
|
|
cuda_available = torch.cuda.is_available()
|
|||
|
|
print(f"CUDA是否可用: {cuda_available}")
|
|||
|
|
|
|||
|
|
if cuda_available:
|
|||
|
|
# 打印CUDA版本
|
|||
|
|
print(f"CUDA版本: {torch.version.cuda}")
|
|||
|
|
|
|||
|
|
# 获取可用的GPU数量
|
|||
|
|
gpu_count = torch.cuda.device_count()
|
|||
|
|
print(f"可用GPU数量: {gpu_count}")
|
|||
|
|
|
|||
|
|
# 打印每个GPU的信息
|
|||
|
|
for i in range(gpu_count):
|
|||
|
|
print(f"\nGPU {i} 信息:")
|
|||
|
|
print(f" 名称: {torch.cuda.get_device_name(i)}")
|
|||
|
|
print(f" 计算能力: {torch.cuda.get_device_capability(i)}")
|
|||
|
|
|
|||
|
|
# 获取GPU内存信息
|
|||
|
|
try:
|
|||
|
|
mem_info = torch.cuda.get_device_properties(i).total_memory / (1024**3)
|
|||
|
|
print(f" 总内存: {mem_info:.2f} GB")
|
|||
|
|
except:
|
|||
|
|
print(" 无法获取内存信息")
|
|||
|
|
|
|||
|
|
# 进行简单的GPU测试
|
|||
|
|
print("\n执行GPU测试...")
|
|||
|
|
try:
|
|||
|
|
# 创建一个在GPU上的张量
|
|||
|
|
x = torch.rand(1000, 1000).cuda()
|
|||
|
|
y = torch.rand(1000, 1000).cuda()
|
|||
|
|
|
|||
|
|
# 计时矩阵乘法操作
|
|||
|
|
import time
|
|||
|
|
start = time.time()
|
|||
|
|
z = torch.matmul(x, y)
|
|||
|
|
torch.cuda.synchronize() # 等待GPU操作完成
|
|||
|
|
end = time.time()
|
|||
|
|
|
|||
|
|
print(f"1000x1000矩阵乘法耗时: {(end-start)*1000:.2f} ms")
|
|||
|
|
print("✅ GPU测试成功!")
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"❌ GPU测试失败: {str(e)}")
|
|||
|
|
else:
|
|||
|
|
print("\n❌ 未检测到支持CUDA的GPU")
|
|||
|
|
print("如果你的计算机有NVIDIA GPU，请确保:")
|
|||
|
|
print("1. 已安装正确的NVIDIA驱动程序")
|
|||
|
|
print("2. 已安装与PyTorch兼容的CUDA版本")
|
|||
|
|
print("3. 安装了正确的PyTorch CUDA版本")
|
|||
|
|
|
|||
|
|
print("\n建议:")
|
|||
|
|
if cuda_available:
|
|||
|
|
print("✅ 你的系统已准备好使用GPU加速药店销售预测模型")
|
|||
|
|
else:
|
|||
|
|
print("🔄 考虑使用GPU版本安装脚本来启用GPU加速:")
|
|||
|
|
if sys.platform == "win32":
|
|||
|
|
print(" 运行 install_dependencies.bat 并选择GPU选项")
|
|||
|
|
else:
|
|||
|
|
print(" 运行 ./install_dependencies.sh 并选择GPU选项")
|
|||
|
|
|
|||
|
|
print("=" * 40)
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
check_gpu()
|