46 lines
1.7 KiB
Python
46 lines
1.7 KiB
Python
|
import torch
|
|||
|
|
|
|||
|
|
def check_gpu_availability():
|
|||
|
|
"""检查GPU可用性并打印详细信息"""
|
|||
|
|
print("="*50)
|
|||
|
|
print("GPU/CUDA 可用性检查")
|
|||
|
|
print("="*50)
|
|||
|
|
|
|||
|
|
# 1. 检查PyTorch是否支持CUDA
|
|||
|
|
if not torch.cuda.is_available():
|
|||
|
|
print("❌ 错误: PyTorch未编译CUDA支持或未检测到NVIDIA驱动")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
# 2. 检查GPU数量
|
|||
|
|
gpu_count = torch.cuda.device_count()
|
|||
|
|
print(f"✅ 找到 {gpu_count} 个GPU设备")
|
|||
|
|
|
|||
|
|
# 3. 打印每个GPU的详细信息
|
|||
|
|
for i in range(gpu_count):
|
|||
|
|
print(f"\nGPU {i} 详细信息:")
|
|||
|
|
print(f" 名称: {torch.cuda.get_device_name(i)}")
|
|||
|
|
print(f" CUDA计算能力: {torch.cuda.get_device_capability(i)}")
|
|||
|
|
print(f" 总显存: {torch.cuda.get_device_properties(i).total_memory/1024**3:.2f} GB")
|
|||
|
|
|
|||
|
|
# 检查当前是否被选中
|
|||
|
|
current_device = torch.cuda.current_device()
|
|||
|
|
print(f" {'⭐' if i == current_device else ' '} 当前使用设备: {'是' if i == current_device else '否'}")
|
|||
|
|
|
|||
|
|
# 4. 执行实际GPU计算测试
|
|||
|
|
try:
|
|||
|
|
print("\n正在执行GPU计算测试...")
|
|||
|
|
x = torch.randn(1000, 1000).cuda()
|
|||
|
|
y = torch.randn(1000, 1000).cuda()
|
|||
|
|
z = (x @ y).mean() # 矩阵乘法
|
|||
|
|
print(f"✅ GPU计算测试成功完成! 结果: {z.item():.4f}")
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"❌ GPU计算测试失败! 错误信息:")
|
|||
|
|
print(str(e))
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
if check_gpu_availability():
|
|||
|
|
print("\n🎉 您的GPU已准备好用于PyTorch训练!")
|
|||
|
|
else:
|
|||
|
|
print("\n⚠️ 您的GPU不可用，请检查以上错误信息")
|