45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
import torch
|
||
|
||
def test_cuda_availability():
|
||
"""全面测试CUDA可用性"""
|
||
|
||
print("="*50)
|
||
print("PyTorch CUDA 测试")
|
||
print("="*50)
|
||
|
||
# 基本信息
|
||
print(f"PyTorch版本: {torch.__version__}")
|
||
print(f"CUDA可用: {torch.cuda.is_available()}")
|
||
|
||
if not torch.cuda.is_available():
|
||
print("CUDA不可用，可能原因：")
|
||
print("1. 未安装CUDA驱动")
|
||
print("2. 安装的是CPU版本的PyTorch")
|
||
print("3. CUDA版本与PyTorch不匹配")
|
||
return False
|
||
|
||
# 设备信息
|
||
device_count = torch.cuda.device_count()
|
||
print(f"发现 {device_count} 个CUDA设备")
|
||
|
||
for i in range(device_count):
|
||
print(f"\n设备 {i}:")
|
||
print(f" 名称: {torch.cuda.get_device_name(i)}")
|
||
print(f" 内存总量: {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB")
|
||
print(f" 计算能力: {torch.cuda.get_device_properties(i).major}.{torch.cuda.get_device_properties(i).minor}")
|
||
|
||
# 简单张量测试
|
||
print("\n运行CUDA测试...")
|
||
try:
|
||
x = torch.randn(3, 3).cuda()
|
||
y = torch.randn(3, 3).cuda()
|
||
z = x + y
|
||
print("CUDA计算测试: 成功!")
|
||
print(f"设备上的张量形状: {z.shape}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"CUDA计算测试失败: {e}")
|
||
return False
|
||
|
||
if __name__ == "__main__":
|
||
test_cuda_availability() |