如何使用Python和PyTorch实现自定义的Dataset类加速数据读取？-Python教程-PHP中文网

直接继承 Dataset 不提速，因它仅定义接口；瓶颈在 getitem 中重复 I/O 和解码，需预加载、内存映射、合理配置 num_workers 和 pin_memory 才能加速。

如何使用python和pytorch实现自定义的dataset类加速数据读取？

为什么直接继承 `torch.utils.data.Dataset` 有时反而变慢？

不是继承了就自动加速——Dataset 本身不提供任何加载优化，它只定义了 __getitem__ 和 __len__ 接口。如果在 __getitem__ 里每次打开文件、解码图像、做 transform，I/O 和 CPU 计算会成为瓶颈，DataLoader 的多进程根本喂不饱 GPU。

真正提速的关键在于：把耗时操作提前固化（如预解码、内存映射）、避免重复计算、配合 DataLoader 的 num_workers 和 pin_memory 才能生效。

常见错误：在 __getitem__ 中用 PIL.Image.open(path).convert('RGB') —— 每次都磁盘读+解码
更优做法：提前把图像 decode 成 np.ndarray 或 torch.Tensor 缓存到内存/内存映射文件中
注意：若数据集太大放不下内存，改用 mmap 或 HDF5 + h5py.File(..., driver='core') 控制缓存粒度

`getitem` 里哪些操作必须同步、哪些可以异步？

DataLoader 的 worker 进程是独立的 Python 进程，不能共享 OpenCV 的全局状态、PyTorch 的 CUDA 上下文，也不能直接传递未序列化的对象（比如带 file handle 的 PIL Image）。

必须同步（只能在主线程做）：torch.cuda.set_device()、torch.backends.cudnn.benchmark = True
必须在每个 worker 内部重置（否则可能 crash）：cv2.setNumThreads(0)（OpenCV 多线程和 DataLoader 冲突）、torch.set_num_threads(1)
可安全异步（推荐放 __getitem__）：Numpy 数值变换、torchvision.transforms（纯 tensor 操作）、label 编码
禁止异步：open() 文件句柄跨进程传递、cv2.VideoCapture 实例、未 pickleable 的自定义类实例

如何让自定义 Dataset 支持 `collate_fn` 而不报错？

默认 default_collate 假设所有样本返回相同结构（比如都是 (tensor, int)），一旦你返回 dict、tuple 混合或变长序列（如不同尺寸图像 pad 后堆叠），就会报 TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists。

macOS UI 自动化技能包 —— 用 Python 控制你的 Mac

让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac，无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启：辅助功能 + 屏幕录制权限

下载

最简修复：在 DataLoader 中显式指定 collate_fn，例如处理图像+mask+meta 的 dict：

def collate_fn(batch):
    images = torch.stack([b["image"] for b in batch])
    masks = torch.stack([b["mask"] for b in batch])
    metas = [b["meta"] for b in batch]  # list of dict, not stacked
    return {"image": images, "mask": masks, "meta": metas}

注意：不要在 collate_fn 里做 heavy transform（如 resize），它运行在主线程，会拖慢整个 pipeline

如果 batch 内 tensor shape 不一致（如检测任务的 variable-size boxes），用 torch.nn.utils.rnn.pad_sequence 或手动 pad，别依赖 stack

使用 `torchvision.io.read_image` 替代 PIL 的实际效果

在 Linux + SSD 场景下，torchvision.io.read_image(path, mode=torchvision.io.ImageReadMode.RGB) 比 PIL.Image.open().convert() 快 2–3 倍，且返回 torch.Tensor（CHW, uint8），省去 to_tensor() 转换，也规避 PIL 的 GIL 锁。

前提：安装带 libjpeg-turbo 支持的 torchvision（pip install torchvision --no-cache-dir 可能装错版本）
验证是否生效：torchvision.io._has_read_img_support() 返回 True
坑点：read_image 不支持 GIF/WEBP，遇到就 fallback 到 PIL；路径必须是 str（不能是 Path 对象）
进阶：配合 torchvision.datasets.ImageFolder 的 loader 参数替换默认 loader，无需重写整个 Dataset

实际部署时最容易被忽略的是 worker 初始化逻辑——DataLoader 的每个子进程第一次调用你的 Dataset 时，都要重新执行一遍 import 和全局配置，OpenCV 线程数、随机 seed、甚至 logger 都得在 getitem 外围或 worker_init_fn 里重置。