ubuntu

Ubuntu下PyTorch支持CUDA吗

小樊
46
2025-05-02 02:49:54
栏目: 智能运维

Ubuntu下PyTorch支持CUDA。以下是在Ubuntu系统上安装并使用PyTorch与CUDA的步骤:

安装CUDA

  1. 检查GPU兼容性
  1. 下载CUDA Toolkit
  1. 安装CUDA Toolkit
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
  1. 配置环境变量
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}}
  1. 验证CUDA安装

安装PyTorch with CUDA支持

  1. 创建虚拟环境(可选)
python3 -m venv pytorch_env
source pytorch_env/bin/activate
  1. 安装PyTorch
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
  1. 验证PyTorch与CUDA的集成
import torch
print(torch.cuda.is_available())  # 应该返回True
print(torch.cuda.current_device())  # 显示当前使用的GPU设备编号
print(torch.cuda.get_device_name(0))  # 显示第一个GPU的设备名称

注意事项

总之,只要按照上述步骤正确配置,Ubuntu下的PyTorch就可以顺利地利用CUDA加速计算了。

0
看了该问题的人还看了