——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之九
上一篇我们讲了 Tech-Renaissance 的数据加载管线抽象:无论是 MNIST、CIFAR 还是 ImageNet,框架都用同一套 DataLoader + Preprocessor + TransferStation 接口把它们串起来。那篇文章侧重的是“接口长什么样”,这一篇我们往下一层钻,看看数据在磁盘上到底是怎么存的。
数据格式这件事,看起来远不如 CUDA Graph 或算子融合那么耀眼。但如果你真把训练吞吐当回事,它往往是第一道关卡。道理很简单:GPU 算得再快,如果 CPU 喂数据跟不上,它就只能空转。而在 ImageNet 这种百万级小文件数据集上,“喂数据”的瓶颈常常不在解码、不在预处理,而在磁盘怎么读。读一百万个 JPEG 文件,和读几个顺序排放的大文件,完全是两回事。
Tech-Renaissance 为此自研了一套名为 DTS 的数据存储格式。DTS 不是通用的键值存储,也不是某个流行格式的简单封装,而是专门为框架的 NHWC 训练管线、双缓冲加载、CUDA Graph 全捕获量身定制的二进制格式。它的核心思想可以概括成一句话:在导出阶段用一次性的磁盘布局优化,换取训练阶段持续的高带宽、低延迟、零竞争数据读取。
一、数据格式为什么是第一道性能关卡
1.1 数据饥饿:GPU 在等 CPU
训练一个神经网络,数据要经历的链路大致是:磁盘 → 内存 → CPU 预处理 → 锁页内存 → GPU 显存 → 计算。这条链路上任何一环落后,都会形成瓶颈。对于 CNN 图像分类任务,当模型不大、batch 不小、GPU 很强时,链路里最容易落后的往往是磁盘读取 + 文件系统开销。
以 ImageNet 为例。训练集约 128 万张 JPEG,平均几十到几百 KB 一张。如果你用 PyTorch 常见的 ImageFolder + DataLoader 方式去读,每个 worker 每次都要发起一次文件系统调用:定位目录、打开文件、读取数据、关闭文件。128 万次这样的调用分布在多个 epoch 里,文件系统元数据操作的总量非常可观。机械硬盘上随机小文件读取的 IOPS 很低,即使 SSD,百万次 open/close 也绝不是免费午餐。
这就是常说的“数据饥饿”(Data Starvation):GPU 算完一个 batch 只要几十毫秒,CPU 却要几百毫秒才能把下一个 batch 的图片从磁盘捞出来。
1.2 主流方案各有利弊
业界针对这个问题,已经沉淀出几种典型方案。
原始文件夹 + 独立 JPEG:最直观,数据保持原样,任何人都看得懂。缺点是训练时大量小文件随机 IO,文件系统压力大;num_workers 开得再高,也可能被 inode 查找和随机 seek 拖住。
LMDB:把键值对打包进一个 memory-mapped 数据库文件。优点是单文件、可 mmap、读取快;缺点是对 Windows 支持不够原生,数据导入过程比较重,而且它的键值语义比较通用,不会为你的图像训练管线做任何专门优化。
TFRecord:TensorFlow 生态常用的二进制记录容器。TFRecord 本身是一种带长度和 CRC 的 record framing 格式,实际训练中常把序列化后的 tf.train.Example / SequenceExample(Protocol Buffer)作为 payload 存进去。优点是顺序读取和 tf.data 集成成熟;缺点是如果使用 protobuf payload,会有相应的反序列化开销,跨框架使用也不如原始图片或简单二进制格式直接。
WebDataset:把样本打包成 tar 归档,图片、标签作为 tar 里的独立文件。优点是天然适合流式读取和云存储;缺点是 tar 格式按顺序扫描,随机访问和 shuffle 需要做额外工作,某些实现下 worker 之间有额外的数据拷贝。
FFCV 的 .beton:这是近年来对 ImageNet 训练提速影响最大的方案之一。FFCV 把数据集写成自定义的 .beton 文件,按固定大小的 page 组织样本,可以选择 JPEG 编码、原始像素或混合编码。它的核心洞察是:把训练时要读的数据在磁盘上排得尽可能紧密,一次大 IO 读多个样本,减少文件系统元数据开销和随机 seek。 Tech-Renaissance 的 DTS 在 ImageNet 制作策略上对标了 FFCV 的思路:提前做缩放、压缩、分块,让训练时的 IO 带宽接近顺序读取峰值。
这些方案各有舞台,但通用格式总要考虑兼容性、生态、跨平台。Tech-Renaissance 作为一个自研训练框架,没有历史包袱,于是选择了一条更直接的路:为自身训练管线设计一个专用格式。
二、DTS 的设计目标
在动手写格式之前,我们先把 DTS 要解决的问题列清楚:
- 顺序大 IO 替代小文件随机 IO:训练时不应再接触文件系统元数据,IO 单元尽可能大。
- 零拷贝暴露样本:加载器读到内存后,应能直接把指针交给预处理器,不需要额外 malloc 和 memcpy。
- 多级压缩可选:用户可以在“原始画质”和“训练速度”之间做 trade-off。
- 可复现的 shuffle:大规模并发加载不能破坏训练的可复现性。
- 内存策略灵活:内存足够时全量加载,内存不足时双缓冲循环加载。
- 数据完整性校验:一百多 GB 的数据集,传输出错要能及时发现。
- 跨平台原生:Windows 和 Linux 都走 Native API,不依赖 mmap 或特定库。
这七条目标贯穿了 DTS 从文件头到 Block 布局、从导出脚本到 C++ 加载器的全部设计。
三、DTS 文件结构
3.1 文件头:一个 144 字节的“身份证”
DTS 文件头采用 #pragma pack(push, 1) 严格打包,固定 144 字节。定义在:
// include/renaissance/data/imagenet_loader_dts.h:57-91
#pragma pack(push, 1)
struct DtsHeader {
char magic[4]; // ".DTS"
uint8_t version[4]; // [3, 0, 0, 0]
char dataset_type[8]; // "IMAGENET" / " CIFAR10" / "CIFAR100" / " MNIST"
uint32_t is_training; // 0=val, 1=train
uint32_t compress_level;
uint32_t val_set_prep;
uint32_t num_classes;
char tensor_layout[4]; // "NHWC"
uint32_t image_width;
uint32_t image_height;
uint32_t num_channels;
char color_type[4]; // " RGB" / "GRAY"
uint32_t num_samples;
uint32_t num_volumes;
uint32_t volume_id;
uint32_t total_blocks;
uint32_t num_blocks;
uint64_t total_bytes;
uint32_t header_bytes;
uint64_t block_bytes;
uint32_t block_size; // 16MB
uint32_t block_header_size; // LV0=4KB, LV123=16KB
uint32_t pic_alignment; // 64
uint32_t max_pic_area;
uint32_t max_pic_per_block;
float compression_ratio;
float normalize_mean[3];
float normalize_std[3];
uint32_t crc_code; // 数据区 CRC-32
};
#pragma pack(pop)
static_assert(sizeof(DtsHeader) == 144, "DtsHeader must be exactly 144 bytes");
这个头里藏了很多关键信息:数据集类型、训练/验证、压缩级别、类别数、NHWC 布局、块大小、对齐方式、CRC 校验码等等。它既是文件的“身份证”,也是加载器初始化时唯一需要解析的元数据。
需要注意的是,DTS 文件头有两种物理尺寸:
- ImageNet:16 MB(
IMAGENET_HEADER_SIZE) - MNIST/CIFAR:256 字节(
CIFAR_MNIST_HEADER_SIZE)
MNIST 和 CIFAR 数据量小,256 字节足矣;ImageNet 头部扩到 16 MB 是为了让 Block 从 16 MB 对齐边界开始,便于大块预读和后续可能的扩展。
3.2 Block:16 MB 的固定大小 IO 单元
DTS 把数据按固定 16 MB 的 Block 组织。以 ImageNet LV0(无压缩)为例,一个 Block 的物理布局是:
[block_magic(4B)] [block_id(4B)] [num_pics(4B)] [offsets 数组] // 每个样本在 Block 内的偏移 [sizes 数组] // 每个样本的字节数 [labels 数组] // 每个样本的标签 [JPEG 数据流 + 64B 对齐填充] [零填充至 16 MB]
block_magic 是 4 字节魔数:LV0B、LV1B、LV2B、LV3B,分别对应四个压缩级别。Block 头部本身也有两种尺寸:LV0 用 4 KB,LV1~LV3 用 16 KB,因为压缩后单 Block 内样本数上限可能更高,需要更大的索引表。
固定 16 MB 的好处是显而易见的:
- 顺序读取友好:磁盘/SSD 处理大顺序块的吞吐远高于小随机读。
- 静态分配简单:加载器在初始化时就知道每个 buffer 需要多少内存,不需要运行时根据样本大小动态调整。
- 多线程无竞争:每个 IO worker 负责固定的一组 Block,没有锁、没有队列竞争。
以 ImageNet LV0 训练集为例,平均每个 Block 装约 146 张图片,最多 177 张;空间浪费控制在约 0.08% 左右(与 5.1 节 13 KB / 16 MB 口径一致)。
3.3 SlotMeta:零拷贝定位每个样本
每个 Block 的元数据被解析成 SlotMeta:
// include/renaissance/data/imagenet_loader_dts.h:97-107
struct SlotMeta {
uint32_t block_id = UINT32_MAX;
uint32_t num_samples = 0;
uint32_t offsets[MAX_SAMPLES_PER_BLOCK]; // 样本在 Block 内偏移
uint32_t sizes[MAX_SAMPLES_PER_BLOCK]; // 样本字节数
int32_t labels[MAX_SAMPLES_PER_BLOCK]; // 样本标签
};
加载器只需要根据 slot_idx * BLOCK_SIZE + offset 就能直接拿到样本的字节流指针。这个指针可以直接交给 Preprocessor,实现零拷贝:
// src/data/imagenet_loader_dts.cpp:1068-1075
const SlotMeta& smeta = buffer_meta.slot_metas[slot_idx];
label = smeta.labels[sample_idx];
data_size = smeta.sizes[sample_idx];
data_ptr = current_set_->full_arena +
global_slot_idx * block_size_ +
smeta.offsets[sample_idx];
data_ptr 指向的是加载器内部 arena 的地址,没有额外堆分配,没有 std::vector 扩容,也没有 memcpy。
四、ImageNet 的四级压缩:画质、体积与速度
DTS 为 ImageNet 提供了四个压缩级别,从 LV0 到 LV3,逐级增加预处理强度。整个导出逻辑都写在仓库的 python/scripts/make_dataset.py 中,完全开源。由于导出流程使用固定随机种子、固定缩放算法(PIL LANCZOS)和固定 JPEG 质量参数,在相同的依赖环境下,只要原始 ImageNet 数据集一致,多次运行生成的 DTS 文件就是字节级一致的。这意味着团队内部可以只分发一次 DTS 文件,所有成员、所有机器都能拿到完全相同的训练输入。
四个级别的具体含义如下:
| 级别 | 处理方式 | 短边限制 | 长边限制 | JPEG 质量 | 与原版 ImageNet 的关系 |
|---|---|---|---|---|---|
| LV0 | 仅做拼接、对齐,不做任何图像处理 | 原图 | 原图 | 原图 | 字节级等价于原版 JPEG |
| LV1 | 短边超过阈值则按比例缩放到 400 px | 400 px | 原比例 | 90 | 有损缩放,保留全部画面内容 |
| LV2 | 在 LV1 基础上对长边做中心裁剪,限制到 600 px | 400 px | 600 px | 90 | 去除边缘,进一步减小体积 |
| LV3 | 在 LV2 基础上把 JPEG 质量降到 80 | 400 px | 600 px | 80 | 体积最小,训练速度最快 |
这里需要特别强调:LV0 是完全无压缩的。它不会重新编码图片,只是把原始的 JPEG 字节流原封不动地拼接进 16 MB 的 Block,再补上 64 字节对齐和索引信息。因此,LV0 的 DTS 文件与 ImageNet 原版数据集在图像内容上是完全一致的——你可以把它理解为“把 128 万张 JPEG 打包成一个文件”。这也解释了为什么 LV0 的 DTS 仍然比 RAW 文件夹更快:快的不是解码,而是读取方式。
LV1~LV3 则借鉴了 FFCV .beton 的思路:训练时最终会把图片 RandomResizedCrop 到 224×224,原图的几百万像素大部分都会被丢掉,那不如在导出阶段先做合理的下采样和裁剪,减少磁盘 IO 量,同时不显著影响最终精度。具体实现位于 process_imagenet_image:
# python/scripts/make_dataset.py:891-950
def process_imagenet_image(src_path):
if compress_level == 0:
# LV0:不做任何处理,直接返回原文件字节流
with open(src_path, 'rb') as f:
return f.read(), 0
else:
with Image.open(src_path) as im:
w, h = im.size
if w > h:
# 短边超过阈值时进行 LANCZOS 缩放
if h > crop_boundary:
scale = SHORT_LIMIT * 1.0 / h
new_w, new_h = int(w * scale), int(h * scale)
im = im.resize((new_w, new_h), algorithm)
# LV2/LV3:如果长边仍超过限制,做中心裁剪
if compress_level >= 2:
final_w = min(new_w, LONG_LIMIT)
final_h = min(new_h, SHORT_LIMIT)
...
buf = BytesIO()
if im.mode in ('RGBA', 'LA', 'P'):
im = im.convert('RGB')
im.save(buf, format='JPEG', quality=jpeg_quality)
return buf.getvalue(), output_w * output_h
实际测试中,ImageNet 训练集 LV0 的大小约 147 GB(约 137 GiB),LV3 可以把它压缩到几十 GB。对于追求极致读取速度、磁盘容量又有限的场景,LV2/LV3 是很有吸引力的选择。我们的 VGG16BN ImageNet 示例默认使用 LV3,以换取更小的磁盘占用和更高的顺序读取带宽;如果用户更在意图像质量或需要与原版 ImageNet 做严格对齐,也可以切换回 LV0 或 LV1。
五、导出阶段的三级随机性
深度学习训练需要 shuffle,但大规模并发加载下的 shuffle 要保证可复现并不容易。DTS 把随机性拆成了三级:
5.1 导出级:greedy arrange + random fill
在制作 DTS 文件时,我们用一种类似装箱的算法把不同大小的图片安排进 16 MB Block:
# python/scripts/make_dataset.py:1149-1195
def greedy_arrange(info_list, minimum_num_blocks, train):
random.seed(RAND_SEED)
if USING_PIC_ALIGNMENT:
apply_pic_alignment_to_all(info_list, PIC_ALIGNMENT_BYTES)
blocks, block_margin = make_blocks_l0(BLOCK_SIZE, HEADER_SIZE_LV0, minimum_num_blocks)
ffd_fill(blocks, block_margin, info_list, minimum_num_blocks)
random_fill(blocks, block_margin, info_list, minimum_num_blocks)
for k in range(MARGINAL_FILL_ITERATIONS):
block_margin, blocks = map(list, (zip(*sorted(zip(block_margin, blocks)))))
marginal_fill(blocks, block_margin, info_list, minimum_num_blocks)
extra_block_fill(blocks, block_margin, info_list, minimum_num_blocks)
return blocks, block_margin
这个算法先用 First-Fit Decreasing(FFD)把大图片放进 Block,再用随机填充把小空隙填满,最后做 32 轮余量精修。为什么要搞得这么复杂?因为 ImageNet 图片的原始大小差异很大:有的只有几十 KB,有的超过 1 MB。如果直接按文件顺序依次填充,最后几个 Block 会极其稀疏,造成大量空间浪费;而把大文件和小文件“混装”,才能把 16 MB 的容器填得更满。
具体流程如下:
- 初始化 Block:根据最小 Block 数(训练集 8700、验证集 400)创建空 Block,并留出 4 KB 的 Block 头部空间。
- First-Fit Decreasing:把所有图片按大小从大到小排序,依次放入“第一个能装得下它”的 Block。FFD 是经典装箱问题的近似最优贪心策略,能在较短时间内得到一个相当紧凑的初始布局。
- 随机填充:打乱剩余小图片的顺序,按顺序尝试填入当前仍有空间的 Block。这一步避免 FFD 可能产生的“前满后空”的系统性偏差。
- 余量精修:把 Block 按剩余空间从小到大排序,反复尝试把剩余图片塞进缝隙,迭代 32 轮。这是压缩率提升最关键的一步。
- 新增 Block:如果仍有图片装不下,再动态追加 Block,直到所有图片都有位置。
最终效果非常理想:ImageNet LV0 训练集约 8768 个 Block,平均每个 Block 只剩 13 KB 余量,空间浪费控制在约 0.08%(13 KB / 16 MB)左右。也就是说,DTS 用几乎可忽略不计的额外存储,换来了完全顺序化的 IO 模式。
5.2 Block 级:Level 2 shuffle
每个 epoch 开始时,加载器会对 Block 的顺序做 shuffle:
// src/data/imagenet_loader_dts.cpp:566-606
bool should_shuffle = is_train ? shuffle_train_ : shuffle_val_;
if (should_shuffle) {
perform_level2_shuffle(*current_set_, epoch_id);
}
这一步保证相邻 Block 在磁盘上不必相邻读取,但每个 Block 内部仍是顺序 IO。种子由全局 seed 和 epoch_id 共同决定,固定 seed 即可复现。
5.3 样本级:Level 3 shuffle
Block 读入内存后,再对每个 buffer 内的样本做 Fisher-Yates 洗牌:
// src/data/imagenet_loader_dts.cpp:1658-1669
uint64_t base_seed = tr::get_default_generator().seed();
uint64_t seed = base_seed ^
(static_cast<uint64_t>(current_epoch_id_.load(std::memory_order_relaxed)) << 32) ^
(static_cast<uint64_t>(start_group_idx) << 16);
shuffle_samples(buffer->shuffled_locations, seed);
三级随机配合 Philox 计数器随机数(系列第 22 篇会详细展开),既能保证统计意义上的随机性,又能在固定 seed 下精确复现。
5.4 为什么 DTS 加载更快:实测数据说话
到这里,我们可以系统性地回答一个问题:DTS 到底比直接从文件夹读 RAW 图片快在哪里?
从工程角度看,DTS 的优势来自几个叠加效应:
- 顺序 IO 替代随机 IO:RAW 文件夹需要频繁地在 1000 个类别目录之间跳转,产生大量随机 seek 和文件系统元数据操作;DTS 把数据合并成少量大文件,读取基本是纯顺序的。
- Block 索引实现 O(1) 定位:给定样本序号,DTS 直接算出 Block 编号和内部偏移,不需要遍历目录、解析文件名或查询数据库。
- Native IO + 大分块:DTS 使用
pread(Linux)或ReadFile(Windows),以 4 MB 为单位读取,减少了系统调用次数;pread不依赖共享文件偏移,天然支持多线程无锁并发。 - 零拷贝样本暴露:数据从磁盘读进 Block buffer 后,
get_next_sample直接把指针交给Preprocessor,没有额外的malloc/memcpy。 - 压缩减少磁盘流量:LV1~LV3 在导出阶段完成缩放、裁剪和 JPEG 质量降低,训练时读取的字节数显著减少;现代 CPU 用 TurboJPEG 解码这些缩小后的 JPEG 非常轻松,省下的磁盘带宽远大于新增的解码开销。
这些设计叠加在一起,效果非常显著。仓库中的 docs/EXTREME_PERFORMANCE.md 记录了一次完整 epoch(train + val)的加载耗时对比,测试平台是智星云 GPU 服务器:CPU 为 2× Intel Xeon Platinum 8480C(112 核 / 224 线程),内存 512 GB,GPU 为 8× A100 40GB,操作系统 Ubuntu 24.04;测试线程配置为 16 个加载线程 + 96 个预处理线程。
无缓存预热(冷启动)结果,单位:秒
| Dataset | Format | Mode | TRAIN | VAL | TRAIN+VAL |
|---|---|---|---|---|---|
| MNIST | RAW | FULLY | 0.203 | 0.122 | 0.325 |
| MNIST | DTS | FULLY | 0.168 | 0.113 | 0.281 |
| CIFAR-10 | RAW | FULLY | 0.701 | 0.204 | 0.905 |
| CIFAR-10 | DTS | FULLY | 0.273 | 0.135 | 0.408 |
| CIFAR-100 | RAW | FULLY | 0.532 | 0.200 | 0.732 |
| CIFAR-100 | DTS | FULLY | 0.243 | 0.131 | 0.374 |
| ImageNet | RAW | FULLY | 92.648 | 4.252 | 96.900 |
| ImageNet | DTS LV0 | FULLY | 76.632 | 3.607 | 80.239 |
| ImageNet | DTS LV1 | FULLY | 35.220 | 1.602 | 36.822 |
| ImageNet | DTS LV2 | FULLY | 35.033 | 1.616 | 36.649 |
| ImageNet | DTS LV3 | FULLY | 24.620 | 1.147 | 25.767 |
| ImageNet | RAW | PARTIAL | 91.593 | 4.238 | 95.831 |
| ImageNet | DTS LV0 | PARTIAL | 75.726 | 3.571 | 79.297 |
| ImageNet | DTS LV1 | PARTIAL | 35.890 | 1.687 | 37.577 |
| ImageNet | DTS LV2 | PARTIAL | 36.249 | 1.635 | 37.884 |
| ImageNet | DTS LV3 | PARTIAL | 25.408 | 1.163 | 26.571 |
从冷启动数据可以观察到三个关键结论:
- 即使是 LV0(完全无压缩),DTS 也比 RAW 快:ImageNet FULLY 模式下,RAW 需要 96.9 秒,而 DTS LV0 只需 80.2 秒;PARTIAL 模式下 RAW 95.8 秒,DTS LV0 79.3 秒。这说明顺序读取 + O(1) 索引 + Native IO本身就足够产生显著收益。
- 压缩级别越高,加载越快:DTS LV3 在 FULLY 模式下把训练集加载时间从 92.6 秒降到 24.6 秒,PARTIAL 模式下从 91.6 秒降到 25.4 秒。磁盘读取字节数的减少对速度的影响非常直接。
- 小数据集同样受益:CIFAR-10 FULLY 模式下,DTS 把 0.905 秒降到 0.408 秒,接近 2.2 倍;CIFAR-100 从 0.732 秒降到 0.374 秒。
有缓存预热(OS 文件系统缓存已热)结果,单位:秒
| Dataset | Format | Mode | TRAIN | VAL | TRAIN+VAL |
|---|---|---|---|---|---|
| MNIST | RAW | FULLY | 0.200 | 0.124 | 0.324 |
| MNIST | DTS | FULLY | 0.151 | 0.110 | 0.261 |
| CIFAR-10 | RAW | FULLY | 0.606 | 0.178 | 0.784 |
| CIFAR-10 | DTS | FULLY | 0.236 | 0.124 | 0.360 |
| CIFAR-100 | RAW | FULLY | 0.486 | 0.187 | 0.673 |
| CIFAR-100 | DTS | FULLY | 0.237 | 0.123 | 0.360 |
| ImageNet | RAW | FULLY | 16.369 | 0.708 | 17.077 |
| ImageNet | DTS LV0 | FULLY | 13.757 | 0.679 | 14.436 |
| ImageNet | DTS LV1 | FULLY | 7.422 | 0.371 | 7.793 |
| ImageNet | DTS LV2 | FULLY | 6.516 | 0.355 | 6.871 |
| ImageNet | DTS LV3 | FULLY | 4.640 | 0.277 | 4.917 |
| ImageNet | RAW | PARTIAL | 6.313 | 0.723 | 7.036 |
| ImageNet | DTS LV0 | PARTIAL | 4.399 | 0.307 | 4.706 |
| ImageNet | DTS LV1 | PARTIAL | 2.099 | 0.182 | 2.281 |
| ImageNet | DTS LV2 | PARTIAL | 2.195 | 0.195 | 2.390 |
| ImageNet | DTS LV3 | PARTIAL | 1.645 | 0.161 | 1.806 |
预热后,RAW 文件夹因为有 OS 缓存加持,差距会缩小,但 DTS 仍然有稳定优势。特别是在 1000 个 epoch 的总时间对比中,ImageNet PARTIAL 模式下 RAW 需要 4377.384 秒,DTS LV0 降到 3650.975 秒,DTS LV3 进一步降到 1548.642 秒;FULLY 模式下 RAW 340.045 秒,DTS LV3 333.137 秒——当数据集能全量进内存时,压缩带来的收益会被内存访问拉平,但 DTS 依然不落后。
更有趣的是笔记本电脑场景(Intel Core i9-14900HX,8 GB 内存,Windows):
| TRAIN | VAL | TRAIN+VAL | |
|---|---|---|---|
| RAW | 1610.463 | 63.005 | 1673.468 |
| DTS LV0 | 37.658 | 1.424 | 39.082 |
| Speed Up | 42.765× | 44.245× | 42.819× |
在笔记本电脑这类 IO 和内存资源都更受限的环境下,RAW 文件夹的随机小文件读取被极度放大,而 DTS LV0 凭借顺序大 IO 直接把训练集加载时间从 1610.463 秒压缩到 37.658 秒,超过 42 倍。这个数字不是特例,而是小文件随机 IO 瓶颈被消除后的典型表现。
六、FULLY 与 PARTIAL:两种内存策略
DTS 加载器支持两种模式:
// include/renaissance/core/global_config.h:51-55
enum class LoadMode {
AUTO, ///< 自动选择(根据内存判断,当前版本保留)
FULLY, ///< 全量加载:整个数据集一次性加载到内存
PARTIAL ///< 部分加载:使用环形缓冲区循环加载
};
6.1 FULLY 模式
内存足够时,把整个数据集一次性读入内存。后续 epoch 不再触碰磁盘,只在内存里做全局 shuffle。这是吞吐最高的模式,因为训练阶段完全没有 IO。
// src/data/imagenet_loader_dts.cpp:609-753
if (current_set_->mode == LoadMode::FULLY) {
allocate_buffers(*current_set_);
build_full_shuffled_locations(*current_set_);
load_one_buffer_batch_fully(*current_set_, 0);
}
第一个 epoch 会流式加载所有 Block 到 full_arena,同时每个 worker 收集自己读过的 SampleInfo。第一个 epoch 结束后,这些数据被汇总成全局样本数组。第二个 epoch 起直接对全局数组做 shuffle,然后按 worker 分段,实现纯内存读取。
6.2 PARTIAL 模式
内存不足时,只保留两个 buffer,每个 buffer 容纳 PF × N 个 Block。N 是 IO worker 数,PF 是预取因子(默认 4)。一个 buffer 在被 preprocessor 消费的同时,另一个 buffer 在后台加载,形成双缓冲。
// src/data/imagenet_loader_dts.cpp:754-803 load_one_buffer_batch(actual_buffer_A, *current_set_, 0); current_set_->ready_buffer = actual_buffer_A;
如果训练集和验证集都使用 PARTIAL 模式,框架还会启用共享缓冲区,让 train 和 val 复用同一对 buffer,减少内存占用:
// src/data/imagenet_loader_dts.cpp:516-528
if (train_set_.mode == LoadMode::PARTIAL &&
val_set_.mode == LoadMode::PARTIAL &&
!train_set_.file_path.empty() && !val_set_.file_path.empty()) {
use_shared_buffers_ = true;
}
6.3 MNIST/CIFAR 强制 FULLY
MNIST 和 CIFAR 数据集很小(MNIST 训练集约 47 MB,CIFAR-10 训练集约 146 MB),直接全量加载收益最大,代码也更简单:
// src/data/mnist_loader_dts.cpp:151-158 train_set_.mode = LoadMode::FULLY; val_set_.mode = LoadMode::FULLY;
七、Native IO 与零竞争线程模型
DTS 的 C++ 加载器不使用 C 标准库的 fread,而是直接调用平台 Native API:
- Windows:
ReadFile+SetFilePointerEx - Linux:
pread
// src/data/imagenet_loader_dts.cpp:1263-1328
void ImageNetLoaderDts::read_block_native(FileHandle& file, uint32_t block_id, uint8_t* dst) {
constexpr size_t CHUNK_SIZE = 4 * 1024 * 1024; // 4MB chunks
size_t file_offset = FILE_HEADER_SIZE + static_cast<size_t>(block_id) * BLOCK_SIZE;
// Windows: ReadFile
// Linux: pread
}
文件句柄的跨平台封装由 FileHandle 类完成,它使用 RAII 管理资源生命周期——打开时自动获取句柄,析构时自动释放。在 Windows 上使用 CreateFileA/ReadFile/CloseHandle,在 Linux 上使用 open/pread/close,对外暴露统一的接口。
每个 IO worker 独立打开一个文件句柄,按 stride=N 的静态方式领取 Block,写到自己专属的 slot:
// src/data/imagenet_loader_dts.cpp:1232-1258
for (int block_seq = start_block + thread_id; block_seq < end_block; block_seq += N) {
uint32_t block_id = ds.epoch_block_order[block_seq];
uint8_t* dst = buffer.data + static_cast<size_t>(slot_idx) * block_size_;
read_block_native(file, block_id, dst);
parse_block_meta(slot_idx, dst, ds, buffer.slot_metas[slot_idx]);
slot_idx++;
}
这种“静态领取”策略是关键:每个线程负责哪些 Block 在加载前就已经确定,不需要任务队列、不需要锁、不需要原子计数器。主线程只需要在所有 IO 线程上做一次 join,就能获得完整的内存屏障。这也为后续 Philox 可复现随机数奠定了基础。
八、CRC-32:大文件传输的守门员
一个 147 GB 的训练集文件,在服务器之间拷贝、解压、上传下载,出错概率并不低。DTS 在文件头里保存了数据区的 CRC-32 校验码:
# python/scripts/make_dataset.py:357-362
def get_partial_crc32(file_path, start_byte_index, return_str=True):
"""从文件的第 N 个字节开始校验直到文件结束"""
return _compute_crc32(file_path, offset=start_byte_index, return_str=return_str)
加载器可以选择在 configure() 时启用校验:
// src/data/imagenet_loader_dts.cpp:427-437
if (verify_crc_) {
LOG_INFO << "Performing CRC-32 verification...";
if (!train_path.empty()) {
if (!verify_dts_crc(train_path)) { ... }
}
}
verify_dts_crc() 跳过文件头,用 zlib 的 crc32() 逐 MB 计算剩余数据:
// src/data/imagenet_loader_dts.cpp:2180-2254
uLong computed_crc = crc32(0L, Z_NULL, 0);
while (file) {
file.read(reinterpret_cast<char*>(buffer.data()), BUFFER_SIZE);
computed_crc = crc32(computed_crc, buffer.data(), bytes_read);
}
虽然首次校验需要读一遍整个文件,但对于大团队共享数据集、云端下载等场景,这笔开销换来的心安是值得的。
九、DTS 与后续管线的衔接
DTS 加载器不是孤立存在的。它通过统一接口把样本交给 Preprocessor:
// include/renaissance/data/data_loader.h:102-106
virtual bool get_next_sample(
int preproc_worker_id,
int32_t& label,
const uint8_t*& data_ptr,
size_t& data_size) = 0;
Preprocessor 拿到 data_ptr 后,用 TurboJPEG 解码,然后执行 PO 链(RandomResizedCrop、ColorJitter、FusedNormalization 等),最后写入 TransferStation 的双缓冲。下一篇讲多线程预处理时会详细展开;这里只需要知道,DTS 的零拷贝返回让 CPU 预处理不必等待一次额外的内存拷贝。
// src/data/preprocess_worker.cpp:476-490
uint8_t* PreprocessWorker::request_transfer_station_slot(int32_t label) {
auto [batch_id, position] = calculate_write_position();
uint8_t* write_ptr = transfer_station_->request_write_slot(position, batch_id, label);
return write_ptr;
}
用户代码里,DTS 的启用也很简单,通过 PREPROCESSOR_SETTING 链式配置即可:
// tests/correction/test_h2d_only_epoch.cpp:534-545
PREPROCESSOR_SETTING
.dataset(dataset_type_str, dataset_path)
.using_dts_format(use_dts, compression_level)
.fully_mode(!partial_mode)
...
.commit();
十、与主流方案的一点对比
DTS 不是要做通用数据库,也不是要替代 FFCV。它做的是:在 Tech-Renaissance 的静态图、NHWC、CUDA Graph、多流训练这条技术栈里,让数据格式与后续每一个环节对齐。
相比 PyTorch 默认的 ImageFolder,DTS 把训练时的小文件随机 IO 变成了顺序大 IO;相比 LMDB,DTS 不依赖 mmap,Windows 和 Linux 都能原生跑;相比 TFRecord,DTS 没有 protobuf 反序列化开销;相比 WebDataset,DTS 的 shuffle 和双缓冲是为训练循环静态图专门设计的;相比 FFCV,DTS 在思路上学习了 .beton 的分块压缩策略,但文件结构和加载逻辑完全围绕本框架的管线定制。
这种专用化带来了取舍:DTS 文件只能用 Tech-Renaissance 读取,不像 TFRecord 那样通用。但对于一个追求极致训练吞吐的框架来说,专用格式是性能的最后几公里。
十一、小结
DTS 数据格式是 Tech-Renaissance 数据引擎的底座。它通过四级压缩、固定 Block、SlotMeta 索引、三级随机 shuffle、FULLY/PARTIAL 双模式、Native IO、CRC 校验这些设计,把 ImageNet 这种百万级小文件数据集的训练读取,变成了一件可控、可复现、高吞吐的事情。
核心要点再强调一遍:
- 磁盘布局决定训练 IO 上限:16 MB Block + 顺序读取,让小文件随机 IO 不再是瓶颈。
- LV0 完全等价于原版 ImageNet:LV0 不重新编码图片,只是把原始 JPEG 字节流拼接对齐,图像内容与原版数据集完全一致。
- 导出阶段一次性优化:LV1~LV3 的缩放、裁剪、JPEG 压缩都在制作 DTS 时完成,训练时只做必要的 RandomResizedCrop 和 Normalize。
- 零拷贝 + 静态领取:
get_next_sample直接返回内部 arena 指针,多 IO worker 无竞争。 - FULLY/PARTIAL 两种内存策略:内存充足则全量加载实现零 IO,内存受限则双缓冲循环加载。
- 可复现 shuffle + CRC 校验:训练结果可复现,数据完整性可验证。
数据格式这层地基打好了,下一篇我们就可以放心地讲多线程预处理、NUMA 感知和异步双缓冲——看看 CPU 侧怎么用 200+ 个线程把 GPU 喂到饱。
