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()
|