【Pytorch】PyTorch 深度学习框架详解：从安装到实战

本文系统梳理 PyTorch 的核心概念与常用 API，涵盖安装配置、张量操作、自动求导、神经网络搭建、数据加载、训练循环、模型保存与加载、GPU 加速、常用技巧等内容，配合可运行代码示例讲解，适合有 Python 基础、希望系统入门深度学习实战的开发者阅读。

📌 前言

PyTorch 是当前学术界和工业界主流的深度学习框架之一，以 动态计算图（Define-by-Run） 、接口简洁 、调试友好 著称，是 Hugging Face 等主流深度学习生态的默认底层框架。本文按照 安装 → 张量基础 → 自动求导 → 模型搭建 → 训练流程 → 进阶技巧 的顺序讲解，所有代码均可直接复制运行。

目录导航：

PyTorch 简介与安装
张量（Tensor）基础操作
自动求导机制（Autograd）
神经网络模块（nn.Module）
常用层与激活函数
损失函数与优化器
数据加载（Dataset / DataLoader）
完整训练流程实战
GPU 加速与设备管理
模型保存与加载
学习率调度器
常用进阶技巧
torch.compile 加速
常见问题排查
总结速查表

一、PyTorch 简介与安装

PyTorch 由 Meta AI 开发，目前已捐赠给 Linux 基金会旗下的 PyTorch Foundation 独立维护。它的核心组件包括：

组件	作用
`torch`	核心张量运算库，支持 GPU 加速
`torch.autograd`	自动求导引擎
`torch.nn`	神经网络层、损失函数等构建模块
`torch.optim`	优化器（SGD、Adam 等）
`torch.utils.data`	数据加载与预处理工具
`torchvision` / `torchaudio`	视觉/音频领域的数据集、模型和预处理工具

1.1 安装

PyTorch 官方建议前往 pytorch.org 官网的安装选择器，根据操作系统、包管理器（pip/conda）、计算平台（CPU/CUDA 版本）自动生成对应的安装命令，因为不同 CUDA 版本对应的安装源不同。以 pip 安装 CPU 版本为例：

复制代码

pip3 install torch torchvision torchaudio

安装 GPU（CUDA）版本时，需要根据本机 CUDA 版本指定安装源，例如：

复制代码

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

⚠️ 重要提示：CUDA 版本、显卡驱动版本、PyTorch 版本三者需要匹配，安装前务必以官网安装选择器生成的命令为准，避免手动拼接导致版本不兼容。

1.2 验证安装

复制代码

import torch

print(torch.__version__)                  # 查看 PyTorch 版本
print(torch.cuda.is_available())           # 查看是否支持 GPU 加速
print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "无可用 GPU")

二、张量（Tensor）基础操作

张量（Tensor）是 PyTorch 中最核心的数据结构，类似 NumPy 的多维数组，但额外支持 GPU 加速和自动求导。

2.1 创建张量

复制代码

import torch

t1 = torch.tensor([1, 2, 3])                  # 从列表创建
t2 = torch.zeros(3, 4)                         # 全零张量，形状 (3, 4)
t3 = torch.ones(2, 3)                          # 全一张量
t4 = torch.rand(2, 3)                          # [0, 1) 均匀分布随机数
t5 = torch.randn(2, 3)                         # 标准正态分布随机数
t6 = torch.arange(0, 10, 2)                    # 类似 range：[0, 2, 4, 6, 8]
t7 = torch.linspace(0, 1, steps=5)             # 等间隔取 5 个数：[0, 0.25, 0.5, 0.75, 1]
t8 = torch.eye(3)                              # 3x3 单位矩阵

# 与 NumPy 互相转换
import numpy as np
np_array = np.array([1, 2, 3])
t_from_np = torch.from_numpy(np_array)         # NumPy -> Tensor（共享内存）
np_from_t = t_from_np.numpy()                  # Tensor -> NumPy（共享内存）

2.2 张量属性

复制代码

t = torch.rand(3, 4)
print(t.shape)      # torch.Size([3, 4])，也可用 t.size()
print(t.dtype)       # torch.float32
print(t.device)      # cpu 或 cuda:0
print(t.ndim)        # 维度数：2
print(t.numel())     # 元素总数：12

2.3 索引、切片与形状变换

复制代码

t = torch.arange(12).reshape(3, 4)

print(t[1, 2])          # 索引单个元素
print(t[:, 1])           # 取第 1 列
print(t[0:2, :])         # 切片取前两行

t2 = t.reshape(4, 3)     # 改变形状（要求元素总数不变）
t3 = t.view(2, 6)        # view 与 reshape 类似，但要求张量在内存中连续
t4 = t.flatten()          # 展平为一维
t5 = t.unsqueeze(0)       # 在第 0 维增加一个维度，形状变为 (1, 3, 4)
t6 = t5.squeeze(0)        # 去掉大小为 1 的维度，还原为 (3, 4)
t7 = t.permute(1, 0)      # 交换维度顺序，转置
t8 = t.T                  # 二维张量的快捷转置写法

2.4 常用数学与聚合运算

复制代码

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])

c = a + b               # 逐元素加法，也可用 torch.add(a, b)
d = a * b                # 逐元素乘法（注意不是矩阵乘法）
e = torch.matmul(a, b)   # 向量点积 / 矩阵乘法，也可用 a @ b

t = torch.rand(3, 4)
t.sum()                  # 所有元素求和
t.sum(dim=0)             # 按第 0 维求和（结果形状为 (4,)）
t.mean()                 # 均值
t.max()                  # 最大值
t.argmax(dim=1)          # 每行最大值对应的索引，分类任务常用
t.std()                  # 标准差

2.5 拼接与堆叠

复制代码

a = torch.rand(2, 3)
b = torch.rand(2, 3)

torch.cat([a, b], dim=0)     # 沿已有维度拼接，形状变为 (4, 3)
torch.stack([a, b], dim=0)   # 沿新维度堆叠，形状变为 (2, 2, 3)

三、自动求导机制（Autograd）

PyTorch 的核心优势之一是 自动求导（Autograd） ：只要将 requires_grad=True 的张量参与运算，PyTorch 就会自动构建计算图，反向传播时自动计算梯度。

复制代码

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1

y.backward()          # 反向传播，计算 dy/dx
print(x.grad)          # tensor(7.)，因为 dy/dx = 2x + 3 = 2*2+3 = 7

3.1 多变量梯度

复制代码

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()

y.backward()
print(x.grad)    # tensor([2., 4., 6.])，即 dy/dxi = 2*xi

3.2 梯度累积与清零

复制代码

w = torch.tensor(1.0, requires_grad=True)

for i in range(3):
    loss = w ** 2
    loss.backward()
    print(w.grad)      # 每次都会累加：2.0, 4.0, 6.0
    w.grad.zero_()      # ⚠️ 必须手动清零，否则梯度会持续累加，训练循环中容易踩坑

3.3 禁用梯度计算（推理阶段常用）

复制代码

x = torch.tensor(2.0, requires_grad=True)

with torch.no_grad():           # 上下文内不构建计算图，节省显存、加速推理
    y = x ** 2
    print(y.requires_grad)       # False

# 或使用装饰器方式
@torch.no_grad()
def predict(model, x):
    return model(x)

# 将张量从计算图中分离，得到一个不追踪梯度的新张量
z = x.detach()

四、神经网络模块（nn.Module）

PyTorch 通过继承 nn.Module 来定义神经网络，需要实现 __init__（定义层）和 forward（定义前向传播逻辑）。

复制代码

import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNet(input_dim=10, hidden_dim=32, output_dim=2)
x = torch.rand(5, 10)     # batch_size=5, input_dim=10
output = model(x)          # 等价于 model.forward(x)，但推荐直接调用 model(x)
print(output.shape)         # torch.Size([5, 2])

4.1 查看模型结构与参数

复制代码

print(model)                          # 打印模型层级结构

for name, param in model.named_parameters():
    print(name, param.shape)          # 查看每层参数的名称与形状

total_params = sum(p.numel() for p in model.parameters())
print(f"模型总参数量: {total_params}")

4.2 nn.Sequential 快速搭建

复制代码

model = nn.Sequential(
    nn.Linear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 2)
)

💡 nn.Sequential 适合层与层之间是简单顺序连接的场景；如果需要跳跃连接（如 ResNet 的残差结构）、多分支等复杂结构，必须使用自定义 nn.Module 子类，在 forward 中手写连接逻辑。

五、常用层与激活函数

5.1 常用层

复制代码

nn.Linear(in_features, out_features)          # 全连接层
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)  # 二维卷积层
nn.MaxPool2d(kernel_size=2)                    # 最大池化
nn.BatchNorm2d(num_features)                   # 批归一化，加速收敛、缓解梯度问题
nn.Dropout(p=0.5)                              # 随机失活，缓解过拟合
nn.Embedding(num_embeddings, embedding_dim)    # 词嵌入层，NLP 任务常用
nn.LSTM(input_size, hidden_size, num_layers)   # 长短期记忆网络
nn.TransformerEncoderLayer(d_model, nhead)     # Transformer 编码层

5.2 常用激活函数

复制代码

nn.ReLU()          # 最常用，计算简单、缓解梯度消失
nn.LeakyReLU(0.1)  # 负区间保留小斜率，缓解"神经元死亡"问题
nn.Sigmoid()        # 输出 (0,1)，常用于二分类输出层
nn.Tanh()           # 输出 (-1,1)
nn.Softmax(dim=1)   # 多分类输出层，将输出转换为概率分布
nn.GELU()           # Transformer 类模型中常用的平滑激活函数

5.3 卷积神经网络示例（图像分类）

复制代码

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2),           # 输出尺寸减半

            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 128),   # 假设输入图像为 32x32
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

六、损失函数与优化器

6.1 常用损失函数

复制代码

criterion = nn.CrossEntropyLoss()   # 多分类任务，内部已包含 Softmax，标签为类别索引（非 one-hot）
criterion = nn.MSELoss()             # 均方误差，回归任务常用
criterion = nn.BCELoss()             # 二分类交叉熵，需配合 Sigmoid 使用
criterion = nn.BCEWithLogitsLoss()   # 内置 Sigmoid 的二分类交叉熵，数值更稳定，推荐优先使用
criterion = nn.L1Loss()              # 平均绝对误差

6.2 常用优化器

复制代码

import torch.optim as optim

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))
optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)  # 主流 Transformer 训练常用

优化器选择建议：

优化器	特点	适用场景
`SGD` + momentum	收敛稳定，泛化性能通常更好	CNN 图像分类等经典任务
`Adam`	收敛快，对学习率不敏感	大多数场景的默认首选
`AdamW`	修正了 Adam 权重衰减实现上的缺陷	Transformer、大模型训练首选

七、数据加载（Dataset / DataLoader）

PyTorch 通过 Dataset 定义数据的获取方式，通过 DataLoader 实现批处理、打乱、多进程加载。

7.1 自定义 Dataset

复制代码

from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

data = torch.randn(100, 10)
labels = torch.randint(0, 2, (100,))
dataset = MyDataset(data, labels)

dataloader = DataLoader(
    dataset,
    batch_size=16,
    shuffle=True,       # 训练时打乱数据顺序
    num_workers=4,       # 多进程加载数据，加速训练
    drop_last=False,     # 是否丢弃最后不足一个 batch 的数据
)

for batch_data, batch_labels in dataloader:
    print(batch_data.shape, batch_labels.shape)
    break

7.2 使用 torchvision 内置数据集与预处理

复制代码

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),                                  # 转为张量，并归一化到 [0,1]
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),  # 标准化
])

train_dataset = datasets.CIFAR10(
    root="./data", train=True, download=True, transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

八、完整训练流程实战

一个标准的 PyTorch 训练流程通常遵循固定套路，建议牢记以下结构：

复制代码

import torch
import torch.nn as nn
import torch.optim as optim

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

model = SimpleNet(10, 32, 2).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
for epoch in range(num_epochs):
    model.train()               # 切换到训练模式（启用 Dropout、BatchNorm 使用 batch 统计量）
    total_loss = 0.0

    for batch_x, batch_y in dataloader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)

        optimizer.zero_grad()             # 1. 清空上一轮梯度
        outputs = model(batch_x)           # 2. 前向传播
        loss = criterion(outputs, batch_y) # 3. 计算损失
        loss.backward()                    # 4. 反向传播，计算梯度
        optimizer.step()                   # 5. 更新参数

        total_loss += loss.item()

    avg_loss = total_loss / len(dataloader)
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}")

    # 验证阶段
    model.eval()                 # 切换到评估模式
    correct, total = 0, 0
    with torch.no_grad():         # 验证阶段不需要计算梯度
        for val_x, val_y in val_loader:
            val_x, val_y = val_x.to(device), val_y.to(device)
            outputs = model(val_x)
            preds = outputs.argmax(dim=1)
            correct += (preds == val_y).sum().item()
            total += val_y.size(0)
    print(f"验证集准确率: {correct/total:.4f}")

⚠️ 易错点总结：

每个 batch 训练前必须 optimizer.zero_grad()，否则梯度会累加导致训练异常。

model.train() / model.eval() 一定要在对应阶段调用，否则 Dropout、BatchNorm 的行为会不一致，导致验证结果异常。

数据和模型必须在 同一设备（CPU 或同一张 GPU）上，否则会报错。

九、GPU 加速与设备管理

复制代码

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

model = model.to(device)          # 将模型参数迁移到指定设备
x = x.to(device)                   # 将张量迁移到指定设备

# 多 GPU 数据并行（单机多卡场景）
model = nn.DataParallel(model)     # 简单场景可用，生产环境推荐 DistributedDataParallel

# 查看显存占用情况
print(torch.cuda.memory_allocated() / 1024**2, "MB")
print(torch.cuda.memory_reserved() / 1024**2, "MB")
torch.cuda.empty_cache()            # 释放未使用的缓存显存（不会影响正在使用的显存）

9.1 混合精度训练（加速 + 省显存）

复制代码

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch_x, batch_y in dataloader:
    batch_x, batch_y = batch_x.to(device), batch_y.to(device)
    optimizer.zero_grad()

    with autocast():                          # 自动混合精度上下文，部分算子使用 fp16 计算
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)

    scaler.scale(loss).backward()               # 缩放损失，避免 fp16 梯度下溢
    scaler.step(optimizer)
    scaler.update()

十、模型保存与加载

10.1 只保存参数（推荐方式）

复制代码

torch.save(model.state_dict(), "model_weights.pth")

# 加载时需要先构建相同结构的模型，再加载参数
model = SimpleNet(10, 32, 2)
model.load_state_dict(torch.load("model_weights.pth", map_location=device))
model.eval()

10.2 保存整个模型（不推荐，兼容性较差）

复制代码

torch.save(model, "full_model.pth")
model = torch.load("full_model.pth", weights_only=False)

10.3 保存训练检查点（Checkpoint）

复制代码

checkpoint = {
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": loss,
}
torch.save(checkpoint, "checkpoint.pth")

checkpoint = torch.load("checkpoint.pth")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
start_epoch = checkpoint["epoch"]

💡 官方推荐 ：优先使用 state_dict() 只保存参数，因为它与具体的类定义解耦，更利于版本管理和跨环境迁移；torch.save(model, ...) 保存整个模型对象会绑定具体的类路径，代码结构变化后可能加载失败。

十一、学习率调度器

复制代码

from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR, ReduceLROnPlateau

scheduler = StepLR(optimizer, step_size=10, gamma=0.1)      # 每 10 轮学习率乘以 0.1
scheduler = CosineAnnealingLR(optimizer, T_max=50)            # 余弦退火，训练常用
scheduler = ReduceLROnPlateau(optimizer, mode="min", patience=5)  # 验证指标停滞时自动降低学习率

for epoch in range(num_epochs):
    train_one_epoch(...)
    scheduler.step()               # 大多数调度器在每个 epoch 后调用一次
    # ReduceLROnPlateau 需要传入监控指标：scheduler.step(val_loss)

十二、常用进阶技巧

12.1 梯度裁剪（防止梯度爆炸）

复制代码

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

12.2 固定随机种子（保证实验可复现）

复制代码

import random
import numpy as np

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed(42)

12.3 冻结部分参数（迁移学习常用）

复制代码

for param in model.features.parameters():
    param.requires_grad = False    # 冻结特征提取层，只训练分类头

optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)

12.4 张量维度检查小技巧

复制代码

assert x.shape == (batch_size, 10), f"输入形状不符，实际为 {x.shape}"

十三、torch.compile 加速

PyTorch 2.x 引入的 torch.compile 可以在几乎不改变代码的情况下，通过图捕获与算子融合显著提升训练/推理速度：

复制代码

model = SimpleNet(10, 32, 2).to(device)
compiled_model = torch.compile(model)     # 一行代码即可启用编译加速

outputs = compiled_model(x)                # 用法与原模型完全一致

💡 首次调用会有编译预热开销，适合训练轮数较多、输入形状相对固定的场景；如果输入形状频繁变化，编译带来的收益可能被反复重新编译的开销抵消。

十四、常见问题排查

报错/现象	常见原因	解决方法
`RuntimeError: Expected all tensors to be on the same device`	模型和数据不在同一设备	检查 `.to(device)` 是否对模型和所有输入张量都执行了
损失值变成 `NaN`	学习率过大 / 数值不稳定	减小学习率、使用梯度裁剪、检查数据是否存在异常值
训练损失不下降	忘记 `zero_grad()`、学习率设置不当	检查训练循环是否遗漏关键步骤
验证集表现异常	忘记调用 `model.eval()`	验证/推理前务必切换模式
显存溢出（OOM）	batch size 过大 / 未释放中间变量	减小 batch size、使用混合精度训练、及时 `del` 不再使用的张量
`view` 报错 tensor 不连续	张量在内存中不连续（如经过 transpose）	改用 `reshape()`，或先调用 `.contiguous()`

十五、总结速查表

场景	核心 API
安装与验证	`torch.__version__`、`torch.cuda.is_available()`
张量创建	`torch.tensor` `torch.zeros/ones/rand/randn`
形状操作	`reshape` `view` `permute` `squeeze/unsqueeze`
自动求导	`requires_grad=True` `backward()` `torch.no_grad()`
模型搭建	`nn.Module` `nn.Sequential` `nn.Linear/Conv2d`
损失与优化	`nn.CrossEntropyLoss` `optim.Adam/AdamW`
数据加载	`Dataset` `DataLoader`
训练流程	`zero_grad → forward → loss → backward → step`
GPU 加速	`.to(device)` `autocast` `GradScaler`
模型保存	`state_dict()` `torch.save/load`
学习率调度	`StepLR` `CosineAnnealingLR`
加速编译	`torch.compile`

PyTorch 的核心思想始终围绕 张量 + 自动求导 + nn.Module 三大基石展开，掌握了标准训练流程后，无论是 CNN、RNN 还是 Transformer 类模型，代码结构都是相通的。建议结合官方文档（pytorch.org/docs）动手实践，加深理解。

📎 如果本文对你有帮助，欢迎点赞收藏，也欢迎在评论区分享你在 PyTorch 使用过程中踩过的坑！