训练基础设施

相关源文件

• docs/experiments/mnist.html
• docs/normalization/batch_norm/index.html
• docs/normalization/batch_norm/readme.html
• docs/normalization/layer_norm/index.html
• docs/normalization/layer_norm/readme.html
• labml_nn/experiments/mnist.py
• labml_nn/experiments/nlp_autoregression.py
• labml_nn/normalization/batch_norm/__init__.py
• labml_nn/normalization/layer_norm/__init__.py
• labml_nn/normalization/layer_norm/readme.md
• labml_nn/optimizers/__init__.py
• labml_nn/optimizers/ada_belief.py
• labml_nn/optimizers/adam.py
• labml_nn/optimizers/adam_warmup.py
• labml_nn/optimizers/adam_warmup_cosine_decay.py
• labml_nn/optimizers/amsgrad.py
• labml_nn/optimizers/mnist_experiment.py
• labml_nn/optimizers/noam.py
• labml_nn/optimizers/performance_test.py
• labml_nn/optimizers/radam.py
• labml_nn/transformers/configs.py
• labml_nn/transformers/gpt/__init__.py

本页介绍 labml_nn 仓库中不同模型使用的核心训练基础设施。该基础设施为深度学习模型的训练提供模块化和可配置的组件，包括

1. 不同任务的训练配置
2. 优化器实现
3. 归一化技术
4. 训练循环和实用工具

有关具体的模型实现，请参阅Transformer 架构、微调技术或生成模型。

训练基础设施概览

该代码库中的训练基础设施遵循模块化设计模式，组件可以针对不同任务进行混合和匹配。其核心是可配置的实验类、优化器以及用于管理训练工作流的实用工具。

来源: labml_nn/experiments/nlp_autoregression.py labml_nn/optimizers/__init__.py labml_nn/normalization/batch_norm/__init__.py labml_nn/normalization/layer_norm/__init__.py

训练配置

该代码库使用配置系统来定义和定制不同任务的训练工作流。

基本训练配置

训练配置的基类是 TrainValidConfigs，它提供训练和验证循环、指标跟踪以及模型管理的共享功能。

特定任务的配置扩展此基类，以定制不同场景的训练过程。

NLP 自回归配置

NLPAutoRegressionConfigs 专为自回归文本生成任务设计，具有

• 带分词的文本数据集
• 序列建模
• 文本生成的可选采样功能
• 带洗牌选项的可配置数据加载器
• 梯度裁剪

关键组件

• step(): 处理训练或验证步骤
• sample(): 在训练期间生成文本样本
• 各种数据加载器选项（顺序或洗牌）

来源: labml_nn/experiments/nlp_autoregression.py41-104

示例配置选项

优化器

该代码库实现了 PyTorch Optimizer 类的扩展，包含一系列优化算法。

优化器架构

来源: labml_nn/optimizers/__init__.py70-164 labml_nn/optimizers/adam.py50-214

通用自适应优化器

GenericAdaptiveOptimizer 作为 Adam 及其变体等自适应优化算法的基类。它提供通用功能

• 参数组管理
• 状态初始化
• 优化步骤的模板方法
• 处理稀疏梯度

来源: labml_nn/optimizers/__init__.py70-164

权重衰减

权重衰减作为单独的类实现，以允许

• 解耦权重衰减（直接应用于参数）
• 常规权重衰减（添加到梯度）
• 绝对或依赖学习率的缩放

来源: labml_nn/optimizers/__init__.py167-218

优化器实现

该仓库包含多种优化器实现

优化器	主要功能	来源
Adam	Adam 基本实现	labml_nn/optimizers/adam.py
AMSGrad	维持二阶矩的最大值	labml_nn/optimizers/amsgrad.py
RAdam	带有方差校正的 Rectified Adam	labml_nn/optimizers/radam.py
AdaBelief	使用梯度方差进行自适应学习	labml_nn/optimizers/ada_belief.py
Noam	带有指数衰减的预热期	labml_nn/optimizers/noam.py
AdamWarmup	学习率线性预热	labml_nn/optimizers/adam_warmup.py
AdamWarmupCosineDecay	预热 + 余弦衰减调度	labml_nn/optimizers/adam_warmup_cosine_decay.py

每种方法都以不同方式扩展优化过程，主要通过修改

1. 自适应学习率的计算方式
2. 学习率的调度方式
3. 动量和方差的跟踪与使用方式

归一化技术

该代码库实现了两种主要的归一化技术：批量归一化 (Batch Normalization) 和层归一化 (Layer Normalization)。

批量归一化

批量归一化在批次维度上对激活进行归一化，这通过减少内部协变量偏移来帮助稳定训练。

主要特性

• 在批次中独立归一化每个特征
• 在训练期间跟踪运行统计数据用于推理
• 通过可学习参数应用缩放和偏移
• 在训练和推理期间行为不同

实现细节

来源: labml_nn/normalization/batch_norm/__init__.py103-206

层归一化

层归一化在特征维度上进行归一化，使其更适用于序列模型和批量大小较小的情况。

主要特性

• 独立地对每个样本的特征维度进行归一化
• 推理时不需要运行统计数据
• 对 NLP 任务和 Transformer 尤其有效
• 适用于可变序列长度和小型批量

实现细节

来源: labml_nn/normalization/layer_norm/__init__.py43-131

训练循环工作流

该训练基础设施实现了一个标准化的训练工作流，并为不同任务提供了定制点。

来源: labml_nn/experiments/nlp_autoregression.py124-172 labml_nn/optimizers/mnist_experiment.py64-105

Step 函数工作流

训练循环的核心是 step 函数，它

1. 将模型设置为训练或评估模式
2. 将数据传输到正确的设备
3. 通过模型运行前向传播
4. 计算损失和额外指标
5. 如果是在训练
   • 计算梯度
   • 如果配置，裁剪梯度
   • 使用优化器更新参数
   • 清除梯度
6. 记录指标

NLPAutoRegressionConfigs 的示例

来源: labml_nn/experiments/nlp_autoregression.py124-172

整合

训练基础设施组件协同工作，创建一个灵活模块化的系统。以下是它们如何用于 NLP 任务的示例

来源: labml_nn/experiments/nlp_autoregression.py labml_nn/transformers/gpt/__init__.py221-260

总结

该训练基础设施为实现和训练深度学习模型提供了一个全面的框架

1. 配置：针对不同任务的灵活配置类
2. 优化器：具有一致接口的一系列优化算法
3. 归一化：批量和层归一化实现
4. 训练循环：带有定制点的标准化训练工作流
5. 实用工具：指标跟踪、设备管理、检查点保存

该基础设施使得在重用通用训练组件和模式的同时，能够实现各种模型架构。