先聊点实际的。做图像分类、检索或者自监督训练的人，迟早会遇到同一个瓶颈：数据加载吃不满GPU。尤其是ImageNet这种上百万张的图片，几十万个.jpg散落在三四层目录里，每个训练step都要把batch从磁盘搬到内存再送进显存。最开始大家都会按路径直接读，也就是写一个 FileListDataset，拿一个list记录每张图的路径和标签，然后训练时一个个open、read、decode。这套方案在小数据上没问题，到了ImageNet级别，GPU利用率掉到三四十是常有的事。我自己调过很多轮之后，最终稳定下来是用LMDB把ImageNet整个打包起来，配合PyTorch自定义Dataset来读。这篇文章就是把这套做法完完整整拆开讲一遍：为什么选LMDB、转换脚本怎么写、Dataset里有哪些坑、多进程要注意什么、性能怎么调。适合正在做大规模图像训练、或者手上有一堆散乱图片想加速读取的开发者，尽量做到可以直接照着落地。
1. 为什么是LMDB而不是直接读文件
1.1 海量小文件的IO根因
很多人以为训练慢是GPU不行，实际上在ImageNet量级的任务上，数据管线经常才是短板。直接读jpg文件时，每次open都要经过文件系统查找、目录项解析、打开文件描述符、读取inode，然后才真正开始读数据。ImageNet平均每张图100KB到200KB，但系统开销跟文件大小无关，它主要花在文件打开和寻址上。你从机械硬盘或普通网络存储上随机读一百万个小文件，整个过程基本就是磁盘寻道在主导，毫无顺序可言。
即使换成SSD，高并发下文件描述符的开销、内核页缓存的碎片化，也会让实际吞吐远低于磁盘标称顺序读速度。PyTorch的DataLoader虽然可以用多进程并行加载，但每个worker还是各自去文件系统里读文件，到几百个worker的时候，系统调用开销已经大得离谱。这类“海量小文件随机读”的IO模式，恰恰是文件系统最不擅长、而键值数据库最擅长的场景。
1.2 LMDB为什么快：内存映射与B+树
LMDB，全称Lightning Memory-Mapped Database，是一个嵌入式键值数据库。它最大的特点是基于内存映射（mmap），打开数据库文件后，整个库的索引和内容都映射到进程的虚拟地址空间。读取数据时不再需要read系统调用，而是直接访问内存地址，由操作系统页缓存来管理热点页。这意味着同一个lmdb文件被反复读取时，第二次开始很大概率直接命中页缓存，磁盘IO被大幅摊薄。
索引结构上，LMDB使用B+树。这个结构保持了键的有序性，随机查找复杂度是O(log N)级别，对百万级键值来说非常稳定。更关键的是，B+树的内部节点会被缓存在页缓存里，当访问模式有一定的局部性时，效果会更好。对比前面说的直接读jpg，LMDB把“打开文件、读取、关闭”的损耗压缩成了一个内存查找加一次memcpy，速度上完全是两个量级。
1.3 为什么不选TFRecord或者WebDataset
提到大规模数据打包，很多人也会想到TFRecord和WebDataset。TFRecord用流式读法，适合顺序访问，但它需要遵循protobuf定义，改格式还要改协议，调试起来麻烦一点；WebDataset把数据打包成tar文件，最大的优点是可以直接用流式数据管道解压。这俩方案都有各自的使用场景，但在PyTorch单机训练里，LMDB有个非常自然契合的优势：它天然支持按键随机访问，刚好匹配Dataset按索引取样本的范式。而且LMDB的事务性写入让转换过程可断点续写，不容易由于意外中断就整个重来。
2. 整套数据管线的设计
2.1 数据流向与模块划分
一个完整的LMDB加速方案，可以拆成这样几条链路：
原始图片目录 -> 转换脚本 -> LMDB文件 -> 自定义Dataset -> DataLoader多进程 -> 数据增强 -> 模型训练
这个划分里，转换脚本是一次性离线工作，之后每个epoch读的都是同一个lmdb文件。Dataset负责按索引取出图片字节和标签，数据增强（RandomResizedCrop、RandomHorizontalFlip之类的操作）放在取到数据之后执行。模块之间界限清楚，哪里出问题也很容易定位。
2.2 key和value的设计
LMDB里键和值都只是bytes，所以怎么设计编码完全取决于你。我见过有人直接用图片路径字符串作为key，这样虽然可读性好，但B+树索引的空间利用率会变差，字符串匹配的开销也比整数大。更推荐的做法是使用从0开始递增的整数索引作为key，编码成bytes。训练时Dataset收到的index就是整数，直接编码成key，一次get就拿到对应数据。
value方面，最通用的是用pickle序列化一个二元组，里面包含图片原始字节和标签。比如pickle.dumps((image_bytes, label))。搞不懂为什么把label放进去？因为这样读取时一次反序列化就能拿回完整样本，不用额外查meta表。还有一种做法是自定义二进制布局，把图片字节长度、标签、图片数据依次排列，读取时手动解析。这个优化后面会细说，如果不想折腾，pickle已经足够稳定。
2.3 样本数量单独存meta键
有一个经常被忽略的设计细节：在转换脚本里，写完所有样本后，额外存一个专门的键来记录样本总数。比如约定b'__len__'对应一个十进制字符串表示的总数。原因是PyTorch的DataLoader在初始化时要调用__len__方法，如果这个值每次都要遍历整个数据库才能得到，启动开销会非常大。有了meta键，Dataset只需要打开环境、读一次这个键、关闭环境，整个过程也就几毫秒。
2.4 图片解码放哪一层
有的方案会在转换阶段把图片直接解码成numpy数组，再存进LMDB。我不推荐这么做，因为解码后的RGB图像体积比JPEG大好几倍，存储成本暴涨，而且解码发生在训练时也没有问题——worker在CPU上做解码，GPU在那边等着吃数据，两者天然并行。所以主流做法就是保留原始JPEG字节，数据管线解码，既能省空间，又保持了灵活性。
3. 实操：转换脚本怎么写出高效
3.1 环境准备
转换阶段需要安装两个基础库。
pip install lmdb pillowLMDB库在PyTorch生态里很常用，没有版本兼容的困扰。Pillow用于图片读取和解码，但转换脚本里其实不需要解码图片，可以直接读原始字节，所以用内置的open就够。唯一要注意的是确保脚本运行在64位系统上，这个在map_size部分会详细解释。
3.2 map_size的计算
map_size是LMDB里最坑、最容易出错的参数。很多第一次用的人以为它是数据文件的大小，随便设个几GB，结果转换到一半就报错。实际上map_size是内存映射的虚拟地址空间大小，它决定了数据库能容纳的数据总量上限。
计算方法并不复杂：LMDB写入数据时会有页对齐开销，但经验上取所有JPEG文件总字节数的1.3到1.5倍作为map_size比较稳妥。如果数据总量不好精确统计，可以先跑一遍脚本统计总字节数，再乘以这个系数。一个可能的命令行实现：
du -sb /data/ImageNet/train比如统计出来总共180GB，那么map_size至少设置到250GB左右。注意这里用的是字节单位，代码里直接乘1024的幂。
MAP_SIZE = 250 * 1024 * 1024 * 1024为什么强调64位系统？因为单个进程的虚拟地址空间在32位系统下只有4GB左右，你不可能映射一个250GB的LMDB文件。所以如果你的机器还是32位，赶紧换环境，这条路走不通。
3.3 批量事务写入
LMDB写入时，每次都调用commit会产生很大的开销，因为commit意味着一次事务同步、可能的日志写入。转换ImageNet这种百万级数据，如果一张图提交一次，时间会拖到一个完全不可接受的程度。正确做法是攒够一批再提交，比如每5000条数据commit一次。
另外还有一个影响转换性能的关键点：先构建样本列表，再开始写入。如果边遍历目录边写库，文件系统的目录扫描和LMDB写入会频繁争抢磁盘IO资源，导致整体速度下降。先把所有图片路径和标签收集到一个Python列表里，然后按顺序读取并写入LMDB，这样更贴近顺序读写的模式。
一个经过实践检验的转换脚本长这样：
import lmdb import os import pickle import glob def build_sample_list(root): samples = [] classes = sorted(os.listdir(root)) for label, cls in enumerate(classes): cls_dir = os.path.join(root, cls) paths = sorted(glob.glob(os.path.join(cls_dir, '*.jpeg'))) for p in paths: samples.append((p, label)) return samples def convert_to_lmdb(sample_list, db_path, map_size): env = lmdb.open(db_path, map_size=map_size) txn = env.begin(write=True) for idx, (path, label) in enumerate(sample_list): with open(path, 'rb') as f: data = f.read() txn.put(str(idx).encode(), pickle.dumps((data, label))) if (idx + 1) % 5000 == 0: txn.commit() txn = env.begin(write=True) print(f'processed {idx + 1}/{len(sample_list)}') txn.put(b'__len__', str(len(sample_list)).encode()) txn.commit() env.sync() env.close() if __name__ == '__main__': sample_list = build_sample_list('/data/ImageNet/train') convert_to_lmdb(sample_list, '/data/ImageNet/train.lmdb', 250 * 1024**3)这里有几个细节值得解释。sorted的作用是让同一类的图片在文件系统里尽量靠在一起，这样读取原图时HDD的顺序性会好很多。你可能会问为什么不加多线程写入，因为LMDB本身是单写者模型，多线程写同一个环境反而不稳定，SSD上单线程写LMDB能跑到接近磁盘顺序写的上限。
3.4 转换耗时与校验
转换百万量级图片整个过程会持续不短时间。SSD上大概一到两个小时，机械硬盘上会更久。所以建议在正式转换之前，先用一个几百张图的小目录跑通整个流程，确认没有逻辑错误，再去转换完整数据集。
转换完成后，打开LMDB随机抽查几个键值，确认能取出来正确的图片和标签。这一步看起来多余，但能省下后面训练跑了一半才发现数据坏了的时间。
4. 自定义Dataset实现与多进程细节
4.1 初始化时机是最大的坑
写LMDB Dataset时最容易犯的错误，就是在__init__方法里直接打开LMDB环境。这个坑我见过太多次了。原因在于PyTorch的DataLoader设置num_workers>0时，每个worker是通过fork来创建的。如果在__init__里已经打开了环境，那么fork出来的所有子进程会继承同一个文件描述符和内存映射区。多进程同时共享这同一个环境对象，轻则提示各种lock冲突报错，重则直接崩溃或者读到错乱的数据。
正确做法是：在__init__里只保存路径，初始化环境放在第一次__getitem__调用时。这样每个worker进程在数据处理时都会自己打开一份环境，互不干扰，也没有继承的问题。
import lmdb import io import pickle from PIL import Image from torch.utils.data import Dataset class ImageNetLMDBDataset(Dataset): def __init__(self, db_path, transform=None): self.db_path = db_path self.transform = transform self.env = None self.txn = None self.length = self._get_length() def _get_length(self): env = lmdb.open(self.db_path, readonly=True, lock=False) with env.begin() as txn: length = int(txn.get(b'__len__').decode()) env.close() return length def _init_db(self): self.env = lmdb.open( self.db_path, readonly=True, lock=False, readahead=False, meminit=False, ) self.txn = self.env.begin(write=False) def _load_sample(self, index): key = str(index).encode() value = self.txn.get(key) img_bytes, label = pickle.loads(value) return img_bytes, label def __getitem__(self, index): if self.env is None: self._init_db() img_bytes, label = self._load_sample(index) img = Image.open(io.BytesIO(img_bytes)).convert('RGB') if self.transform is not None: img = self.transform(img) return img, label def __len__(self): return self.length这个类的设计有几个关键参数：readonly=True保证训练时不会对数据库做任何写操作，lock=False让多个worker可以无锁并发访问同一个文件，readahead=False禁止LMDB一次性预读大量连续页，meminit=False跳过新页面初始化时的清零动作。这几个组合起来，在多进程读场景下性能是最稳的。
4.2 图片解码的细节
读取时需要通过Image.open(io.BytesIO(img_bytes)).convert('RGB')来做解码。.convert('RGB')这一步别省。因为有的JPEG保存的是灰度图，有的带透明度通道，直接送入后续transform会造成通道数不一致或者报错。显式转换到RGB能保证下游输入统一。
解码本身是在CPU上执行的，每个worker并行独立解码，所以整体吞吐量可以随着worker数量提升。由于lmdb读取走内存映射，少了文件IO的系统调用，整个value->解码->tensor链路的开销大头就落在JPEG解码上。如果头部图细碎导致解码很慢，可以考虑用OpenCV的cv2.imdecode替换Pillow，但要注意OpenCV解码出来是BGR顺序，后面需要显式转成RGB。
4.3 DataLoader参数配合
Dataset写好了还不够，DataLoader参数如果不搭配合适，同样发挥不出LMDB的优势。我常用的参数组合是这样：
dataloader = DataLoader( dataset, batch_size=128, shuffle=True, num_workers=8, pin_memory=True, prefetch_factor=4, persistent_workers=True, )prefetch_factor控制每个worker预取几个batch，适当调大可以让数据准备和训练迭代重叠得更充分。persistent_workers=True避免每个epoch重新创建worker，省掉fork的开销。这两个参数配合LMDB低延迟读取，训练循环的数据等待时间会明显下降。
4.4 增强后的数据流
当整个链路跑通之后，数据流大概是这样的：每个worker维护自己的LMDB环境，拿到index后读取键值、反序列化、解码图片，再执行transform。由于每个worker独立操作自己的环境，没有锁竞争，而且页缓存可以跨进程共享，所以多worker并发时吞吐量接近线性增长。GPU那边只要等着拿batch，几乎不会再因为数据问题空转。
5. 性能对比与调优实录
5.1 评测方法
要评估数据管线快不快，不要直接跑完整训练，那样难以分离瓶颈。正确做法是写一个极简评测脚本：只迭代DataLoader，按批次把数据从CPU搬到GPU，记录每秒处理的图片数。
import time from torch.utils.data import DataLoader dataset = ImageNetLMDBDataset('/data/ImageNet/train.lmdb', transform=None) loader = DataLoader(dataset, batch_size=256, num_workers=8) start = time.time() for i, (imgs, labels) in enumerate(loader): if i >= 50: break end = time.time() print(f'images/s: {50 * 256 / (end - start)}')跑这个脚本时只测读取性能，不训练模型。这样能忽略模型的干扰，专门看数据管线能提供多大吞吐。
5.2 FileList、LMDB加pickle、LMDB加自定义二进制
我分别在同样的机器上（普通双路服务器、SATA SSD、8核CPU）跑过三种方案，结论可以用一张表说明。
| 读取方案 | 主要开销 | 8线程下实测吞吐（经验值） | 备注 |
|---|---|---|---|
| 直接按路径读jpg | open/read/close系统调用 | 约500-700 img/s | HDD上更低，目录层级深时更差 |
| LMDB + pickle | 页缓存命中+反序列化 | 约1000-1300 img/s | 稳定，容易上手 |
| LMDB + 自定义字节布局 | 页缓存命中+内存拷贝 | 约1200-1500 img/s | 解码仍然是最大瓶颈 |
这个表格是经验值，不同机器上数字会有出入，但比例关系可以参考。LMDB的提速不只是因为少了文件打开，更重要的是页缓存命中率远高于直接读目录。
5.3 去掉pickle的自定义二进制
如果你的数据模块在profile时发现pickle.loads花费占比不低，可以改成自定义二进制布局。具体做法是转换时这样编码：图片字节长度占4字节、标签占4字节，后面接图片原始数据。读取时先用struct.unpack取头部信息，再用索引切片取出图片字节。
import struct def pack_sample(img_bytes, label): return struct.pack('II', len(img_bytes), label) + img_bytes def unpack_sample(data): img_len, label = struct.unpack('II', data[:8]) return data[8:8 + img_len], label这个方案能省掉pickle构造tuple的额外开销，代码也更直接。不过提升幅度在百分之十几左右，如果当前吞吐已经足够喂饱GPU，不折腾也行。
5.4 页缓存带来的“越读越快”现象
用LMDB还会观察到一个现象：第一个epoch偏慢，之后会变快。原因很简单，首次访问时lmdb文件还没有全部进入页缓存，读的过程本质上是在磁盘IO；但多次读取之后，操作系统的页缓存已经保留了热点数据，后续访问基本都命中内存。如果你的数据量特别大，机器内存装不下整个lmdb文件，这个效果会打折扣。
5.5 解码仍是最终瓶颈
把读取做到极致之后，最终卡住你的往往是JPEG解码。一旦DataLoader的所有worker都忙于解码，加大num_workers也未必有提升，因为CPU核心数有限。这时候可以考虑三种途径：换用更快的解码库、在转换阶段将图片统一缩放或转成另一种格式（比如WebP），或者让数据增强部分也放到GPU上执行。最省事的还是增大worker数量，让训练进程和Worker进程保持在不同的核心上运行。
6. 常见问题与排查技巧
6.1 问题速查表
| 现象 | 原因 | 解决方案 |
|---|---|---|
| 转换时报MDB_MAP_FULL | map_size设置过小 | 重新计算总字节数，留1.5倍余量重建 |
| 多进程一起读时进程卡住或报锁冲突 | fork继承了相同环境指针 | 把env初始化移至第一次__getitem__，并设置lock=False |
| 训练时内存占用异常高 | readahead=True导致预读大量数据页 | 设置readahead=False |
| 读取图片时报无法解码 | 转换时读入了损坏文件 | 校验图片完整性，跳过损坏文件 |
| 复制lmdb文件到另一台机器打不开 | 复制前未正常close环境 | env.close()后再复制，并同时复制到SSD上运行 |
| 所有worker都在解码但吞吐上不去 | CPU核心已是瓶颈 | 加大worker数超过CPU核数没有意义，考虑换解码器 |
6.2 MDB_MAP_FULL的完整现场
我当年第一次转换的时候，随手设了个map_size=10 * 1024**3，结果转换到一半直接抛异常。后来养成了习惯：写转换脚本前先统计文件夹大小，再给1.5倍余量。另一个注意点是如果在Windows上跑，需要使用Linux的WSL或Docker环境，LMDB在原生Windows下的兼容性有各种坑，没必要在这种环境上浪费时间。
6.3 多进程读取时的锁冲突
LMDB对环境加锁有自己的机制。当你以默认方式打开环境并且多个进程同时读取，可能会看到类似锁文件异常或者进程阻塞。解决方法就是我前面Dataset代码里的那三件套：readonly=True, lock=False, readahead=False。每加一个worker就多一点概率踩到锁冲突，能尽早规避就别等到跑训练时才去排查。
6.4 训练前预热
如果数据量大到第一个epoch慢到影响体验，可以考虑在正式训练前加一个预热循环。不调用反向传播，只主动遍历一遍DataLoader。这一步能做两件事：一是让页缓存预热，二是确保数据管线在训练前已经稳定，而不是训练跑到一半才暴露问题。
6.5 多卡时的特别提醒
多卡训练时每张卡会有几个worker进程，整体并发会更高。类设计不需要改变，因为不同GPU上的DataLoader每个worker都独立打开环境，文件级别的共享不会带来一致性风险。唯一要考虑的是机器内存：多个worker的页缓存大小与数据文件大小之间需要做好平衡。如果内存实在不够，可以把一个大数据集拆成几个小shard，比如train-part0.lmdb、train-part1.lmdb，每个epoch轮流读取其中一部分，或者用多个Dataset拼接来解决。
写到这里，结合我反复用过这个方案的经验，再留几个小看法。第一，在转换脚本里把__len__的meta键保存下来是绝对值得的，否则每次初始化Dataset都要遍历几十万条数据，慢得让人怀疑人生。第二，readahead=False一定不要省掉，很多看起来“LMDB内存占用异常高”的问题就是这么来的。第三，如果你的数据规模只有几万张甚至更少，FileList方案空转一小会儿就能读完，完全没必要引入LMDB；这个方案真正的价值体现在百万级图片、每次采样随机性强、页缓存命中能带来持续收益的场景。另外，后期如果觉得读取还是卡，检查一下num_workers和解码库的开销，往往比继续折腾存储层更有效果。
