DeepEpochs

DeepEpochs 是 PyTorch 深度学习模型训练工具。将模型、损失函数、优化器和 DataLoader 交给 Trainer，即可执行训练、验证和测试。指标聚合、Checkpoint、日志、混合精度与分布式训练由训练器管理，特定任务可以通过 Callback、EpochTask 和自定义步骤扩展。

安装

运行环境要求 Python >= 3.8、PyTorch >= 2.0。在包含 pyproject.toml 的仓库根目录安装：

python -m pip install .

开发时使用可编辑安装，使源码修改直接用于运行和测试：

python -m pip install -e .

库源码位于 src/deepepochs，测试位于 tests，示例位于 examples。MNIST 示例需要额外安装 torchvision，图节点分类示例需要 DGL。安装完成后，从仓库根目录运行示例和测试脚本。

基本使用

Trainer 接受普通 PyTorch 模型、损失函数和优化器。fit 执行训练并返回逐轮记录；传入验证 DataLoader 后，每轮训练还会计算验证指标。训练完成后调用 test 评估模型，最后使用 close() 结束训练器的使用并释放资源。

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset
from deepepochs import Trainer

x = torch.randn(256, 16)
y = (x[:, 0] > 0).long()
train_dl = DataLoader(TensorDataset(x[:192], y[:192]), batch_size=32, shuffle=True)
val_dl = DataLoader(TensorDataset(x[192:], y[192:]), batch_size=32)

model = nn.Linear(16, 2)
opt = torch.optim.Adam(model.parameters(), lr=0.001)
trainer = Trainer(model, F.cross_entropy, opt, epochs=5)
try:
    progress = trainer.fit(train_dl, val_dl)
    results = trainer.test(val_dl)
finally:
    trainer.close()

epochs 设置训练的总轮数，test 返回测试任务的指标字典。fit 和 test 可以在训练器关闭前连续调用，因此一次运行可以完成训练、模型评估和日志记录。使用 try/finally 可在运行结束时关闭训练器，也适用于循环创建多个训练器的场景。

数据与指标

训练器根据 batch 的结构准备模型输入和标签：

• 指定损失函数时，tuple/list batch 的最后一项作为标签，其余项作为模型输入。标签可以包含嵌套的 tuple、list 或 dict。
• 不指定损失函数时，tuple/list batch 的所有项作为模型输入，模型需要返回标量损失或包含 loss 的字典。
• dict batch 的全部字段作为关键字参数传给模型，不自动拆出标签。模型可以读取字典中的标签并计算损失；指标需要标签时，可由自定义步骤或模型输出提供。

默认会将输入和标签中的 Tensor 递归搬到目标设备。自定义步骤接收单个输入时使用 self.model(batch_x)，多个输入为 tuple 时使用 self.model(*batch_x)，字典输入则使用 self.model(**batch_x)。字典输入输出示例 展示了由模型接收标签并返回损失的用法。

指标函数通过 metrics 传入，通常接受 (preds, targets)，返回标量或字典。除了项目提供的指标，也可以使用 torchmetrics.functional 定义函数式指标。每个 batch 产生的指标由 Patch 汇总为整个数据集的结果，metric_patch 用于选择聚合方式：

聚合方式应与指标的数学性质一致。Precision、recall、F1 需要依据完整预测或混淆矩阵计算，直接平均每批结果不能得到数据集指标。ConfusionMetrics 可在自定义步骤中一次计算多个分类指标；特殊数据则可通过 batch_size 函数提供用于加权的样本数。

设备与训练配置

训练器默认选择可用的 CUDA 设备，没有 CUDA 时使用 CPU。通过 device 可以显式指定 PyTorch 设备字符串或 torch.device；使用 MPS 时需要显式设置该参数。

控制台输出由 log_batch、log_tqdm、log_long 和 show_info 配置，分别控制逐批指标、进度条、指标小数位数和训练信息显示。默认显示逐批指标及训练信息，指标采用四位小数。

混合精度与梯度累积

precision 控制模型前向和损失计算的 autocast。CUDA FP16 使用 GradScaler 缩放梯度；BF16 使用目标设备支持的 autocast。可以在构造训练器时同时配置精度和梯度累积：

trainer = Trainer(
    model,
    F.cross_entropy,
    opt,
    device='cuda',
    precision='fp16',
    gradient_accumulation_steps=2,
)
try:
    trainer.fit(train_dl, val_dl)
finally:
    trainer.close()

gradient_accumulation_steps=2 将两个 DataLoader batch 的梯度用于一次更新。训练器按窗口中的实际样本数归一化梯度，轮末不足一个完整窗口时也会更新参数。分布式训练在需要更新时同步梯度，使各进程使用一致的累计结果。

梯度裁剪作用于反缩放后的梯度。多优化器的 norm 裁剪联合处理有限梯度组，GradScaler 按各优化器的溢出结果决定是否更新参数及对应的 step 调度器。

分布式训练

使用 torchrun 启动同一份训练脚本，每个进程构造模型、优化器和 DataLoader，再交给 Trainer。训练器读取 RANK、WORLD_SIZE 和 LOCAL_RANK，初始化进程组、选择当前设备并包装 DDP；CPU 使用 Gloo，CUDA 使用 NCCL。

trainer = Trainer(
    model,
    F.cross_entropy,
    opt,
    epochs=5,
    device='cuda',
    precision='fp16',
    split_batches=True,
)
try:
    trainer.fit(train_dl, val_dl)
    results = trainer.test(test_dl)
finally:
    trainer.close()

分布式训练示例 包含模型、数据准备和完整运行入口，可以按设备数量启动：

# 每张 GPU 一个进程
torchrun --standalone --nproc-per-node=2 examples/21-distributed.py

# CPU 双进程
CUDA_VISIBLE_DEVICES="" torchrun --master-addr=127.0.0.1 --master-port=29500 --nproc-per-node=2 examples/21-distributed.py

普通 DataLoader 由训练器自动分片，并在每轮更新采样器的 epoch。训练分片保持各进程等长，验证与测试分片按实际样本划分。任务结束后，各进程聚合全局指标，并依据同一结果进行学习率调度和早停；日志与 Checkpoint 由主进程写入。

普通 DataLoader 自动分片时，split_batches=False 的名义全局有效批量为 batch_size × world_size × gradient_accumulation_steps。设置为 True 时，DataLoader 的 batch_size 表示全局批量，有效批量为 batch_size × gradient_accumulation_steps，且 batch_size 必须能被进程数整除。调用方提供 DistributedSampler、自定义 sampler、batch_sampler 或 DataLoader 子类时，训练器直接使用该 DataLoader，批量和分片方式由调用方配置。

共享数据只需准备一次时，将下载或缓存创建放入 with trainer.main_process_first():。主进程完成准备后，其他进程再读取相同缓存；单进程脚本也可以使用该上下文。

进程组由创建方管理。Trainer.close() 释放训练器创建的进程组，调用方事先初始化的进程组由调用方关闭。训练器关闭时还会关闭共享日志 writer，确保待写入的日志完成保存。

Checkpoint 与恢复训练

CheckCallback 监控训练或验证指标，按指标变化保存模型，并通过 patience 控制早停。save_best=True 保存指标最佳时的快照，save_best=False 保存最近一次监控时的快照；patience=0 表示关闭早停。

from deepepochs import CheckCallback

checker = CheckCallback('loss', on_stage='val', mode='min', patience=5)
trainer = Trainer(model, F.cross_entropy, opt, epochs=5,
                  callbacks=[checker], resume=True)
try:
    trainer.fit(train_dl, val_dl)
    results = trainer.test(test_dl)
finally:
    trainer.close()

resume=True 汇总全部监控器的 ckpt_dir，按文件修改时间选择最近的 Checkpoint；传入运行 ID 时，仅在各目录的该次运行中选择。没有 Checkpoint 的首次启动从头训练，指定运行缺失或文件无法读取时会报告错误。通过 trainer.load_checkpoint(path) 可以显式加载指定文件。

Checkpoint 保存模型、优化器、调度器、已完成轮次、早停状态和启用的 AMP GradScaler 状态，使恢复训练从相应的下一轮开始。running_id 为本次运行指定名称，用于日志和 Checkpoint 子目录。多个 CheckCallback 各自保存独立文件，每个快照包含本轮全部监控器的状态；after_epoch 回调请求停止时，训练器会完成本轮已登记的保存请求。

有可用的 Checkpoint 时，test 按声明顺序选择第一个已有文件的监控器，并临时载入其中的模型权重。测试完成或异常退出后，模型恢复测试前的权重，训练进度保持不变，因此可以继续训练。

文件由 torch.save 写入，各字段含义如下：

单进程、DDP 和编译模型使用相同的参数格式。单个 Checkpoint 可以移动或改名后加载；多个监控器的文件应按其相对目录关系一起搬迁，以便各监控器找到对应文件。

定制训练流程

Callback

继承 Callback 并实现需要的事件方法，再将实例传给 Trainer(callbacks=...)，可以在训练流程中执行初始化、记录和监控等跨步骤操作。分布式训练会在初始化回调完成后同步模型参数与 buffer，使各进程从相同模型状态开始。

TrainerBase 与 EpochTask

继承 TrainerBase 或 EpochTask，实现 step、train_step、val_step、test_step 或 evaluate_step，可以定义各阶段的计算过程。步骤方法接受 (batch_x, batch_y, **step_args)，返回 None 或 {指标名: Patch对象}。

执行时先选择 step，再选择相应的 train_step、val_step 或 test_step；验证和测试未定义阶段专用方法时，使用 evaluate_step。同名方法优先采用 EpochTask 的实现，其次采用 Trainer 的实现。创建 EpochTask 时提供的额外关键字参数会传给步骤方法。

一个阶段可以安排多个 EpochTask，分别配置 DataLoader、指标和步骤参数。train_dl 与 train_tasks 二选一，val_dl 可以与 val_tasks 合并，test_dl 可以与 tasks 合并。

在训练步骤中，self.loss(...) 默认完成损失计算、反向传播和优化。设置 do_optimize=False 后，可以使用 self.optimize(optimize_index=...) 选择优化器并控制更新时机，适用于多优化器任务和手动梯度累积。

手动切分一个 DataLoader batch 时，每个子 batch 调用 self.loss(..., loss_adjust=子批样本数 / 整批样本数, do_optimize=False)，最后调用一次 self.optimize()。非末子批的前向和反向可以一起放在 self.no_sync() 中，单进程下也可使用这一上下文。手动子批示例 展示了完整循环，跨批累积示例 使用 gradient_accumulation_steps 管理 DataLoader batch 之间的累积。

Patch

继承 PatchBase 并实现 load、add 和 forward，可以为指标定义聚合状态。load 装载当前 batch 数据并返回自身，add 合并状态，forward 返回指标或字典。对于已经计算出的标量，ValuePatch 按样本数进行聚合。

图示为 tuple/list batch 中输入与标签分开的数据流。EpochTask 从 DataLoader 读取 batch，调用模型获得预测，再将损失和指标装入对应 Patch。任务结束时，各 Patch 计算汇总结果，形成阶段指标字典。

日志与预测解释

LogCallback 将指标和学习率写入 TensorBoard。设置 log_graph=True 可以记录模型结构图，example_input 用于指定结构图的输入；未指定时采用第一个训练 batch。通过 Trainer 的 hyper_params 传入超参数字典，测试完成后可将超参数与测试指标一起记录。

InterpreteCallback 根据逐样本指标选择样本，并通过 top_samples() 返回样本列表，每条记录为 [指标值, 预测, 标签, 输入特征]。输入特征与样本对应：单输入为 Tensor，多输入为 tuple 或 dict。入选样本的输入、预测和标签各自拥有独立的 CPU Tensor 存储。设置 image_data=True 可以记录图像，多输入时采用首个输入，字典按字段插入顺序选择。

需要观察模型内部计算时，AnalyzeCallback 可记录各层的前向输出或反向梯度。它在训练结束后移除 hooks，日志 writer 由训练器统一管理。各日志回调共享同一个 writer，fit 后的 test 会继续写入同一日志，close() 则完成日志保存并退出写线程。

示例

验证与开发

环境与回归结果

验证涵盖 CPU 训练、AMP、多进程通信和发行包安装。开发环境及其测试结果为：

CPU 回归覆盖基础训练、二分类指标、Patch 聚合、嵌套容器转换、未排序 key 分组、多输入预测解释、学习率调度、Checkpoint 多目录恢复与搬迁、停止收尾、日志线程释放以及双进程全局指标。CPU AMP 定向回归验证溢出缓存复用、有限组联合 norm 与非有限组隔离；CUDA 回归验证 FP32 更新对照、FP16/BF16、裁剪、尾批累积、多优化器、溢出跳步、AMP 状态恢复和嵌套标签设备往返。

从仓库根目录安装开发环境、运行测试并构建发行包：

python -m pip install build
python -m pip install -e .
MPLBACKEND=Agg python -m unittest discover -s tests -v
python -m build

小数据联调脚本执行3轮训练、全局验证和 Checkpoint 临时加载测试，可以用于检查单进程、CPU 多进程或 CUDA 环境：

python tests/smoke_training.py --output /tmp/deepepochs-smoke-python
torchrun --master-addr=127.0.0.1 --master-port=29500 --nproc-per-node=2 tests/smoke_training.py --output /tmp/deepepochs-smoke-ddp
python tests/smoke_training.py --device cuda --precision fp16 --output /tmp/deepepochs-smoke-cuda

多 GPU 联调

在2卡、4卡与 FP32、FP16、BF16 的组合中，NCCL 训练与评估共98项检查全部通过，其中训练56项、评估42项。检查覆盖全局梯度参考、梯度累积、多优化器、裁剪与溢出、全局指标、空分片、调度、早停、多监控器目录统一恢复、Checkpoint 和预测解释。

FP32、FP16 的四卡 Checkpoint 分别通过4项单卡加载、搬迁、续训和测试状态恢复检查，FP16 同时恢复 GradScaler 状态。四卡 FP16 的普通模式与 Inductor 编译模式完成训练联调，示例21完成 MNIST 下载、四卡两轮训练和测试。各项检查的配置与结果保存在验证汇总及原始JSON记录中。

可以先生成四卡训练和评估结果，再将评估脚本生成的 Checkpoint 交给单卡脚本继续训练：

torchrun --standalone --nproc-per-node=4 tests/multigpu_training.py --expected-world-size 4 --precision fp16 --output /tmp/deepepochs-training-4-fp16.json
torchrun --standalone --nproc-per-node=4 tests/multigpu_evaluation.py --expected-world-size 4 --precision fp16 --output /tmp/deepepochs-evaluation-4-fp16
python tests/single_gpu_checkpoint.py --checkpoint /tmp/deepepochs-evaluation-4-fp16/checkpoints/multiple/checkpoint-2.ckpt --source-world-size 4 --precision fp16 --device cuda:0 --output /tmp/deepepochs-single-checkpoint

吞吐测量

吞吐测试使用同一四卡环境，GPU 间通过 PCIe PHB 互联，无 NVLink。脚本以 CPU 缓存的合成图像驱动 Trainer 的完整批次训练路径，模型为5,244,968参数的九层卷积网络，输入大小224×224，TF32关闭。

每卡 batch 固定为32，因此1卡、2卡、4卡的全局 batch 分别为32、64、128，测量属于弱扩展。预热20步后测量200步，计时窗口包括 CPU 数据整理、H2D、模型计算、参数更新和 batch 统计。全局吞吐按全部样本数除以最慢 rank 的测量时间计算，倍率以同精度的单卡结果为基准。

Epoch 汇总和日志位于批次计时窗口之外，完整 fit 耗时另行记录。JSON 同时包含各 rank 的步时、吞吐和峰值显存，可用于检查测量范围及各进程的运行情况。复现测量时，依次运行各精度与卡数组合：

for precision in fp32 fp16 bf16; do
  for gpus in 1 2 4; do
    torchrun --standalone --nproc-per-node="$gpus" tests/multigpu_benchmark.py --expected-world-size "$gpus" --precision "$precision" --warmup 20 --steps 200 --batch-size 32 --output "/tmp/deepepochs-benchmark-${gpus}-${precision}.json"
  done
done

适用范围与约束

• DDP 训练在 drop_last=False 时会为等长分片补齐样本，训练指标对应实际处理的数据；验证与测试不重复补齐样本。
• 自定义采样器及 DataLoader 子类由调用方保证分布式分片正确。各进程的任务安排、训练步数及调用顺序需要一致。
• TensorPatch 保存并传输完整预测与标签。数据量较大时，可加权平均或混淆矩阵类指标适合使用 MeanPatch 或 ConfusionPatch。
• 吞吐测量采用合成 CPU 缓存数据。实际输入管线中的读取、解码和预处理会影响端到端吞吐。