PyTorch在Linux上如何使用GPU加速

在Linux上使用PyTorch进行GPU加速，需按以下步骤操作：

1. 安装NVIDIA驱动
   确保安装与显卡兼容的驱动，可通过nvidia-smi命令检查驱动是否正常。
2. 安装CUDA Toolkit
   从NVIDIA官网下载对应版本的CUDA Toolkit（需与驱动兼容），安装后配置环境变量（PATH和LD_LIBRARY_PATH）。
3. 安装cuDNN库
   下载与CUDA版本匹配的cuDNN，解压后复制文件至CUDA安装目录并赋予权限。
4. 安装支持GPU的PyTorch
   通过pip或conda安装对应CUDA版本的PyTorch（如pip install torch --extra-index-url https://download.pytorch.org/whl/cu117），安装时需指定cudatoolkit版本。
5. 验证安装
   运行Python代码import torch; print(torch.cuda.is_available())，若输出True则配置成功。
6. 使用GPU加速
   在代码中将模型和数据移动到GPU设备：

   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
   model.to(device)  
   input_tensor = input_tensor.to(device)  
   

注意：

• 需确保CUDA、cuDNN与PyTorch版本兼容，可参考PyTorch官方版本对照表。
• 多GPU并行计算可使用torch.nn.DataParallel或DistributedDataParallel。