——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之八
前面我们已经讲过张量布局、静态图编译和框架的整体架构。静态图确实给框架带来了全局优化的空间，但有一张”全局地图”我还没展开——那就是数据。再精巧的计算图、再极致的 CUDA Kernel，如果 GPU 拿不到足够快的数据，也只能空转。训练框架里常常出现一种尴尬：算子优化了半天，profiler 一开，发现时间全耗在 cudaMemcpyH2D 或者 CPU 预处理上。
这一篇，我们把镜头对准训练管线的最前端，讲讲 Tech-Renaissance 的数据加载管线（Data Loading Pipeline）。它的核心目标很简单：用同一套抽象，把 MNIST、CIFAR、ImageNet 这三类规模、格式、存储方式完全不同的数据集，统一装进训练流程；同时让 CPU 解码、预处理、组 batch、H2D 传输这几个阶段高效衔接，不让 GPU 挨饿。
一、数据管线是训练框架的”厨房”
先回到常识。一个深度学习训练迭代，数据通常要经历这么一段旅程：
- 从磁盘读取原始样本（可能是 ubyte、bin、JPEG，也可能是自定义打包格式）；
- 解码成原始像素（如果是压缩图像）；
- 做数据增强与预处理（裁剪、缩放、翻转、归一化、类型转换等）；
- 按 batch 组织样本，把标签和图像数据排列整齐；
- 通过 PCIe/NVLink 从主机内存异步拷贝到 GPU 显存；
- GPU 拿到张量，开始真正的前向、反向计算。
如果第 1~5 步跟不上第 6 步的速度，GPU 就会处于”数据饥饿（data starvation）”状态。对于 MNIST 这种小数据集，数据几乎可以全部塞进内存，问题不大；但对于 ImageNet 这种上百万张 JPEG 的训练集，单线程读盘加解码再预处理，速度根本不够看。一块 A100 处理一个 256 样本的 batch 可能只要二三十毫秒，而 CPU 侧光解码一张 224×224 的 JPEG 就要好几毫秒，单线程跑完一个 batch 几百毫秒就过去了。
所以现代框架都会在这一层下功夫。PyTorch 的做法大家都很熟悉：torch.utils.data.Dataset 负责单样本读取，DataLoader 通过 num_workers 拉起多个子进程并行做预处理，再用 pin_memory=True 把 CPU 内存锁页，方便后续异步 H2D 传输。TensorFlow 则主推 TFRecord，把数据预先打包成顺序读取的二进制记录，减少小文件随机 IO。此外，Caffe 时代常用的 LMDB（Lightning Memory-Mapped Database）也是一种键值存储方案，利用内存映射文件实现高效的随机读取。这些方案的共同思路是：把磁盘上的随机读取（打开文件、seek、read）变成顺序读取，大幅减少 IO 开销。但代价是：每次修改数据增强策略，可能都需要重新生成这些二进制文件，灵活性受限。
Tech-Renaissance 面对的挑战并没有不同，但我们选择了一条更”框架内聚”的路线：不是把数据加载做成一个独立的外部组件，而是把它纳入整个静态图编译与运行时体系，让 DataLoader、Preprocessor、TransferStation 三个模块紧密配合，形成一条零拷贝、可复现、可扩展的管线。
二、DataLoader：统一接口，后端随意换
数据管线的第一层抽象是 DataLoader。在 include/renaissance/data/data_loader.h 中，它被定义为一个抽象基类：
class DataLoader {
public:
virtual void configure(
int num_load_workers,
int num_preproc_workers,
const std::string& train_path,
const std::string& val_path,
bool shuffle_train = true,
bool shuffle_val = false,
bool skip_first = false,
bool verify_crc = false) = 0;
virtual void begin_epoch(int epoch_id, bool is_train) = 0;
virtual void end_epoch() = 0;
virtual bool get_next_sample(
int preproc_worker_id,
int32_t& label,
const uint8_t*& data_ptr,
size_t& data_size) = 0;
virtual void set_train_mode(LoadMode mode) = 0;
virtual void set_val_mode(LoadMode mode) = 0;
virtual bool verify_dts_crc(const std::string& file_path) const = 0;
virtual void download(const std::string& save_path) = 0;
// ...
};
所有具体加载器——MnistLoaderRaw、MnistLoaderDts、CifarLoaderRaw、CifarLoaderDts、ImageNetLoaderRaw、ImageNetLoaderDts，以及用于部署的 SampleLoader——都继承自这个接口。设计上有几个关键约定：
- 单例模式：每个具体 Loader 都是全局单例，避免训练与验证、不同 epoch 之间重复创建；
- 线程安全：多个 Preprocessor worker 可以同时调用
get_next_sample()，Loader 内部保证互不干扰； - 零拷贝：
get_next_sample()返回的data_ptr指向 Loader 内部缓冲区，外部只读，不复制。Preprocessor 拿到指针后原地解码、原地处理，处理完的数据再交给 TransferStation 走 H2D 传输。整条链路中，同一个样本的原始数据在内存中只有一份，没有中间拷贝； - 生命周期明确：
begin_epoch()/end_epoch()用来切换 epoch，重置 shuffle 状态，区分训练/验证模式。
用户其实很少直接操作 DataLoader。通常只需要在 PREPROCESSOR_SETTING 里写一行 .dataset("imagenet", "/data/imagenet")，Preprocessor 在内部就会根据数据集名称和是否启用 DTS 格式，自动选择对应的具体加载器：
switch (dataset_type_) {
case DatasetType::mnist:
current_dataloader_ = dts_format ? &MnistLoaderDts::instance()
: &MnistLoaderRaw::instance();
break;
case DatasetType::cifar_10:
case DatasetType::cifar_100:
current_dataloader_ = dts_format ? &CifarLoaderDts::instance()
: &CifarLoaderRaw::instance();
break;
case DatasetType::imagenet:
current_dataloader_ = dts_format ? &ImageNetLoaderDts::instance()
: &ImageNetLoaderRaw::instance();
break;
}
这就是标题里说的”一套抽象”。对上层 Preprocessor 来说，MNIST 的 60,000 张 28×28 灰度图、CIFAR 的 50,000 张 32×32 彩色图、ImageNet 的 128 万张高分辨率 JPEG，获取方式都是同一个 get_next_sample()。后端的差异被彻底封在 Loader 内部。
三、RAW 与 DTS：两种加载路径
Tech-Renaissance 为每个内置数据集都提供了两种加载器：RAW 和 DTS，总共六种实现：
| 数据集 | RAW 格式 Loader | DTS 格式 Loader |
|---|---|---|
| MNIST | MnistLoaderRaw | MnistLoaderDts |
| CIFAR-10/100 | CifarLoaderRaw | CifarLoaderDts |
| ImageNet | ImageNetLoaderRaw | ImageNetLoaderDts |
RAW 加载器读取数据集的公开原始格式。MNIST 对应官方的 train-images-idx3-ubyte / train-labels-idx1-ubyte；CIFAR 对应 data_batch_1.bin ~ data_batch_5.bin / test_batch.bin；ImageNet 对应原始 JPEG 文件夹，外加一个预生成的 summary.bin 索引文件，记录每张图片的路径、标签、文件大小等元信息。RAW 的优点是无需额外转换，拿到数据集就能跑；缺点是小文件随机 IO 和解码开销较大。
DTS 加载器读取 Tech-Renaissance 自研的 .dts 打包格式。DTS 会在下一篇文章里详细讲，这里只说它对 Loader 接口的影响：DTS 文件有固定的 header、CRC-32 校验、可选压缩，并且按照 Block/Slot 的方式组织数据。对 Loader 而言，DTS 路径只需要跳过 header、按索引读取块即可，逻辑比解析 JPEG 更简单。MNIST/CIFAR 的 DTS 文件同样只有 256 字节 header，后面紧跟全部样本数据；ImageNet 的 DTS 则按 16MB 的 Block 组织，支持 FULLY 全量加载和 PARTIAL 流式加载两种策略。
两种路径通过 using_dts_format(true, compression_level) 一键切换：
PREPROCESSOR_SETTING
.dataset("imagenet", "/data/imagenet")
.using_dts_format(true, 3) // 启用 DTS，压缩级别 3
.fully_mode(false) // PARTIAL 流式加载
.commit();
同一个 DataLoader 抽象，让 RAW 和 DTS 可以无缝互换。这是框架层抽象的价值：上层不需要关心数据是以什么物理形式存在，只需要保证接口语义一致。
四、FULLY 与 PARTIAL：小数据集全进内存，大数据集流式加载
LoadMode 是数据加载的另一个重要维度，定义在 include/renaissance/core/global_config.h：
enum class LoadMode {
AUTO, // 自动选择
FULLY, // 全量加载到内存
PARTIAL // 环形缓冲区循环加载
};
对于 MNIST 和 CIFAR，Tech-Renaissance 会强制使用 FULLY 模式。原因很简单：训练集只有几十 MB 到几百 MB，一次性读进内存后，后续每个 epoch 都不再触碰磁盘，shuffle 也可以在内存中完成，速度最快。FULLY 模式下，第一个 epoch 加载完成后，Loader 会构建一份 SampleInfo 表，记录每个样本的标签、数据指针、大小；从第二个 epoch 开始，worker 直接按表领取，几乎是纯内存访问。
MNIST 的 FULLY 加载尤为简单：解析 .ubyte 文件头（magic number + 维度信息）后，将图像和标签分别存入 images_region 和 labels_region 两块连续内存中。CIFAR 的 FULLY 加载则多一步格式转换：官方二进制文件中图像是 CHW 格式（通道–高–宽），而框架统一使用 NHWC 布局。这个格式转换必须做，但关键问题是什么时候做。如果做在预处理阶段，每次取样本都要重新排列一次，是重复计算。Tech-Renaissance 的选择是：做在加载阶段。CifarLoaderRaw 在加载时一次性完成 CHW→HWC 转换，之后所有访问拿到的都是 NHWC 格式，保证了后续链路中的格式一致性。此外，CIFAR-10 和 CIFAR-100 的文件结构完全相同，只是 class 数量不同，CifarLoaderRaw 通过自动检测类别数来共用同一套代码，不需要两个独立的 Loader 类。
对于 ImageNet，默认使用 PARTIAL 模式。因为 128 万张原始 JPEG 全塞进内存可能要占用数百 GB，不是每台机器都受得了。PARTIAL 模式采用双缓冲/环形缓冲思想：Loader 每次只把一部分数据读入内存缓冲区，Preprocessor 消费完一块后，再异步加载下一块。ImageNet 的 RAW 版本按 64MB 的 Part Slot 切分（共 16 个 PART），DTS 版本按 16MB 的 Block 切分，两者都支持后台 IO 线程持续填缓冲区，前台 Preprocessor 持续消费，实现 IO 与预处理的重叠。
这里值得一提的是，FULLY/PARTIAL 不是简单地在”快但费内存”和”慢但省内存”之间二选一。PARTIAL 模式下的关键优化是静态领取：每个 Preprocessor worker 知道自己该取第几个样本，不需要抢一个全局原子计数器，从而把并发竞争降到最低。
五、Epoch 边界与 Shuffle：确定性随机
每个 epoch 开始时，数据需要重新打乱。这是深度学习训练的基本操作——如果每个 epoch 都按相同顺序喂数据，模型容易学到”顺序相关”的 bias，影响泛化。Tech-Renaissance 的 shuffle 机制分为两级：
第一级：全局 shuffle。在 begin_epoch() 中，框架用 Philox 随机数生成器对样本索引进行全局打乱。Philox 是基于计数器的——给定相同的 seed 和 epoch_id，生成的随机序列完全一致，无论用多少个线程、在什么平台上运行。这就是”确定性训练”的基石之一。
以 MNIST 为例，begin_epoch() 的核心流程是：先调用 perform_global_shuffle() 将 60000 个样本的 SampleInfo 数组打乱，然后调用 distribute_to_threads() 将打乱后的索引均匀分配给 M 个 Preprocessor worker：
void MnistLoaderRaw::begin_epoch(int epoch_id, bool is_train) {
// 全局 shuffle
perform_global_shuffle(global_info, epoch_id);
// 分配给 M 个 worker
distribute_to_threads(global_info, thread_info);
}
第二级：静态领取。分配之后，每个 worker 拿到的是自己独享的一份索引列表。Worker i 的第 k 次调用 get_next_sample()，就读取自己列表中的第 k 个元素。这种设计被称作”静态样本领取机制”——每个 worker 负责固定间隔的样本（i + k×M），不存在竞争。Worker 之间不需要任何锁或同步——它们各自拿着自己的索引列表往前走，互不干扰。这是 Tech-Renaissance 多线程预处理架构的核心设计，由第九篇详细展开。但它的根基——索引的分配——就发生在 begin_epoch() 这个阶段。
epoch 边界处理：get_next_sample() 返回 false 时，表示当前 epoch 的所有样本已被消费完毕。此时 Preprocessor 会调用 end_epoch()，然后进入下一个 begin_epoch()，开始新一轮 shuffle+分配。对于 PARTIAL 模式（ImageNet），每个 epoch 的 shuffle 是在”虚拟索引”层面完成的——shuffled_locations 数组记录了所有样本的访问顺序，DataLoader 按这个顺序去对应的 Buffer 和 Slot 中取数据。
configure() 中还有一个 skip_first 参数。当设置为 true 时，第一个 epoch 不执行 shuffle——这在 warmup 阶段很有用，可以快速验证数据管线是否正常工作，而不需要等待 shuffle 的计算开销。
六、标签管理：容易被忽视的细节
标签管理看似简单，但有几个容易踩坑的细节，Tech-Renaissance 在设计时都做了处理。
标签类型统一为 int32。get_next_sample() 返回的标签固定为 int32_t 类型。Tech-Renaissance 的 DType 系统只有 FP32/FP16/INT8/INT32（见第七篇），标签统一为 int32 既与框架内 top-1/top-5 指标算子的期望类型一致，也比 int64 减半标签传输量。PyTorch 的 nn.CrossEntropyLoss 在处理类别索引（硬标签）时强制要求标签为 int64（LongTensor），在本框架中我们通过统一 int32 避免了这一约束带来的额外带宽开销。
class_to_idx 映射的一致性。对于 ImageNet 这类文件夹结构的数据集，框架在扫描目录时自动构建 label_to_folder 和 folder_to_label 两个映射表。训练集和验证集共用同一个映射表，确保标签空间一致——不会出现”训练集里 n01440764 是第 0 类，但验证集里它变成了第 1 类”这种隐性错误。这是一个在 PyTorch 社区中反复被问到的问题——很多人在 ImageNet 上训练时发现验证准确率很低，排查到最后发现是 class_to_idx 映射不一致导致的。
CRC-32 完整性校验。DTS 格式的 Loader 支持 verify_dts_crc() 方法，对 DTS 文件进行 CRC-32 校验。DTS 文件的文件头中存储了 CRC-32 校验码，加载后会重新计算并与文件头中的值比对。这保证了数据在传输和存储过程中没有损坏——标签和图像数据因此不会出现”静默错误”。
七、Preprocessor：把数据集配置翻译成一条流水线
DataLoader 负责把数据从磁盘搬到内存，Preprocessor 则负责把原始样本变成 GPU 能直接消费的 batch。它通过一个流畅的 Setup 构建器对外暴露配置：
PREPROCESSOR_SETTING
.dataset("imagenet", "/data/imagenet")
.preprocess_workers(128)
.normalization(NormMode::IMAGENET)
.train_transforms(
RandomResizedCrop(224, 0.08f, 1.0f),
RandomHorizontalFlip(),
ColorJitter(0.2f, 0.2f, 0.2f, 0.1f),
RandomErasing(0.25f, {0.02f, 0.33f}, {0.3f, 3.3f})
)
.val_transforms(
Resize(256),
CenterCrop(224)
)
.commit();
这个配置做了几件事：指定数据集、设置预处理 worker 总数、选择归一化预设、分别定义训练与验证的变换链。commit() 会按照一个内部状态机依次完成：选择 DataLoader、下载/解压/校验、配置 Loader 模式、计算 batch/分辨率/缓冲区大小、分配 staging memory、注入归一化、注册变换链、计算 steps_per_epoch。
注意其中的 normalization(NormMode::IMAGENET)。框架不允许用户把 Normalize 或 FusedNormalization 直接写进 train_transforms，而是强制通过 NormMode 来指定。原因是 Tech-Renaissance 会把 Normalize / 除以 255 / 数据类型转换 这些步骤在 PO 链末尾统一合并成一次 FusedNormalization，避免中间结果反复读写。这是数据侧的一项重要优化，第十篇文章会专门讲它，这里先埋下伏笔。
所有具体变换都继承自 PreprocessOperation：
virtual void execute(
const uint8_t* input_ptr,
int32_t input_width, int32_t input_height, size_t input_stride,
uint8_t* output_ptr,
int32_t& output_width, int32_t& output_height, size_t& output_stride,
Generator* rng = nullptr,
bool execute_from_full = false,
bool forced_compact_output = true) = 0;
每个 Preprocessor worker 在启动时都会克隆一份自己的 PO 副本，拥有独立的输入/输出缓冲区和 TurboJPEG 句柄。worker 之间不共享状态，也就没有 false sharing 和锁竞争。这种”无共享”设计，是把预处理扩展到 128 甚至 200+ 线程的基础。
八、TransferStation：CPU 与 GPU 之间的双缓冲驿站
预处理完成后，数据还需要到达 GPU。Tech-Renaissance 没有让用户手动调用 tensor.to(device)，而是把 H2D 传输也封装进了框架内部，核心就是 TransferStation。
TransferStation 是每个 GPU Engine 一个的双缓冲管理器。它在 configure() 阶段从 GlobalRegistry 的 staging memory pool 里拿到两块已经分配好的锁页内存，每块被划分为标签区和图像数据区，并且都按 256 字节对齐：
std::atomic<int> current_buffer_{0};
int32_t* buffer_labels_[2] = {nullptr, nullptr};
uint8_t* buffer_data_[2] = {nullptr, nullptr};
std::atomic<bool> buffer_0_is_readable_{false};
std::atomic<bool> buffer_1_is_readable_{false};
std::atomic<bool> buffer_0_is_writeable_{true};
std::atomic<bool> buffer_1_is_writeable_{true};
工作方式很直观：
- Preprocessor worker 把处理好的样本写入当前 buffer 的对应位置；
- 当当前 buffer 被一个 batch 的样本填满后，
TransferStation把它标记为 readable，并通知 GPU 侧读取； - 同时切换到另一块 buffer，继续让 Preprocessor 写入；
- GPU 读取完毕后，把用过的 buffer 重新标记为 writeable，等待下一轮填充。
这样 CPU 预处理、CPU→GPU 异步拷贝、GPU 计算这三个阶段就可以流水起来：当 GPU 在计算 buffer A 对应的 batch 时，CPU 正在往 buffer B 里填下一个 batch。双缓冲消除了 CPU 等待 GPU、GPU 等待 CPU 的串行空闲。
对齐要求也不容忽视。单区大小计算为：
label_aligned_ = align_up_256(local_batch_size * sizeof(int32_t) + 16); data_aligned_ = align_up_256(local_batch_size * max_sample_bytes + 16); bytes_per_transfer_ = label_aligned_ + data_aligned_;
之所以加 16 字节再向上对齐到 256 字节，是为了与 GPU 端 DTensor 的内存布局兼容，确保 H2D 异步传输后不需要再做数据重排。这个细节看起来很小，但它体现了 Tech-Renaissance 一贯的作风：数据侧的任何布局都要和显存侧的布局对齐，否则性能会一点点漏掉。
九、与 PyTorch DataLoader 的几点区别
写到这里，可能有人会问：这不就是一个 C++ 版的 PyTorch DataLoader 吗？功能上确实有相似之处，但设计定位完全不同。
第一，worker 总数的语义不同。PyTorch 的 num_workers 通常是每张 GPU 的进程数；而 Tech-Renaissance 的 preprocess_workers 是跨所有 GPU 的总线程数。例如 8 卡云服务器、PyTorch 每张卡 16 个 worker，总共是 128 个进程；在 Tech-Renaissance 里直接写 preprocess_workers(128)。框架内部再通过 engine_id = worker_id % world_size 把样本划分到各个 GPU。这个设计让多卡场景下的 CPU 资源分配更统一，也便于做跨卡的 NUMA/CPU 绑核。
第二，数据格式原生支持。PyTorch 依赖用户自己写 Dataset 来读原始文件或自定义二进制；Tech-Renaissance 内置了 RAW 和 DTS 两条路径，ImageNet RAW 还需要预生成 summary.bin 索引。DTS 不是简单的打包，它会把下一篇文章要讲的 Block 索引、CRC 校验、压缩策略都考虑进去，让 Loader 可以按块随机访问。
第三，GPU 传输内聚。PyTorch DataLoader 把预处理后的 tensor 传回主进程，用户再自己 to(device) 发起 H2D；Tech-Renaissance 的 TransferStation 直接和深度学习引擎共享 staging memory，H2D 由运行时的 CUDA Graph 子图统一调度，不需要 Python 侧参与。
第四，可复现性贯穿全程。Tech-Renaissance 的 shuffle、数据增强、worker 领取顺序都基于 Philox 计数器随机数，保证同样种子下每次运行结果一致。这一点在第 22 篇会专门展开，但在数据管线里它已经体现为”静态领取”和”确定性 shuffle”的设计。
十、小结：数据管线是性能的第一道关卡
数据加载管线看起来不如算子融合、CUDA Graph 那么”硬核”，但它决定了 GPU 能不能满负荷运转。Tech-Renaissance 在这一层的设计思路可以概括为三句话：
- 统一抽象：用
DataLoader一个接口覆盖 MNIST、CIFAR、ImageNet，RAW 与 DTS 随意切换，共六种 Loader 实现； - 按需加载：小数据集全进内存（FULLY），大数据集流式双缓冲（PARTIAL），CIFAR 的 CHW→HWC 转换做在加载阶段而非预处理阶段；
- 零拷贝贯通：从 Loader 内部缓冲区到 Preprocessor worker，再到 TransferStation 的 staging memory，最后通过 H2D 进入 GPU 计算图，中间尽量少做内存复制。
这套管线的下一环，就是 DTS 格式本身。为什么我们要自研数据格式？DTS 的文件头、Block 索引、CRC 校验、FULLY/PARTIAL 策略到底长什么样？它又是怎么把磁盘布局的精心设计转化为内存带宽节省的？这些问题，我们下一篇《DTS 数据格式：为高速训练定制的存储方案》里接着聊。
