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))