如何安装正确版本的PyTorch

一. 概述

当你首次使用带有GPU的云主机时，通常需要安装一个与主机CUDA版本和GPU架构兼容的PyTorch版本。选择正确的版本可以充分利用GPU的计算能力，避免运行时错误或性能问题。

二.环境检查

在安装PyTorch之前，需要确认当前环境中的GPU、CUDA版本以及PyTorch版本。

2.1 检查CUDA版本

在终端中运行以下命令：

nvcc --version

你会看到类似如下输出：

Cuda compilation tools, release 11.8, V11.8.89

从输出中可以确认主机的CUDA工具链版本为 11.8。

2.2 检查当前PyTorch版本

进入Python环境，运行以下代码：

import torch
print("PyTorch Version:", torch.__version__)
print("CUDA Version:", torch.version.cuda)
print("Is CUDA Available:", torch.cuda.is_available())

如有之前安装，则会输出类似：

PyTorch Version: 1.10.1+cu113
CUDA Version: 11.3
Is CUDA Available: True

这说明当前PyTorch版本为 1.10.1，自带 CUDA 11.3 支持，与主机的CUDA 11.8不完全匹配。

三. 选择并安装正确版本的PyTorch

3.1 从PyTorch官网获取安装命令

在深度学习环境中，确保 Python、PyTorch 和 CUDA 版本的兼容性至关重要。以下是它们的版本对应关系：

访问 PyTorch官网：

1. 选择你的配置：

   • OS: Linux
   • Package: pip
   • Compute Platform: CUDA 11.8

2. 获得安装命令（适配CUDA 11.8）：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

   如果使用集群，无法科学上网，则建议先选择合适的版本，使用如下网站：
   https://pytorch.org/get-started/previous-versions/
   然后改为使用指定的国内镜像站：

   pip install torch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 --index-url https://pypi.mirrors.ustc.edu.cn/simple

3.2 卸载旧版本的PyTorch（如适用）

为避免版本冲突，先卸载当前的PyTorch版本：

pip uninstall torch torchvision torchaudio

然后再次检查是否完全卸载：

pip list | grep torch

3.3 安装新版本

使用上面从官网获得的命令安装适配的PyTorch版本：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

3.4 验证安装

安装完成后，再次进入Python环境，运行以下代码确认新版本是否正确安装：

import torch
print("PyTorch Version:", torch.__version__)
print("CUDA Version:", torch.version.cuda)
print("Is CUDA Available:", torch.cuda.is_available())

你应该看到输出如下：

PyTorch Version: 2.x.x+cu118
CUDA Version: 11.8
Is CUDA Available: True

这说明新版本已成功安装，且与CUDA 11.8完全兼容。

四. 问题排查

4.1 如果出现“CUDA不可用”

1. 确保主机GPU驱动和CUDA工具链安装正确：

   nvidia-smi

   检查GPU是否被正确识别。

2. 确保PyTorch安装与CUDA版本匹配。

4.2 如果出现性能问题

• 设置以下环境变量来调试：

  export CUDA_LAUNCH_BLOCKING=1

  这会让CUDA错误同步显示，便于调试。