PyTorch深度学习入门：从张量、自动微分到完整训练流程实战

1. 项目概述：为什么PyTorch是深度学习的“瑞士军刀”？

如果你刚开始接触深度学习，面对TensorFlow、PyTorch、JAX这些框架可能会有点懵。我当年也一样，每个框架都说自己好，但真正上手做项目、搞研究，尤其是需要快速验证想法、灵活调试模型的时候，PyTorch几乎成了我和身边大多数同事的首选。它不像一个庞大而精密的工业流水线，更像一把趁手的“瑞士军刀”——设计直观，用起来灵活，出了问题也容易找到症结所在。这个“基础1”，我们就来聊聊怎么把这把刀从刀鞘里抽出来，磨锋利，并完成第一次切削。这不是一份照本宣科的API文档，而是我踩过无数坑后，总结出的“生存指南”。无论你是想复现一篇顶会论文，还是为自己手头的图像、文本数据搭建第一个神经网络，从这里开始，都能少走很多弯路。

PyTorch的核心魅力在于它的“动态计算图”和“Python原生”的风格。简单说，它让你用写Python程序的方式去构建神经网络，每一步计算都是实时发生的，你可以随时用print()查看张量的值，用Python的调试器设置断点。这种即时反馈的体验，对于学习和研究来说是无价的。我们将从最核心的数据结构——张量（Tensor）开始，逐步搭建起数据加载、模型定义、训练循环的完整流程。你会发现，原来深度学习框架的基础，并没有想象中那么遥不可及。

2. 核心基石：透彻理解张量（Tensor）与自动微分

2.1 张量：不止是多维数组

很多人把张量简单理解为N维数组，这没错，但理解不到位。在PyTorch的语境下，张量是承载数据和计算的核心载体，它自带两个至关重要的属性：device（设备）和requires_grad（是否需要梯度）。

设备（device）：这决定了你的张量是在CPU上还是GPU上。深度学习计算量巨大，GPU的并行计算能力是关键。一个常见的低级错误就是，模型在GPU上，而数据却在CPU上，导致运行时错误。创建张量时就要有意识地进行管理。

import torch # 在CPU上创建张量 cpu_tensor = torch.tensor([1, 2, 3]) print(cpu_tensor.device) # 输出：cpu # 在GPU上创建张量（如果有可用的CUDA设备） if torch.cuda.is_available(): gpu_tensor = torch.tensor([1, 2, 3], device='cuda') # 或者更常用的方式：创建后移动到GPU gpu_tensor = cpu_tensor.to('cuda') print(gpu_tensor.device) # 输出：cuda:0

梯度（requires_grad）：这是PyTorch自动微分（Autograd）系统的入口。当你将张量的requires_grad属性设置为True时，PyTorch就会开始跟踪在其上执行的所有操作，构建一个动态计算图。在反向传播时，系统会根据这个图自动计算梯度。这对于模型参数来说是必须的。

x = torch.tensor([1.0, 2.0], requires_grad=True) y = x ** 2 + 2 * x + 1 print(y) # tensor([4., 9.], grad_fn=<AddBackward0>) # 注意看，y有一个grad_fn属性，它指向了创建y的操作（加法），这就是计算图的节点。

注意：只有浮点型和复数类型的张量才能设置requires_grad=True。整数类型是不支持自动求导的。

2.2 自动微分（Autograd）实战：从原理到调试

理解了requires_grad，我们来实际完成一次反向传播。假设我们有一个简单的函数z = x * y + sin(x)，我们需要求z对x和y的梯度。

# 1. 创建叶子节点（leaf tensor），并开启梯度追踪 x = torch.tensor(2.0, requires_grad=True) y = torch.tensor(3.0, requires_grad=True) # 2. 执行前向计算，构建计算图 z = x * y + torch.sin(x) print(f'z = {z}') # z = 6.9093 # 3. 执行反向传播，计算梯度 z.backward() # 等价于 z.backward(torch.tensor(1.0))，因为z是标量 # 4. 查看梯度 print(f'∂z/∂x = {x.grad}') # ∂z/∂x = y + cos(x) = 3 + cos(2) ≈ 3 - 0.416 = 2.584 print(f'∂z/∂y = {y.grad}') # ∂z/∂y = x = 2

这里有几个极易出错的点：

1. 梯度累加：默认情况下，张量的.grad属性会累加每次backward()计算的梯度。所以在每次参数更新前，必须手动将梯度清零（optimizer.zero_grad()）。
2. 非标量输出反向传播：如果z不是标量（例如是一个向量），调用z.backward()需要传入一个与z形状相同的“梯度权重”向量，作为链式法则的起始点。绝大多数情况下，我们的损失函数是标量，所以直接调用loss.backward()即可。
3. 中断梯度追踪：有时我们不需要某些计算被追踪，比如在模型评估或冻结部分层时。可以用with torch.no_grad():上下文管理器，或者对张量调用.detach()方法。

# 中断梯度追踪的两种方式 x = torch.tensor(2.0, requires_grad=True) # 方法一：.detach()，返回一个共享数据但无梯度历史的新张量 y = x.detach() ** 2 # y的运算不会被追踪 z = x + y z.backward() print(x.grad) # 输出：1.0，因为y被视为常数，只对x求导 # 方法二：torch.no_grad() 上下文，块内所有计算无追踪 with torch.no_grad(): w = x * 10 # w的requires_grad为False

3. 数据流水线：用Dataset和DataLoader构建高效数据流

模型和算法再漂亮，没有高质量、高效率的数据输入也是白搭。PyTorch用torch.utils.data.Dataset和DataLoader这两个抽象，把数据处理的脏活累活封装得明明白白。

3.1 自定义Dataset：从原始数据到模型可读的样本

Dataset是一个抽象类，你必须继承它并实现三个方法：__init__,__len__,__getitem__。我习惯把它想象成一个“智能数据索引器”。

假设我们有一个简单的任务：读取一个文件夹下的所有图片（.jpg文件）并进行分类，标签写在文件名里（如cat_001.jpg,dog_005.jpg）。

import os from PIL import Image import torch from torch.utils.data import Dataset from torchvision import transforms class AnimalDataset(Dataset): def __init__(self, root_dir, transform=None): """ Args: root_dir (string): 图片根目录。 transform (callable, optional): 一个可选的图像变换函数。 """ self.root_dir = root_dir self.transform = transform self.classes = ['cat', 'dog'] # 类别列表 self.class_to_idx = {c: i for i, c in enumerate(self.classes)} # 类别名到索引的映射 # 收集所有文件路径和标签 self.samples = [] for filename in os.listdir(root_dir): if filename.endswith('.jpg'): # 假设文件名格式为“类别_编号.jpg” label_str = filename.split('_')[0] if label_str in self.class_to_idx: filepath = os.path.join(root_dir, filename) label = self.class_to_idx[label_str] self.samples.append((filepath, label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): filepath, label = self.samples[idx] # 1. 加载图像 image = Image.open(filepath).convert('RGB') # 确保是三通道 # 2. 应用变换（数据增强和归一化） if self.transform: image = self.transform(image) # 3. 将标签也转为Tensor（通常是int64） label = torch.tensor(label, dtype=torch.long) return image, label

关键细节与避坑：

• __getitem__返回的应该是(data, target)的元组形式，这是PyTorch社区的约定俗成。
• 图像变换（transform）非常重要。通常包括：调整大小（Resize）、随机裁剪（RandomCrop）、随机水平翻转（RandomHorizontalFlip）等数据增强操作，以及将PIL图像或NumPy数组转换为张量（ToTensor）和归一化（Normalize）。务必注意：ToTensor()会将像素值从[0, 255]的整数范围转换到[0.0, 1.0]的浮点数范围，并调整维度顺序为(C, H, W)。
• 标签通常转换为torch.long类型，因为后续的交叉熵损失函数（CrossEntropyLoss）期望这样的输入。

3.2 DataLoader：批处理、打乱与并行加载的艺术

Dataset负责提供单个样本，DataLoader则负责把这些样本打包成批（batch），并提供打乱（shuffle）、并行加载（multiprocessing）等高级功能。

from torch.utils.data import DataLoader # 定义图像变换管道 transform = transforms.Compose([ transforms.Resize((224, 224)), # 调整大小 transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转（数据增强） transforms.ToTensor(), # 转为Tensor，并归一化到[0,1] transforms.Normalize(mean=[0.485, 0.456, 0.406], # ImageNet均值 std=[0.229, 0.224, 0.225]) # ImageNet标准差 ]) # 实例化Dataset dataset = AnimalDataset(root_dir='./animals', transform=transform) # 创建DataLoader dataloader = DataLoader( dataset, batch_size=32, # 批大小 shuffle=True, # 每个epoch开始时打乱数据 num_workers=4, # 用于数据加载的子进程数 pin_memory=True, # 如果使用GPU，将数据锁页内存可加速CPU到GPU的传输 drop_last=False # 是否丢弃最后一个不完整的batch ) # 使用方式：像一个可迭代对象 for batch_idx, (images, labels) in enumerate(dataloader): # images的形状: [batch_size, 3, 224, 224] # labels的形状: [batch_size] print(f'Batch {batch_idx}, Images shape: {images.shape}, Labels shape: {labels.shape}') # 这里可以将images和labels送入模型进行前向传播 break # 仅演示，跳出循环

num_workers和pin_memory的实战经验：

• num_workers：这个参数设置多少个并行进程来加载数据。原则是，在CPU和内存允许的情况下，设置得越高，数据加载越快，GPU等待数据的时间就越短。一个经验法则是设置为CPU核心数。但要注意，如果设置过高，可能会因进程间通信开销反而变慢，甚至导致内存不足。调试时，建议先设为0，确保代码逻辑正确，再尝试调大。
• pin_memory：当使用GPU时，将其设为True可以显著加速数据从CPU内存到GPU显存的传输。原理是它使用了“锁页内存”，GPU可以直接通过DMA访问，省去了复制到可分页内存的步骤。只要你的主机内存足够，在GPU训练时永远应该打开它。

4. 模型构建：用nn.Module搭建你的神经网络

PyTorch中所有的神经网络模块都继承自torch.nn.Module。你要做的就是像搭积木一样，在__init__中声明你用到的“积木”（层），在forward方法中定义这些“积木”是如何连接起来处理输入数据的。

4.1 定义一个简单的卷积神经网络（CNN）

我们以经典的LeNet-5结构为蓝本，构建一个用于图像分类的CNN。

import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes=10): super(SimpleCNN, self).__init__() # 特征提取器 self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1) self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1) self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) # 分类器 # 经过两次池化，图像尺寸变为 H/4, W/4。假设输入是32x32，则这里是8x8 self.fc1 = nn.Linear(in_features=32 * 8 * 8, out_features=128) self.fc2 = nn.Linear(in_features=128, out_features=num_classes) # 可选：Dropout层用于防止过拟合 self.dropout = nn.Dropout(p=0.5) def forward(self, x): # x 形状: [batch_size, 3, H, W] # 卷积 -> 激活 -> 池化 x = self.pool1(F.relu(self.conv1(x))) # [batch_size, 16, H/2, W/2] x = self.pool2(F.relu(self.conv2(x))) # [batch_size, 32, H/4, W/4] # 展平操作：将特征图拉成一维向量 x = x.view(x.size(0), -1) # [batch_size, 32 * (H/4) * (W/4)] # 全连接层 x = F.relu(self.fc1(x)) x = self.dropout(x) # 只在训练时生效 x = self.fc2(x) # 输出层通常不加激活函数，损失函数会处理 return x

必须理解的几个要点：

1. super().__init__()：这行代码必须放在__init__方法的第一行，它调用了父类nn.Module的初始化方法，这是PyTorch能够管理你定义的层（如self.conv1）的关键。
2. forward是定义计算流程：你只需要定义前向传播。反向传播的计算图是由Autograd根据forward中的操作自动构建的。永远不要直接调用model.forward(x)，而是使用model(x)。因为model(x)内部会先调用model.forward(x)，但还会处理一些其他的钩子（hooks）和检查。
3. 展平操作（Flatten）：卷积层输出是4D张量[batch, channel, height, width]，而全连接层输入需要2D张量[batch, features]。x.view(x.size(0), -1)是经典的展平方式，-1表示让PyTorch自动计算该维度的大小。
4. 激活函数的选择：F.relu是函数式接口，nn.ReLU()是模块化接口。在forward中直接使用F.relu更简洁。对于有可学习参数的层（如Dropout,BatchNorm），必须使用模块形式（self.dropout），因为它们在训练和评估模式下的行为不同。

4.2 模型参数初始化与状态管理

默认情况下，PyTorch的层使用一种启发式方法初始化参数。但对于深层网络，不恰当的初始化可能导致梯度消失或爆炸。我们可以自定义初始化。

def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) # Xavier初始化，适用于tanh/sigmoid nn.init.constant_(m.bias, 0) # 偏置初始化为0 elif isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') # Kaiming初始化，适用于ReLU if m.bias is not None: nn.init.constant_(m.bias, 0) model = SimpleCNN(num_classes=2) model.apply(init_weights) # 递归地将init_weights函数应用到所有子模块

模型模式切换：这是新手常忽略但至关重要的点。像Dropout和BatchNorm这样的层在训练和评估（推理）时行为不同。

• model.train()：将模型设置为训练模式。Dropout会随机丢弃神经元，BatchNorm会使用当前批次的统计量。
• model.eval()：将模型设置为评估模式。Dropout不生效，BatchNorm会使用训练阶段累积的运行均值/方差。

# 训练循环开始前 model.train() for data, target in train_loader: # ... 训练步骤 # 在验证或测试时 model.eval() with torch.no_grad(): # 同时关闭梯度计算，节省内存和计算 for data, target in val_loader: output = model(data) # ... 计算验证指标

5. 训练循环：将损失、优化器和迭代组合起来

训练一个神经网络本质上是一个优化问题：找到一组模型参数，使得损失函数在训练数据上的值最小。这个过程通过“前向传播 -> 计算损失 -> 反向传播 -> 参数更新”的循环来实现。

5.1 核心组件拆解：损失函数与优化器

损失函数（Loss Function）：衡量模型预测与真实标签的差距。选择取决于任务。

• 分类任务：常用nn.CrossEntropyLoss。它内部集成了Softmax，所以你的模型最后一层不需要再加nn.Softmax。输入是原始分数（logits），目标是与类别索引。
• 回归任务：常用nn.MSELoss（均方误差）或nn.L1Loss（平均绝对误差）。
• 二分类任务：可以用nn.BCEWithLogitsLoss（内置Sigmoid，数值稳定），同样模型最后一层不需要Sigmoid。

优化器（Optimizer）：负责根据梯度更新模型参数。torch.optim提供了多种选择。

• optim.SGD：随机梯度下降，是很多优化器的基础。配合动量（momentum）使用效果很好。
• optim.Adam：自适应学习率算法，通常被认为是默认的、效果不错的优化器，尤其适合初学者。

import torch.optim as optim model = SimpleCNN(num_classes=2).to('cuda') # 将模型移到GPU criterion = nn.CrossEntropyLoss() # 损失函数 optimizer = optim.Adam(model.parameters(), lr=0.001) # 优化器，传入模型参数和学习率 # 学习率调度器（可选但推荐）：在训练过程中动态调整学习率 scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # 每10个epoch学习率乘以0.1

5.2 手把手编写第一个训练循环

下面是一个完整的、带有验证环节的训练循环模板，包含了所有关键步骤和注释。

def train_one_epoch(model, train_loader, criterion, optimizer, device): """训练一个epoch""" model.train() # 设置为训练模式 running_loss = 0.0 correct = 0 total = 0 for batch_idx, (inputs, labels) in enumerate(train_loader): # 1. 数据迁移到设备 inputs, labels = inputs.to(device), labels.to(device) # 2. 梯度清零！！！这是必须的，否则梯度会累加 optimizer.zero_grad() # 3. 前向传播 outputs = model(inputs) loss = criterion(outputs, labels) # 4. 反向传播 loss.backward() # 5. 参数更新 optimizer.step() # 6. 统计信息（用于打印日志） running_loss += loss.item() * inputs.size(0) # loss.item()是标量，乘以batch size得到本批总损失 _, predicted = outputs.max(1) # 获取预测类别（最大值的索引） total += labels.size(0) correct += predicted.eq(labels).sum().item() # 可选：每N个batch打印一次进度 if batch_idx % 50 == 0: print(f' Batch [{batch_idx}/{len(train_loader)}], Loss: {loss.item():.4f}') epoch_loss = running_loss / total epoch_acc = 100. * correct / total return epoch_loss, epoch_acc def validate(model, val_loader, criterion, device): """验证/测试""" model.eval() # 设置为评估模式 running_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): # 关键：关闭梯度计算，节省内存和计算 for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) running_loss += loss.item() * inputs.size(0) _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() epoch_loss = running_loss / total epoch_acc = 100. * correct / total return epoch_loss, epoch_acc # 主训练循环 num_epochs = 20 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) train_losses, val_losses = [], [] train_accs, val_accs = [], [] for epoch in range(num_epochs): print(f'Epoch {epoch+1}/{num_epochs}') # 训练阶段 train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device) train_losses.append(train_loss) train_accs.append(train_acc) # 验证阶段 val_loss, val_acc = validate(model, val_loader, criterion, device) val_losses.append(val_loss) val_accs.append(val_acc) # 更新学习率（如果使用了调度器） scheduler.step() print(f' Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%') print(f' Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%') print('-' * 60)

这段代码里的“坑”与经验：

1. optimizer.zero_grad()的位置：必须在每个batch的loss.backward()之前调用。如果忘记，梯度会不断累加，导致模型更新方向错误，无法收敛。
2. loss.item()vsloss：loss是一个包含计算图的张量，而loss.item()是获取其Python标量值。在统计和打印时使用.item()，避免不必要的计算图留存，节省内存。
3. model.eval()与torch.no_grad()：在验证/测试时，这两者通常一起使用。model.eval()改变某些层（如Dropout）的行为；torch.no_grad()则禁用梯度计算，是更大的性能优化。
4. 设备管理：务必确保模型、输入数据、标签都在同一个设备上（CPU或GPU）。.to(device)是标准做法。

6. 模型保存、加载与推理部署

训练好的模型需要保存下来，以便后续评估、继续训练或部署到生产环境。

6.1 保存与加载的两种方式

PyTorch推荐使用.pt或.pth作为模型文件后缀。有两种主要的保存方式：

方式一：保存整个模型（包括结构和参数）

# 保存 torch.save(model, 'model_complete.pth') # 加载（需要模型类定义在当前位置可访问） model_loaded = torch.load('model_complete.pth', map_location=device) # map_location指定加载设备 model_loaded.eval()

• 优点：简单直接，连模型结构一起保存。
• 缺点：1) 文件较大；2) 加载依赖于原始的模型类定义，如果类定义有修改，可能无法加载；3) 在不同PyTorch版本间可能不兼容。不推荐作为长期保存或分享的方式。

方式二：仅保存模型状态字典（推荐）

# 保存 torch.save(model.state_dict(), 'model_state_dict.pth') # 加载 model = SimpleCNN(num_classes=2) # 必须先实例化一个结构相同的模型 model.load_state_dict(torch.load('model_state_dict.pth', map_location=device)) model.to(device) model.eval()

• 优点：文件小，只包含参数。加载灵活，只要模型结构一致即可，与类定义方式解耦。是最推荐的方式。
• 注意：load_state_dict要求保存和加载的模型结构严格一致（层名、参数形状）。

6.2 进行单样本推理

将模型用于实际预测，需要注意将输入数据预处理成与训练时相同的格式。

def predict_single_image(image_path, model, transform, device, class_names): """ 对单张图片进行预测 Args: image_path: 图片路径 model: 加载好的模型 transform: 与训练时相同的图像变换 device: CPU或GPU class_names: 类别名称列表 """ model.eval() # 1. 加载和预处理图像 image = Image.open(image_path).convert('RGB') image_tensor = transform(image).unsqueeze(0) # 增加一个batch维度 -> [1, C, H, W] image_tensor = image_tensor.to(device) # 2. 前向传播 with torch.no_grad(): outputs = model(image_tensor) # 获取概率（应用Softmax） probabilities = F.softmax(outputs, dim=1) # 获取预测类别和置信度 confidence, predicted_class = torch.max(probabilities, 1) # 3. 返回结果 predicted_label = class_names[predicted_class.item()] confidence_score = confidence.item() return predicted_label, confidence_score # 使用示例 pred_label, conf = predict_single_image( image_path='./test_dog.jpg', model=model, transform=transform, # 使用与训练时完全相同的transform device=device, class_names=['cat', 'dog'] ) print(f'预测结果: {pred_label}, 置信度: {conf:.2%}')

推理时的关键细节：

1. unsqueeze(0)：模型输入通常期望是4D张量[batch, C, H, W]。单张图片是3D的[C, H, W]，需要用unsqueeze(0)在第0维（批次维）增加一个维度。
2. 相同的变换：务必使用与训练时完全相同的transform管道，特别是归一化所用的均值和标准差。如果预处理不一致，模型性能会严重下降。
3. torch.no_grad()和model.eval()：在推理时同样重要，确保计算高效且某些层行为正确。

7. 调试与性能优化实战指南

掌握了基础流程后，你会遇到模型不收敛、精度低、速度慢等问题。这里分享几个最实用的调试和优化技巧。

7.1 模型不收敛？先检查数据与损失

1. 数据检查：这是第一位的。用几行代码可视化一下你的输入数据和标签。

   # 从DataLoader中取一个batch看看 images, labels = next(iter(train_loader)) print(f'Batch图像范围: [{images.min():.3f}, {images.max():.3f}]') # 应该是接近[0,1]或归一化后的值 print(f'标签示例: {labels[:10]}') # 检查标签是否在有效范围内 # 可以显示几张图片看看

   常见问题：图像未归一化、标签错乱、数据增强过于激进导致图片无法辨认。

2. 损失值变化：训练初期，观察损失值。

   • 如果损失纹丝不动：可能是学习率太小、梯度消失（检查初始化、激活函数）、优化器参数未正确传入（model.parameters()）。
   • 如果损失变成NaN：可能是学习率太大导致梯度爆炸、数据中存在非法值（如NaN或inf）、损失函数对输入敏感（如交叉熵输入了非法概率）。
   • 一个技巧：先用一个极小的数据集（比如2-3个batch）跑一下，看损失是否能快速下降（过拟合这个小数据集）。如果能，说明模型有能力学习，问题可能出在数据或超参数上。

3. 梯度检查：在训练循环中插入代码，打印某些层权重的梯度范数。

   # 在loss.backward()之后，optimizer.step()之前 total_norm = 0 for p in model.parameters(): if p.grad is not None: param_norm = p.grad.data.norm(2) # L2范数 total_norm += param_norm.item() ** 2 total_norm = total_norm ** 0.5 print(f'梯度范数: {total_norm}')

   梯度范数非常小（如1e-7）可能意味着梯度消失；非常大（如1e+10）则意味着梯度爆炸。

7.2 性能优化：让训练飞起来

1. 使用torch.backends.cudnn.benchmark = True：如果你的网络结构是固定的（输入尺寸不变），在程序开头设置这个标志，可以让CuDNN自动寻找最适合你硬件和网络的最优卷积算法，显著提升训练速度。

   if torch.cuda.is_available(): torch.backends.cudnn.benchmark = True

2. 混合精度训练（AMP）：这是现代GPU（Volta架构及以后）的一大福利。使用半精度（float16）进行计算，可以几乎不减精度的情况下，大幅减少显存占用并提升训练速度。PyTorch提供了自动混合精度包torch.cuda.amp。

   from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() # 梯度缩放，防止半精度下的梯度下溢 for data, target in train_loader: optimizer.zero_grad() with autocast(): # 自动为操作选择半精度或单精度 output = model(data) loss = criterion(output, target) scaler.scale(loss).backward() # 缩放损失，反向传播 scaler.step(optimizer) # 缩放梯度，更新参数 scaler.update() # 更新缩放因子

   使用AMP通常可以获得1.5-2倍的速度提升，并减少近一半的显存占用。

3. 梯度累加（Gradient Accumulation）：当你的GPU显存不足以容纳你期望的大batch size时，这是一个非常有用的技巧。它通过多次前向传播累积梯度，模拟大batch的效果，然后再进行一次参数更新。

   accumulation_steps = 4 # 累积4个batch的梯度 optimizer.zero_grad() for i, (data, target) in enumerate(train_loader): output = model(data) loss = criterion(output, target) loss = loss / accumulation_steps # 损失按累积步数缩放 loss.backward() # 梯度累积 if (i + 1) % accumulation_steps == 0: optimizer.step() # 执行参数更新 optimizer.zero_grad() # 清零梯度

   这样，虽然实际每次更新用的数据量是batch_size * accumulation_steps，但显存占用只相当于一个batch_size。

从理解张量和自动微分，到构建数据管道、定义模型、编写训练循环，再到保存模型和实战调试，这套流程构成了PyTorch使用的完整闭环。我个人的体会是，最初几次照着这个流程走可能会觉得步骤繁多，但一旦跑通几次，它就会变成肌肉记忆。最关键的是养成好习惯：比如始终管理好设备和梯度，在eval模式时加上torch.no_grad，保存模型时优先用state_dict。遇到问题，多从数据本身和损失曲线入手排查，往往比盲目调整模型结构更有效。PyTorch的灵活性是把双刃剑，它给了你极大的自由，但也要求你对这些底层细节有清晰的掌控。希望这份“基础1”指南，能帮你打下扎实的根基，在后续更复杂的项目里游刃有余。