PyTorch 环境配置

本章将详细介绍如何安装和配置 PyTorch 开发环境。我们会覆盖不同操作系统、不同硬件配置的安装方式。

系统要求

硬件要求

配置	最低要求	推荐配置
CPU	x86_64 架构	现代多核 CPU
内存	4 GB	8 GB 或更多
存储	10 GB 可用空间	SSD 20 GB 以上
GPU	可选	NVIDIA GPU（6GB+ 显存）

软件要求

• 操作系统：Windows 10+/macOS 10.14+/Linux（Ubuntu 18.04+）
• Python：3.8 - 3.11（推荐 Python 3.10）
• CUDA（使用 GPU）：11.8 或 12.1

Python 版本选择

推荐使用 Python 3.10，因为它是目前 PyTorch 官方推荐版本，兼容性最好。如果你是 Python 新手，建议直接从 Python 官网 下载 3.10.x 版本。

安装方式

方式一：使用 pip 安装（推荐）

这是最简单直接的安装方式。

安装 CPU 版本

如果你的电脑没有 NVIDIA GPU，或者只是学习目的，可以使用 CPU 版本：

# 首先升级 pip
pip install --upgrade pip

# 安装最新稳定版 PyTorch（CPU only）
pip install torch torchvision torchaudio

安装 GPU 版本

如果你的电脑有 NVIDIA GPU，需要安装支持 CUDA 的版本：

# CUDA 12.1 版本（推荐）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# CUDA 11.8 版本（适合较旧的 GPU）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

如何选择 CUDA 版本？

CUDA 版本	适用 GPU 架构	说明
CUDA 12.1	Ampere (30xx), Ada (40xx), Hopper (H100)	最新，性能最好
CUDA 11.8	Turing (16xx, 20xx), Ampere (30xx)	兼容性好

方式二：使用 conda 安装

# 创建新环境（推荐）
conda create -n pytorch python=3.10
conda activate pytorch

# 安装 PyTorch（CPU 版本）
conda install pytorch torchvision torchaudio cpuonly -c pytorch

# 安装 PyTorch（GPU 版本）
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

方式三：使用 Docker

如果你熟悉 Docker，可以使用官方提供的 Docker 镜像：

# 拉取最新的 PyTorch Docker 镜像
docker pull pytorch/pytorch:latest

# 运行容器
docker run --gpus all -it pytorch/pytorch:latest bash

注意

Docker 方式需要正确配置 NVIDIA Docker（nvidia-docker），这需要 NVIDIA 驱动和 CUDA 运行时支持。

验证安装

安装完成后，需要验证 PyTorch 是否正确安装。

基础验证

# Python 代码
import torch

# 打印 PyTorch 版本
print(f"PyTorch 版本: {torch.__version__}")

# 验证 CUDA 是否可用
print(f"CUDA 可用: {torch.cuda.is_available()}")

# 如果 CUDA 可用，显示 GPU 信息
if torch.cuda.is_available():
    print(f"CUDA 版本: {torch.version.cuda}")
    print(f"GPU 数量: {torch.cuda.device_count()}")
    print(f"GPU 名称: {torch.cuda.get_device_name(0)}")

预期输出（无 GPU）：

PyTorch 版本: 2.1.0
CUDA 可用: False

预期输出（有 GPU）：

PyTorch 版本: 2.1.0+cu121
CUDA 可用: True
CUDA 版本: 12.1
GPU 数量: 1
GPU 名称: NVIDIA GeForce RTX 3080

GPU 计算验证

如果安装了 GPU 版本，可以测试 GPU 计算：

import torch

# 创建一个张量并移动到 GPU
if torch.cuda.is_available():
    # 创建张量
    x = torch.randn(1000, 1000)
    
    # 移动到 GPU
    x_gpu = x.cuda()
    
    # 执行矩阵乘法
    result = torch.mm(x_gpu, x_gpu)
    
    print("GPU 计算测试成功！")
    print(f"结果张量设备: {result.device}")
else:
    print("CUDA 不可用，将使用 CPU 计算")

cuDNN 验证

cuDNN 是 NVIDIA 提供的深度学习库，PyTorch 的 GPU 计算依赖它：

import torch

print(f"cuDNN 版本: {torch.backends.cudnn.version()}")
print(f"cuDNN 启用: {torch.backends.cudnn.enabled}")

开发环境配置

使用 Jupyter Notebook

Jupyter Notebook 是学习 PyTorch 的好工具：

# 安装 Jupyter
pip install jupyter

# 启动 Jupyter
jupyter notebook

使用 VS Code

VS Code 是最受欢迎的 Python 开发 IDE 之一：

1. 安装 VS Code：从 https://code.visualstudio.com/ 下载

2. 安装 Python 扩展：

   • 打开 VS Code
   • 点击左侧扩展图标
   • 搜索 "Python" 并安装
   • 搜索 "Jupyter" 并安装

3. 配置 Python 环境：

   • 按 Ctrl+Shift+P
   • 输入 "Python: Select Interpreter"
   • 选择你安装的 PyTorch 环境

PyCharm 配置

如果你更喜欢 PyCharm：

1. 创建新项目或打开现有项目
2. 进入 File → Settings → Project → Python Interpreter
3. 点击齿轮图标 → Add
4. 选择 Existing environment，浏览到你的 PyTorch 环境

依赖管理

requirements.txt

创建 requirements.txt 文件管理项目依赖：

torch>=2.0.0
torchvision>=0.15.0
torchaudio>=2.0.0
numpy>=1.21.0
pandas>=1.3.0
matplotlib>=3.4.0
scikit-learn>=1.0.0

安装所有依赖：

pip install -r requirements.txt

使用虚拟环境

强烈建议使用虚拟环境隔离项目依赖：

# 使用 venv
python -m venv pytorch-env
source pytorch-env/bin/activate  # Linux/macOS
pytorch-env\Scripts\activate     # Windows

# 安装 PyTorch
pip install torch torchvision torchaudio

GPU 环境配置

NVIDIA 驱动安装

使用 GPU 前需要安装 NVIDIA 驱动：

Windows：

1. 访问 https://www.nvidia.com/Download/index.aspx
2. 选择你的 GPU 型号和操作系统
3. 下载并安装驱动

Linux：

# Ubuntu/Debian
sudo apt update
sudo apt install nvidia-driver-535

# 验证安装
nvidia-smi

CUDA 安装

如果需要从源码编译 PyTorch，需要安装 CUDA Toolkit：

1. 从 https://developer.nvidia.com/cuda-downloads 下载
2. 按照安装向导完成安装
3. 添加环境变量：

# ~/.bashrc 或 ~/.zshrc
export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

多 GPU 配置

如果有多 GPU，可以指定使用哪个：

import torch

# 查看所有 GPU
for i in range(torch.cuda.device_count()):
    print(f"GPU {i}: {torch.cuda.get_device_name(i)}")

# 指定使用第一个 GPU
device = torch.device('cuda:0')

# 张量移动到指定 GPU
x = torch.randn(100, 100).to(device)

常见问题

问题 1：pip 安装失败

症状：安装超时或下载失败

解决方案：

• 使用国内镜像源：

  pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

• 检查网络连接
• 使用 conda 安装作为替代

问题 2：CUDA 不可用

症状：torch.cuda.is_available() 返回 False

解决方案：

1. 确认 NVIDIA 驱动已安装：运行 nvidia-smi
2. 确认 CUDA 版本与 PyTorch 版本匹配
3. 重新安装对应版本的 PyTorch

问题 3：GPU 内存不足

症状：训练时出现 CUDA out of memory 错误

解决方案：

1. 减小 batch size
2. 使用梯度累积
3. 启用混合精度训练
4. 使用 torch.cuda.empty_cache() 清理缓存

问题 4：安装后导入失败

症状：ImportError: DLL load failed

解决方案：

• Windows：安装 Visual C++ Redistributable
• 重新安装 PyTorch
• 检查 Python 版本兼容性

卸载 PyTorch

如果需要卸载：

# pip 卸载
pip uninstall torch torchvision torchaudio

# conda 卸载
conda uninstall pytorch torchvision torchaudio pytorch-cuda

快速开始模板

安装完成后，可以使用以下模板开始项目：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 检查设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 示例：创建一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = x.view(-1, 784)  # 展平输入
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建模型并移动到设备
model = SimpleNet().to(device)
print(model)

# 打印模型参数设备
for param in model.parameters():
    print(f"参数设备: {param.device}")

下一步

环境配置完成后，让我们开始学习 PyTorch 的核心数据结构——张量（Tensor）。