get_dataloader¶

class torchrl.data.get_dataloader(batch_size: int, block_size: int, tensorclass_type: type, device: torch.device, dataset_name: str | None = None, infinite: bool = True, prefetch: int = 0, split: str = 'train', root_dir: str | None = None, from_disk: bool = False, num_workers: int | None = None)[原始碼]¶

建立資料集並從中返回一個數據載入器。

引數:

batch_size (int) – 資料載入器樣本的批次大小。
block_size (int) – 資料載入器中序列的最大長度。
tensorclass_type (tensorclass 類) – 一個具有 from_dataset() 方法的 tensorclass，該方法必須接受三個關鍵字引數：split（見下文）、max_length（將用於訓練的塊大小）和 dataset_name（一個指示資料集的字串）。還應支援 root_dir 和 from_disk 引數。
device (torch.device 或等效) – 樣本應被強制轉換到的裝置。
dataset_name (str, optional) – 資料集名稱。如果未提供且 tensorclass 支援，則將為正在使用的 tensorclass 收集預設資料集名稱。
infinite (bool, optional) – 如果為 True，則迭代將是無限的，這樣 next(iterator) 始終會返回一個值。預設為 True。
prefetch (int, optional) – 如果使用多執行緒資料載入，要預取的專案數。
split (str, optional) – 資料分割。可以是 "train" 或 "valid"。預設為 "train"。
root_dir (path, optional) – 儲存資料集的路徑。預設為 "$HOME/.cache/torchrl/data"
from_disk (bool, optional) – 如果為 True，將使用 datasets.load_from_disk()。否則，將使用 datasets.load_dataset()。預設為 False。
num_workers (int, optional) – 在分詞過程中呼叫的 datasets.dataset.map() 的工作執行緒數。預設為 max(os.cpu_count() // 2, 1)。

示例

>>> from torchrl.data.llm.reward import PairwiseDataset
>>> dataloader = get_dataloader(
...     batch_size=256, block_size=550, tensorclass_type=PairwiseDataset, device="cpu")
>>> for d in dataloader:
...     print(d)
...     break
PairwiseDataset(
    chosen_data=RewardData(
        attention_mask=Tensor(shape=torch.Size([256, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        input_ids=Tensor(shape=torch.Size([256, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        rewards=None,
        end_scores=None,
        batch_size=torch.Size([256]),
        device=cpu,
        is_shared=False),
    rejected_data=RewardData(
        attention_mask=Tensor(shape=torch.Size([256, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        input_ids=Tensor(shape=torch.Size([256, 550]), device=cpu, dtype=torch.int64, is_shared=False),
        rewards=None,
        end_scores=None,
        batch_size=torch.Size([256]),
        device=cpu,
        is_shared=False),
    batch_size=torch.Size([256]),
    device=cpu,
    is_shared=False)

get_dataloader¶

文件

教程

資源