PyTorch在Ubuntu上的多GPU支持

PyTorch在Ubuntu上支持多GPU，主要通过数据并行或模型并行的方式实现。数据并行是将模型和数据分布在多个GPU上进行训练，而模型并行则是将模型的不同部分放在不同的GPU上进行训练。以下是在Ubuntu上配置PyTorch以支持多GPU的基本步骤：

安装NVIDIA驱动

首先，确保安装了正确版本的NVIDIA显卡驱动。可以通过以下命令安装推荐的驱动：

sudo ubuntu-drivers autoinstall

安装CUDA和cuDNN

接下来，安装与PyTorch兼容的CUDA和cuDNN版本。可以通过PyTorch官网或NVIDIA官网下载并安装。

创建Conda环境

建议使用Conda来管理Python环境，因为它可以更好地处理依赖关系。创建一个新的Conda环境并激活它：

conda create -n pytorch_env python=3.8
conda activate pytorch_env

安装PyTorch

在Conda环境中安装PyTorch，可以通过以下命令安装支持多GPU的PyTorch版本：

conda install pytorch torchvision torchaudio cudatoolkit -c pytorch -c conda-forge

验证多GPU支持

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

import torch
print(torch.cuda.device_count())  # 输出GPU数量
print(torch.cuda.is_available())  # 输出True，表示CUDA可用

此外，PyTorch提供了torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel类来实现多GPU训练。使用DataParallel可以轻松地在多个GPU上并行化模型，而DistributedDataParallel则提供了更强大的分布式训练支持。

请注意，具体的安装步骤和命令可能会随着PyTorch和CUDA版本的更新而发生变化，因此在安装时请参考最新的官方文档。