Python PyTorch InMemoryBinaryCriteoIterDataPipe用法及代码示例

本文简要介绍python语言中 torchrec.datasets.criteo.InMemoryBinaryCriteoIterDataPipe 的用法。

用法:

class torchrec.datasets.criteo.InMemoryBinaryCriteoIterDataPipe(dense_paths: List[str], sparse_paths: List[str], labels_paths: List[str], batch_size: int, rank: int, world_size: int, shuffle_batches: bool = False, hashes: Optional[List[int]] = None, path_manager_key: str = 'torchrec')

参数：

• dense_paths(List[str]) -密集 npy 文件的路径字符串列表。

• sparse_paths(List[str]) -稀疏 npy 文件的路径字符串列表。

• labels_paths(List[str]) -标签 npy 文件的路径字符串列表。

• batch_size(int) -批量大小。

• rank(int) -秩。

• world_size(int) -世界规模。

• shuffle_batches(bool) -是否洗牌

• hashes(可选的[int]) -每个特征的最大分类特征值列表。此列表的长度应为CAT_FEATURE_COUNT。

• path_manager_key(str) -用于从不同文件系统加载的路径管理器 key 。

基础：torch.utils.data.dataset.IterableDataset

Datapipe 旨在对 Criteo 数据集的二进制 (npy) 版本进行操作。将整个数据集加载到内存中，以防止磁盘速度影响整个过程。每个等级仅读取它负责的数据集部分的数据。

torchrec/datasets/scripts/preprocess_criteo.py 脚本可用于将 Criteo tsv 文件转换为该数据集预期的 npy 文件。

例子：

template = "/home/datasets/criteo/1tb_binary/day_{}_{}.npy"
datapipe = InMemoryBinaryCriteoIterDataPipe(
    dense_paths=[template.format(0, "dense"), template.format(1, "dense")],
    sparse_paths=[template.format(0, "sparse"), template.format(1, "sparse")],
    labels_paths=[template.format(0, "labels"), template.format(1, "labels")],
    batch_size=1024,
    rank=torch.distributed.get_rank(),
    world_size=torch.distributed.get_world_size(),
)
batch = next(iter(datapipe))