TorchData

TorchData 是什么？ | 有状态的 DataLoader | 安装指南 | 贡献 | 许可证

TorchData 是什么？

TorchData 项目是对 PyTorch torch.utils.data.DataLoader 和 torch.utils.data.Dataset/IterableDataset 的迭代增强，旨在使其成为可扩展、高性能的数据加载解决方案。我们将要在 torchdata 仓库中对这些增强功能进行迭代。

我们的第一个改动始于为 torch.utils.data.DataLoader 添加检查点功能，该功能位于 stateful_dataloader 中，它是 torch.utils.data.DataLoader 的直接替代品。通过定义 load_state_dict 和 state_dict 方法，它实现了 mid-epoch checkpointing（Epoch 中途检查点），并提供 API 供用户跟踪自定义迭代进度，以及来自 dataloader workers（数据加载器工作进程）的其他自定义状态，例如 token buffers（令牌缓冲区）和/或 RNG states（随机数生成器状态）。

Stateful DataLoader

torchdata.stateful_dataloader.StatefulDataLoader 是 torch.utils.data.DataLoader 的直接替代品，它提供了 state_dict 和 load_state_dict 功能。请参阅 Stateful DataLoader 主页 获取更多信息和示例。也可以查看 此 Colab notebook 中的示例。

torchdata.nodes

torchdata.nodes 是一个由可组合 iterators（迭代器，注意不是 iterables/可迭代对象！）组成的库，允许您将常见的数据加载和预处理操作链接在一起。它遵循流式编程模型，不过如果您需要，仍然可以配置 "sampler + Map-style"（采样器 + Map 风格）模式。有关更多详细信息，请参阅 torchdata.nodes 主页。敬请期待即将推出的 torchdata.nodes 教程！

安装

版本兼容性

以下是相应的 torchdata 版本和支持的 Python 版本。

torch	torchdata	python
master / nightly	main / nightly	>=3.9, <=3.13
2.6.0	0.11.0	>=3.9, <=3.13
2.5.0	0.10.0	>=3.9, <=3.12
2.5.0	0.9.0	>=3.9, <=3.12
2.4.0	0.8.0	>=3.8, <=3.12
2.0.0	0.6.0	>=3.8, <=3.11
1.13.1	0.5.1	>=3.7, <=3.10
1.12.1	0.4.1	>=3.7, <=3.10
1.12.0	0.4.0	>=3.7, <=3.10
1.11.0	0.3.0	>=3.7, <=3.10

本地 pip 或 conda

首先，设置一个环境。我们将安装 PyTorch 二进制文件以及 torchdata。如果您使用 conda，请创建一个 conda 环境：

conda create --name torchdata
conda activate torchdata

如果您希望改用 venv：

python -m venv torchdata-env
source torchdata-env/bin/activate

安装 torchdata：

使用 pip：

pip install torchdata

使用 conda：

conda install -c pytorch torchdata

从源码安装

pip install .

如果从源码构建 TorchData 失败，请按照 贡献页面 上的链接指南安装 PyTorch 的 nightly（每夜构建）版本。

从 nightly 版本安装

TorchData 的 nightly 版本也已提供，并每天从主分支更新。

使用 pip：

pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu

使用 conda：

conda install torchdata -c pytorch-nightly

贡献

我们欢迎 PR（Pull Request）！请参阅 CONTRIBUTING 文件。

Beta 版本使用与反馈

我们非常乐意听取早期采用者的意见并与之合作，以完善我们的设计。如果您有兴趣在项目中使用此工具，请通过提交 issue 与我们联系。

许可证

TorchData 采用 BSD 许可证，详见 LICENSE 文件。

TorchData 快速上手指南

TorchData 是 PyTorch DataLoader 和 Dataset 的迭代增强库，旨在提供可扩展、高性能的数据加载解决方案，支持状态保存与恢复。

环境准备

在安装 TorchData 前，请确保您的环境满足以下依赖要求：

• Python：建议 3.9 及以上版本（不同 TorchData 版本对 Python 支持不同，最新版支持至 3.13）。
• PyTorch：必须预先安装，且版本需与 TorchData 严格匹配。

版本兼容参考（部分）：

torch 版本	torchdata 版本
2.6.0	0.11.0
2.5.0	0.10.0 / 0.9.0
2.4.0	0.8.0

安装步骤

推荐在虚拟环境中进行安装。

1. 创建并激活虚拟环境

使用 Conda（推荐）：

conda create --name torchdata
conda activate torchdata

或使用 venv：

python -m venv torchdata-env
source torchdata-env/bin/activate

2. 安装 TorchData

使用 pip 安装：

pip install torchdata

或使用 Conda 安装：

conda install -c pytorch torchdata

基本使用

TorchData 的核心组件是 StatefulDataLoader，它是 PyTorch 原生 DataLoader 的直接替代品，增加了 state_dict 和 load_state_dict 接口，支持断点续训。

示例：使用 StatefulDataLoader

import torch
from torchdata.stateful_dataloader import StatefulDataLoader

# 1. 准备数据集
dataset = torch.arange(100)

# 2. 初始化 StatefulDataLoader (用法同 torch.utils.data.DataLoader)
dataloader = StatefulDataLoader(dataset, batch_size=10)

# 3. 模拟训练过程
for batch in dataloader:
    print(batch)
    # 4. 在训练过程中保存状态 (例如每个 epoch 结束或中断时)
    state = dataloader.state_dict()
    
    # 5. 恢复状态 (例如从检查点恢复训练)
    # dataloader.load_state_dict(state)