Ubuntu下PyTorch支持CUDA。以下是在Ubuntu系统上安装并使用PyTorch与CUDA的步骤:
sudo dpkg -i cuda-repo-ubuntu2004-11-7-local_11.7.0-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-7-local/7fa2af80.pub
sudo apt-get update
sudo apt-get install cuda
~/.bashrc
文件,添加以下行:export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
source ~/.bashrc
nvidia-smi
命令,确认GPU信息和驱动状态正常。python3 -m venv pytorch_env
source pytorch_env/bin/activate
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
import torch
print(torch.cuda.is_available()) # 应该返回True
print(torch.cuda.current_device()) # 显示当前使用的GPU设备编号
print(torch.cuda.get_device_name(0)) # 显示第一个GPU的设备名称
sudo ubuntu-drivers autoinstall
命令自动安装或更新驱动。总之,只要按照上述步骤正确配置,Ubuntu下的PyTorch就可以顺利地利用CUDA加速计算了。