这次我们来看 PyTorch 中一个核心且必须掌握的技能如何正确地封装一个完整的深度学习模型。这不仅仅是把网络层堆叠起来而是构建一个包含前向传播、损失计算和参数更新的可训练、可复用单元。对于任何想在 PyTorch 框架下进行有效开发的工程师或研究者来说理解nn.Module、损失函数和优化器的协同工作方式是绕不开的一步。很多教程会分开讲解这些概念但实际项目中它们是一个密不可分的整体。一个封装良好的模型能让你在调试、训练、保存和部署时省去大量麻烦。本文将聚焦于如何从零开始构建一个结构清晰、功能完整的模型类并深入探讨损失函数与优化器的选择与配置最后通过一个完整的训练循环示例让你彻底掌握这套“组合拳”。本文适合已经了解 PyTorch 张量基础操作和简单网络层定义希望系统学习模型构建与训练流程的读者。我们将重点关注以下几个实操点nn.Module的正确继承与结构设计如何组织网络层、定义前向传播并管理子模块。损失函数的选择与调用针对不同任务如分类、回归如何选取合适的损失函数并理解其输入输出。优化器的配置与使用如何将模型参数传递给优化器并设置学习率等关键超参数。完整的训练循环搭建将前三点串联起来形成一个可运行的训练流程并观察损失下降。通过本文你将能独立封装出自己的第一个 PyTorch 模型并理解训练过程中数据是如何流动、损失是如何计算、参数是如何被更新的。1. 核心能力速览模型封装三要素在深入代码之前我们先通过一个表格快速把握本次内容的核心组成部分及其作用这相当于你模型工具箱里的“三大件”。组件核心类/模块主要职责关键特点模型容器torch.nn.Module1. 定义网络结构层、子模块2. 实现前向传播逻辑 (forward)3. 自动管理可训练参数 (parameters())4. 支持模型保存与加载 (state_dict)所有神经网络模块的基类。通过继承它来构建自定义模型。其forward方法定义了数据如何通过网络。损失衡量torch.nn中的各种损失函数 (如MSELoss,CrossEntropyLoss)1. 计算模型预测输出与真实标签之间的差异损失值2. 为优化器提供梯度计算的方向损失值是一个标量值越小通常表示模型预测越准。不同的任务对应不同的损失函数。参数优化torch.optim中的各种优化器 (如SGD,Adam)1. 持有需要更新的模型参数引用2. 根据损失函数计算出的梯度按照特定算法更新参数以最小化损失优化器决定了参数更新的策略如步长、动量学习率是其最关键的超参数。协同工作流程在训练循环的每次迭代中数据经过模型容器(forward) 得到预测值预测值与真实值送入损失函数计算出损失调用损失值的backward()方法自动计算所有模型参数的梯度最后优化器执行step()方法利用这些梯度更新模型参数。optimizer.zero_grad()用于在每次更新前清空上一轮的梯度防止累积。2. 适用场景与使用边界适合谁解决什么问题深度学习初学者希望系统理解 PyTorch 训练流程而非仅仅复制代码。算法工程师/研究员需要快速原型验证新网络结构nn.Module提供了灵活的构建方式。需要模型复现与共享的开发者封装好的模型易于保存 (torch.save)、加载和分享。进行模型调试与分析的从业者良好的封装便于插入钩子 (hooks) 来监控中间层激活或梯度。核心解决的问题代码组织与复用将网络结构、损失计算、优化逻辑模块化避免代码冗长混乱。参数自动化管理nn.Module自动追踪所有通过nn.Parameter或子模块注册的参数无需手动维护参数列表。训练流程标准化掌握损失函数-反向传播-优化器更新这一标准范式是进行任何模型训练的基础。设备无缝迁移调用model.to(device)即可将模型所有参数移至 GPU 或 CPU封装简化了设备管理。不适合什么场景极度追求推理速度的轻量级部署最终部署时可能会将 PyTorch 模型转换为如 TorchScript、ONNX 等其他格式原始的nn.Module封装可能不是最终形态但仍是转换的起点。仅使用预训练模型进行微调 (Fine-tuning)虽然也涉及优化器但重点在于加载预训练权重和修改最后一层模型封装本身通常已由库如torchvision.models提供。非神经网络的传统机器学习模型如决策树、SVM 等其训练范式与基于梯度的深度学习不同无需使用nn.Module和torch.optim。安全与合规边界本文涉及的模型封装与训练是通用的深度学习编程实践不涉及特定数据、隐私或版权内容。但在实际应用中需注意数据合规确保用于训练的数据集已获得合法授权不包含个人隐私信息或受版权保护的内容。模型用途使用训练好的模型进行预测或生成时需符合法律法规和伦理准则特别是用于人脸识别、内容生成等敏感领域时。代码借鉴学习开源项目代码时需遵守其对应的开源协议如 MIT、Apache-2.0。3. 环境准备与前置条件在开始动手编码前请确保你的开发环境已就绪。以下是基于 PyTorch 模型开发的通用环境清单。3.1 基础软件环境操作系统Windows 10/11, Linux (Ubuntu 18.04), macOS。本文示例代码在各系统上通用。Python推荐 Python 3.8 至 3.11 版本。可使用python --version检查。包管理工具pip或conda。建议使用虚拟环境如venv,conda env隔离项目依赖。3.2 PyTorch 安装这是核心依赖。安装命令取决于你的系统、包管理工具以及是否需要 GPUCUDA支持。通过 PyTorch 官网获取安装命令是最可靠的方式访问 https://pytorch.org/get-started/locally/ 选择你的配置如 Stable, Linux, Pip, Python, CUDA 11.8网站会生成对应的pip或conda命令。例如对于使用pip和 CUDA 11.8 的 Linux/Windows 用户pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118对于仅使用 CPU 的用户pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu对于使用 Conda 和 CUDA 11.8 的用户conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia验证安装import torch print(fPyTorch version: {torch.__version__}) print(fCUDA available: {torch.cuda.is_available()}) if torch.cuda.is_available(): print(fCUDA version: {torch.version.cuda}) print(fGPU device: {torch.cuda.get_device_name(0)})3.3 硬件要求非必须但推荐CPU现代多核处理器即可。内存 (RAM)至少 8GB处理较大数据集或模型时建议 16GB 以上。GPU (可选但强烈推荐用于训练)NVIDIA GPU支持 CUDA。常见型号如 GTX 1060, RTX 2060/3060/4060, RTX 3090/4090 等均可。显存入门级训练如 MNIST、CIFAR-104GB-6GB 显存可能足够。更复杂的模型和更大的批量大小需要 8GB 或更多。50 系显卡未来的 NVIDIA 50 系显卡预计将继续支持 CUDAPyTorch 届时会提供适配版本。当前安装时选择正确的 CUDA 版本即可。AMD/Apple SiliconPyTorch 已支持 AMD ROCm 和 Apple Metal (MPS) 后端但安装和配置流程与 CUDA 不同需参考官方文档。磁盘空间至少预留 2-5GB 空间用于安装 PyTorch 及相关库。3.4 开发工具可选代码编辑器/IDEVS Code, PyCharm, Jupyter Notebook 等。必要的 Python 库我们主要使用torch及其子模块torch.nn,torch.optim。为了数据可视化可以安装matplotlib和numpy。pip install matplotlib numpy4. 核心组件详解与代码封装现在我们进入核心部分一步步构建一个完整的模型。4.1 继承nn.Module构建模型容器nn.Module是所有神经网络模块的基类。你的自定义模型必须继承它。关键步骤继承nn.Module在__init__方法中调用父类的__init__。定义网络层在__init__中将网络层定义为类的属性。使用nn中的预定义层如nn.Linear,nn.Conv2d,nn.ReLU。实现forward方法定义数据从输入到输出的计算图。切勿直接调用model.forward(x)而是使用model(x)因为后者会调用__call__方法其中包含了forward以及一些内部钩子管理。示例构建一个简单的全连接网络用于分类import torch import torch.nn as nn import torch.nn.functional as F class SimpleClassifier(nn.Module): 一个简单的三层全连接分类网络 def __init__(self, input_size, hidden_size, num_classes): 初始化网络层。 Args: input_size: 输入特征维度 (e.g., 784 for flattened MNIST) hidden_size: 隐藏层神经元数量 num_classes: 分类类别数 (e.g., 10 for MNIST) super(SimpleClassifier, self).__init__() # 必须调用父类初始化 # 定义网络层 self.fc1 nn.Linear(input_size, hidden_size) # 全连接层1 self.relu nn.ReLU() # 激活函数 self.fc2 nn.Linear(hidden_size, hidden_size // 2) # 全连接层2 self.fc3 nn.Linear(hidden_size // 2, num_classes) # 输出层 # 注意我们没有将 nn.Softmax 定义为层通常在计算损失时由损失函数处理或最后单独应用。 def forward(self, x): 定义前向传播路径 out self.fc1(x) out self.relu(out) out self.fc2(out) out self.relu(out) out self.fc3(out) # 输出 raw logits (未归一化的分数) # 不在这里做 softmax因为 CrossEntropyLoss 内部包含了 softmax 计算 return out # 实例化模型 model SimpleClassifier(input_size784, hidden_size128, num_classes10) print(model)代码解析nn.Linear是线性层全连接层参数为(输入特征数, 输出特征数)。我们将nn.ReLU实例化为一个层 (self.relu)当然也可以在forward中直接使用F.relu(out)。forward方法清晰地描述了数据流x - fc1 - relu - fc2 - relu - fc3 - out。打印模型会显示其层次结构这是调试的好方法。4.2 损失函数衡量预测与目标的差距损失函数将模型的预测输出和真实标签映射为一个标量值这个值代表了模型当前的“错误程度”。常见损失函数选择nn.CrossEntropyLoss最常用的分类损失。它结合了LogSoftmax和NLLLoss因此模型的最后一层不需要加 Softmax 激活。输入是 raw logits (形状为[batch_size, num_classes])目标是类别索引 (形状为[batch_size]类型为torch.LongTensor)。nn.MSELoss(均方误差)用于回归任务计算预测值和目标值之间差的平方的平均值。nn.BCELoss(二元交叉熵)用于二分类输出层需使用nn.Sigmoid将值压缩到 [0,1]。nn.BCEWithLogitsLoss结合了 Sigmoid 和 BCELoss数值上更稳定推荐用于二分类。如何使用# 假设我们有一个批量的预测和标签 batch_size 4 num_classes 10 dummy_logits torch.randn(batch_size, num_classes) # 模型输出的 raw logits dummy_labels torch.randint(0, num_classes, (batch_size,)) # 真实标签每个样本一个类别索引 # 定义损失函数 criterion nn.CrossEntropyLoss() # 创建损失函数实例 # 计算损失 loss criterion(dummy_logits, dummy_labels) print(fCalculated loss: {loss.item():.4f})关键点损失函数实例 (criterion) 是可调用的。调用时传入(预测值, 目标值)返回一个标量张量。调用loss.backward()会从这个标量开始自动计算图中所有需要梯度的张量的梯度。4.3 优化器更新模型参数以最小化损失优化器持有模型参数的引用并根据梯度更新它们。常见优化器optim.SGD随机梯度下降可以添加动量 (momentum) 来加速收敛。optim.Adam自适应矩估计目前最常用的优化器通常能提供较好的默认性能。optim.AdamWAdam 的一个变体修正了权重衰减 (weight decay) 的实现在许多任务上表现更好。optim.RMSprop常用于循环神经网络。配置与使用import torch.optim as optim # 将模型的所有可训练参数传递给优化器 optimizer optim.Adam(model.parameters(), lr0.001) # 学习率是关键超参数 # 也可以为不同的层设置不同的学习率精细调优时常用 # optimizer optim.SGD([ # {params: model.fc1.parameters(), lr: 0.01}, # {params: model.fc2.parameters(), lr: 0.005}, # {params: model.fc3.parameters(), lr: 0.001}, # ], momentum0.9)关键参数model.parameters()获取模型中所有需要梯度可训练的参数。lr(learning rate)学习率控制每次参数更新的步长。太大可能导致震荡不收敛太小则收敛慢。通常从1e-3,1e-4等值开始尝试。其他如weight_decay(L2正则化)betas(Adam中的动量参数) 等可根据需要调整。5. 组装训练循环让模型动起来训练循环是将数据、模型、损失函数和优化器串联起来的引擎。一个典型的训练步骤一个 batch如下# 假设我们已经有了 # model: 定义好的模型 # criterion: 损失函数 # optimizer: 优化器 # train_loader: 数据加载器 (提供 (data, target) 的迭代器) # device: 设备 (cuda 或 cpu) model.train() # 将模型设置为训练模式影响某些层如 Dropout, BatchNorm 的行为 for batch_idx, (data, target) in enumerate(train_loader): # 1. 准备数据 data, target data.to(device), target.to(device) # 2. 前向传播 (Forward Pass) output model(data) # 调用 model.forward(data) # 3. 计算损失 loss criterion(output, target) # 4. 反向传播 (Backward Pass) optimizer.zero_grad() # 清空上一轮计算的梯度非常重要 loss.backward() # 计算损失关于模型参数的梯度 # 5. 参数更新 optimizer.step() # 根据梯度更新模型参数 # 可选打印训练信息 if batch_idx % 100 0: print(fTrain Epoch: [{epoch}] [{batch_idx * len(data)}/{len(train_loader.dataset)} f({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f})训练循环五步详解准备数据将数据移至 GPU如果可用。前向传播输入数据通过模型得到预测输出。计算损失用损失函数比较预测输出和真实标签。反向传播optimizer.zero_grad()必须做因为 PyTorch 会累积梯度如果不清零本次计算的梯度会与上次的累加。loss.backward()自动计算图中所有requires_gradTrue的张量的梯度并存储在每个参数的.grad属性中。参数更新optimizer.step()根据优化器算法如 Adam, SGD和参数的.grad来更新参数值。6. 完整示例在 MNIST 数据集上训练分类器让我们用一个完整的、可运行的例子来整合所有概念。我们将使用经典的 MNIST 手写数字数据集。import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms from torch.utils.data import DataLoader import matplotlib.pyplot as plt # 1. 定义设备 device torch.device(cuda if torch.cuda.is_available() else cpu) print(fUsing device: {device}) # 2. 数据预处理与加载 transform transforms.Compose([ transforms.ToTensor(), # 将 PIL Image 或 numpy.ndarray 转换为 Tensor并归一化到 [0.0, 1.0] transforms.Normalize((0.1307,), (0.3081,)) # MNIST 数据集的均值和标准差 ]) train_dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(./data, trainFalse, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size1000, shuffleFalse) # 测试时不需要 shuffle # 3. 定义模型使用之前定义的 SimpleClassifier class SimpleClassifier(nn.Module): def __init__(self, input_size, hidden_size, num_classes): super(SimpleClassifier, self).__init__() self.fc1 nn.Linear(input_size, hidden_size) self.relu nn.ReLU() self.fc2 nn.Linear(hidden_size, hidden_size // 2) self.fc3 nn.Linear(hidden_size // 2, num_classes) def forward(self, x): out self.fc1(x) out self.relu(out) out self.fc2(out) out self.relu(out) out self.fc3(out) return out model SimpleClassifier(input_size28*28, hidden_size128, num_classes10).to(device) # 4. 定义损失函数和优化器 criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 5. 训练函数 def train(epoch): model.train() train_loss 0 for batch_idx, (data, target) in enumerate(train_loader): data, target data.to(device), target.to(device) data data.view(data.size(0), -1) # 将图像展平 [batch, 1, 28, 28] - [batch, 784] optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() train_loss loss.item() if batch_idx % 100 0: print(fTrain Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} f({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}) avg_loss train_loss / len(train_loader) return avg_loss # 6. 测试函数 def test(): model.eval() # 将模型设置为评估模式 test_loss 0 correct 0 with torch.no_grad(): # 关闭梯度计算节省内存和计算 for data, target in test_loader: data, target data.to(device), target.to(device) data data.view(data.size(0), -1) output model(data) test_loss criterion(output, target).item() # 累加批次损失 pred output.argmax(dim1, keepdimTrue) # 获取预测类别 correct pred.eq(target.view_as(pred)).sum().item() test_loss / len(test_loader) accuracy 100. * correct / len(test_loader.dataset) print(f\nTest set: Average loss: {test_loss:.4f}, fAccuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n) return test_loss, accuracy # 7. 开始训练与测试 num_epochs 5 train_losses [] test_losses [] accuracies [] for epoch in range(1, num_epochs 1): avg_train_loss train(epoch) train_losses.append(avg_train_loss) test_loss, accuracy test() test_losses.append(test_loss) accuracies.append(accuracy) # 8. 可视化训练过程 plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(range(1, num_epochs1), train_losses, labelTrain Loss) plt.plot(range(1, num_epochs1), test_losses, labelTest Loss) plt.xlabel(Epoch) plt.ylabel(Loss) plt.legend() plt.title(Training and Test Loss) plt.subplot(1, 2, 2) plt.plot(range(1, num_epochs1), accuracies, labelTest Accuracy, colorgreen) plt.xlabel(Epoch) plt.ylabel(Accuracy (%)) plt.legend() plt.title(Test Accuracy) plt.tight_layout() plt.show()运行预期程序会下载 MNIST 数据集到./data目录。你会看到每个训练周期 (Epoch) 中训练损失在逐步下降。在测试集上准确率会随着训练逐渐上升5 个周期后可能达到 97% 以上。最后会生成两张图表分别展示训练/测试损失和测试准确率的变化曲线。7. 资源占用与性能观察在本地运行上述代码时你可以观察系统资源的使用情况这对于理解模型规模和硬件需求很有帮助。7.1 如何观察资源占用GPU 显存与利用率如果你使用 CUDA可以在命令行使用nvidia-smi命令Windows/Linux来监控。在代码中可以使用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()来查看当前和峰值显存占用。CPU 与内存使用系统任务管理器Windows、htopLinux或活动监视器macOS来观察。7.2 影响性能的关键因素批量大小 (Batch Size)DataLoader中的batch_size参数。增大 batch size 可以提高 GPU 利用率加速训练但也会增加单次迭代的显存占用。如果遇到CUDA out of memory错误首先尝试减小batch_size。模型复杂度网络层数、每层神经元/通道数。更复杂的模型参数更多计算量更大显存占用更高。数据维度输入图像/特征的大小。例如将 MNIST 的 28x28 图像换成 224x224 的 ImageNet 图像显存占用会急剧增加。数据类型默认的torch.float32(32位浮点数)。在某些场景下可以使用model.half()或torch.float16(半精度) 来减少显存占用和加速计算但可能会影响数值精度和稳定性。7.3 针对本示例的预期资源占用模型参数量可以使用sum(p.numel() for p in model.parameters())计算。对于我们的SimpleClassifier(输入784隐藏层128输出10)参数量大约在 10 万左右是一个非常小的模型。显存占用在 batch size64 的情况下即使在集成显卡或低端独立显卡上也能轻松运行峰值显存占用可能只有几十到几百 MB。训练时间在 CPU 上5 个 epoch 可能只需几分钟。在 GPU 上可能只需几十秒。性能调优提示如果项目模型更大数据更复杂可以使用torch.utils.data.DataLoader的num_workers参数进行多进程数据加载防止数据准备成为瓶颈。使用torch.cuda.amp进行自动混合精度训练在支持的 GPU 上提升速度并节省显存。定期使用torch.cuda.empty_cache()清理未使用的显存缓存。8. 常见问题与排查方法在封装和训练模型时你可能会遇到以下典型问题。这里提供一个快速排查指南。问题现象可能原因排查方式解决方案运行时错误RuntimeError: size mismatch, m1: [a x b], m2: [c x d]网络层输入/输出维度不匹配。例如全连接层nn.Linear期望的输入特征数与实际传入的不同。1. 检查forward方法中每一层的输入输出形状。2. 在forward开头用print(x.shape)打印输入形状。3. 检查数据预处理后的形状如是否忘记view展平图像。1. 修正网络层定义中的输入/输出维度。2. 确保数据在传入模型前被正确处理成正确的形状如data.view(batch_size, -1)。训练损失不下降或为 NaN1. 学习率 (lr) 设置过大或过小。2. 数据未归一化或存在异常值。3. 损失函数选择错误如分类任务用了MSELoss。4. 梯度爆炸。1. 打印前几个批次的损失值观察趋势。2. 检查输入数据的范围是否在 0-1 或已标准化。3. 检查标签格式是否与损失函数要求匹配如CrossEntropyLoss需要类别索引而非 one-hot。1. 尝试更小的学习率如1e-4,1e-5。2. 对输入数据进行标准化。3. 确保使用正确的损失函数。4. 使用梯度裁剪 (torch.nn.utils.clip_grad_norm_)。optimizer.zero_grad()忘记调用梯度在批次间累积导致参数更新方向错误训练不稳定或发散。观察损失曲线剧烈震荡或无法收敛。务必在loss.backward()之前调用optimizer.zero_grad()。GPU 内存不足 (CUDA out of memory)1. 批量大小 (batch_size) 太大。2. 模型太大。3. 在训练循环中累积了中间张量未释放。使用nvidia-smi观察显存占用情况。1. 减小batch_size。2. 简化模型结构。3. 使用with torch.no_grad():包裹不需要梯度的计算。4. 使用del删除不再需要的变量并调用torch.cuda.empty_cache()。测试准确率远低于训练准确率模型过拟合 (Overfitting)。观察训练损失持续下降但测试损失在某个点后开始上升。1. 增加训练数据。2. 在模型中添加正则化层如nn.Dropout。3. 为优化器添加 L2 正则化 (weight_decay参数)。4. 使用数据增强。model.eval()模式未切换在测试或推理时忘记调用model.eval()导致Dropout、BatchNorm等层仍处于训练行为模式。对比调用model.eval()前后的测试结果。在测试和推理前调用model.eval()在训练前调用model.train()。保存/加载模型后性能不一致只保存了模型参数 (state_dict)未保存优化器状态或超参数或加载时模型结构不匹配。检查保存和加载的代码。1. 保存时同时保存model.state_dict()、optimizer.state_dict()和 epoch 等信息。2. 加载时先实例化一个结构相同的模型再model.load_state_dict(torch.load(PATH))。9. 最佳实践与使用建议遵循以下实践能让你的 PyTorch 模型开发更加高效和稳健。模块化设计将复杂的网络拆分成有意义的子模块也继承自nn.Module然后在主模型中组合它们。这提高了代码的可读性和复用性。class ResidualBlock(nn.Module): def __init__(self, in_channels, out_channels, stride1): super().__init__() self.conv1 nn.Conv2d(...) self.bn1 nn.BatchNorm2d(...) # ... 其他层 def forward(self, x): identity x out self.conv1(x) # ... 前向传播 out identity # 残差连接 return out class MyResNet(nn.Module): def __init__(self): super().__init__() self.layer1 ResidualBlock(64, 64) self.layer2 ResidualBlock(64, 128) # ...参数初始化nn.Module中的参数默认使用某种初始化如nn.Linear使用 Kaiming Uniform。对于深层网络合理的初始化至关重要。可以使用torch.nn.init中的函数手动初始化。def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.constant_(m.bias, 0) model.apply(init_weights)使用nn.Sequential简化顺序结构对于简单的层堆叠nn.Sequential可以简化代码。self.net nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.2), nn.Linear(256, 10) ) # forward 中直接: out self.net(x)管理设备使用model.to(device)和data.to(device)。可以定义一个全局的device变量来统一管理。device torch.device(cuda:0 if torch.cuda.is_available() else cpu) model MyModel().to(device)保存与加载完整检查点为了能够从中断处继续训练建议保存一个包含模型、优化器、epoch、损失等所有状态的字典。# 保存 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]) epoch checkpoint[epoch]版本控制与依赖记录使用requirements.txt或environment.yml记录项目依赖的库及其版本特别是torch和torchvision的版本以确保代码可复现。掌握nn.Module、损失函数和优化器的封装与使用是解锁 PyTorch 强大能力的关键第一步。它为你构建更复杂模型如 CNN、RNN、Transformer和实现更高级训练技巧如学习率调度、梯度裁剪、多任务学习奠定了坚实的基础。建议在理解本文示例的基础上尝试修改网络结构、更换损失函数、调整优化器参数并观察训练结果的变化这是加深理解的最佳途径。