Ubuntu上PyTorch支持GPU吗

是的，PyTorch在Ubuntu上支持GPU。为了在Ubuntu上使用PyTorch的GPU功能，你需要安装CUDA和cuDNN，并确保你的NVIDIA显卡驱动已经正确安装。以下是详细步骤：

安装NVIDIA显卡驱动

首先，你需要安装与你的显卡型号相匹配的NVIDIA显卡驱动。你可以通过以下命令查看显卡型号：

lspci | grep -i nvidia

然后，根据你的显卡型号选择合适的驱动版本进行安装。

安装CUDA

接下来，你需要安装CUDA Toolkit。你可以从NVIDIA官网下载并安装适合你Ubuntu版本的CUDA Toolkit。

安装cuDNN

在安装了CUDA Toolkit之后，你还需要安装与CUDA版本相匹配的cuDNN库。你可以从NVIDIA官网下载并安装cuDNN。

安装PyTorch GPU版本

在完成了上述步骤之后，你可以使用pip或conda安装PyTorch的GPU版本。以下是使用pip安装PyTorch GPU版本的命令：

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

请将cu113替换为你安装的CUDA版本号。

验证安装

安装完成后，你可以通过以下Python代码验证PyTorch是否支持GPU：

import torch
print(torch.__version__)
print("CUDA is available. You can use GPU.") if torch.cuda.is_available() else print("CUDA is not available. Using CPU.")

如果输出显示CUDA可用，那么你就成功地在Ubuntu上配置了支持GPU的PyTorch环境。