Ubuntu上的PyTorch与CUDA通常是兼容的,需注意以下要点:
- 版本匹配:PyTorch版本需与CUDA版本对应,如PyTorch 2.7.1需CUDA 12.8,2.6.0需CUDA 12.6。可通过
nvidia-smi
查看驱动支持的CUDA版本,确保其≥PyTorch所需版本。
- 安装顺序:先安装NVIDIA驱动,再安装对应版本的CUDA Toolkit和cuDNN,最后通过conda或pip安装匹配的PyTorch版本。
- 环境隔离:建议使用conda虚拟环境管理不同CUDA版本的PyTorch,避免冲突。
- 验证方法:安装后通过
import torch; print(torch.cuda.is_available())
检查是否可用,返回True
则表示兼容。