PyTorch在Ubuntu上支持GPU吗

是的，PyTorch在Ubuntu上支持GPU。为了在Ubuntu上使用PyTorch的GPU功能，需要安装CUDA Toolkit和相应的PyTorch GPU版本。以下是详细步骤：

安装CUDA Toolkit

1. 检查显卡驱动：首先，确保你的系统已经安装了NVIDIA显卡驱动。可以通过在终端中输入 nvidia-smi 来检查驱动是否安装以及GPU信息。

2. 下载并安装CUDA Toolkit：

   • 访问NVIDIA CUDA Toolkit下载页面，根据你的Ubuntu系统版本选择合适的CUDA版本进行下载。
   • 例如，对于Ubuntu 22.04，可以下载CUDA 12.1版本。
   • 下载完成后，按照官方文档提供的步骤进行安装。

3. 配置环境变量：

   • 编辑 ~/.bashrc 文件，添加以下内容：

     export PATH=/usr/local/cuda-12.1/bin:$PATH
     export LD_LIBRARY_PATH=/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
     

   • 保存文件并执行 source ~/.bashrc 使环境变量生效。

4. 验证安装：

   • 在终端中输入 nvcc -V 验证CUDA是否安装成功。

安装PyTorch GPU版本

1. 创建并激活Conda环境（推荐使用Conda管理环境）：

   conda create -n torch_env python=3.8
   conda activate torch_env
   

2. 安装PyTorch GPU版本：

   • 根据你的CUDA版本，从PyTorch官方网站选择相应的安装命令。例如，使用CUDA 12.1：

     pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 -i https://pypi.tuna.tsinghua.edu.cn/simple
     

   • 或者使用Conda安装：

     conda install pytorch torchvision torchaudio cudatoolkit=12.1 -c pytorch
     

3. 验证PyTorch GPU支持：

   • 在Python中运行以下代码检查GPU是否可用：

     import torch
     print(torch.cuda.is_available())
     

   • 如果返回 True，则表示PyTorch已经正确配置并可以使用GPU。

通过以上步骤，你可以在Ubuntu系统上成功安装并配置PyTorch以使用GPU。