torch.utils.data

PyTorch 資料載入工具的核心是 torch.utils.data.DataLoader 類別。 它代表一個資料集的 Python 可迭代物件，支援

• map-style 和 iterable-style 資料集,

• 自定義資料載入順序,

• 自動批次處理,

• 單進程和多進程數據載入,

• 自動記憶體釘選.

這些選項由 DataLoader 的建構函式參數配置，其簽名如下：

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None, *, prefetch_factor=2,
           persistent_workers=False)

以下章節將詳細描述這些選項的效果和用法。

資料集類型

DataLoader 建構函式最重要的參數是 dataset，它指定要從中載入數據的資料集物件。 PyTorch 支援兩種不同的資料集類型

• 映射式資料集,

• 可迭代式資料集.

映射式資料集

映射式資料集是指實現了 __getitem__() 和 __len__() 協議的資料集，並且表示從（可能非整數）索引/鍵到數據樣本的映射。

例如，當使用 dataset[idx] 存取這樣的資料集時，可以從磁碟上的資料夾讀取第 idx 個圖像及其對應的標籤。

有關更多詳細資訊，請參閱 Dataset。

可迭代式資料集

可迭代式資料集是 IterableDataset 子類的一個實例，它實現了 __iter__() 協議，並且表示數據樣本上的可迭代物件。 此類型的資料集特別適合於隨機讀取成本高昂甚至不可能的情況，以及批量大小取決於提取的數據的情況。

例如，當調用 iter(dataset) 時，這樣的資料集可以返回從資料庫、遠端伺服器甚至即時產生的日誌讀取的數據流。

有關更多詳細資訊，請參閱 IterableDataset。

注意

當使用 IterableDataset 配合 多進程數據載入 時。 同一個資料集物件會在每個 worker 進程上複製，因此必須以不同的方式配置副本以避免重複數據。 有關如何實現這一點，請參閱 IterableDataset 文件。

數據載入順序和 Sampler

對於 可迭代式資料集，數據載入順序完全由使用者定義的可迭代物件控制。 這使得更容易實現區塊讀取和動態批量大小（例如，每次產生一個批處理樣本）。

本節的其餘部分涉及 映射式資料集 的情況。 torch.utils.data.Sampler 類別用於指定數據載入中使用的索引/鍵的順序。 它們表示數據集索引上的可迭代物件。 例如，在隨機梯度下降 (SGD) 的常見情況下，Sampler 可以隨機排列索引列表並一次產生一個，或者為小批量 SGD 產生少量索引。

將基於 DataLoader 的 shuffle 參數自動建構一個循序或隨機的 sampler。 或者，使用者可以使用 sampler 參數來指定一個自定義的 Sampler 物件，該物件每次都產生下一個要獲取的索引/鍵。

一次產生一批索引列表的自定義 Sampler 可以作為 batch_sampler 參數傳遞。 也可以通過 batch_size 和 drop_last 參數啟用自動批處理。 有關此的更多詳細資訊，請參閱 下一節。

注意

sampler 和 batch_sampler 都不與可迭代式資料集相容，因為此類資料集沒有鍵或索引的概念。

載入批處理和非批處理數據

DataLoader 支持通過參數 batch_size、drop_last、batch_sampler 和 collate_fn（具有默認函數）將單獨提取的數據樣本自動整理成批次。

自動批次處理 (預設)

這是最常見的情況，對應於提取一個小批次的資料並將它們整理成批次樣本，即包含張量 (Tensors)，其中一個維度是批次維度（通常是第一個）。

當 batch_size (預設 1) 不是 None 時，資料載入器會產生批次樣本，而不是個別樣本。 batch_size 和 drop_last 參數用於指定資料載入器如何獲取資料集鍵的批次。對於 map-style 資料集，使用者也可以指定 batch_sampler，它一次產生一個鍵的列表。

注意

batch_size 和 drop_last 參數本質上是用於從 sampler 建構一個 batch_sampler。對於 map-style 資料集，sampler 要么由使用者提供，要么基於 shuffle 參數建構。對於 iterable-style 資料集，sampler 是一個虛擬的無限 sampler。有關 sampler 的更多詳細資訊，請參閱 此節。

注意

當從具有 多進程的 iterable-style 資料集提取資料時，drop_last 參數會丟棄每個 worker 的資料集副本的最後一個非完整批次。

在使用 sampler 中的索引提取樣本列表後，作為 collate_fn 參數傳遞的函式用於將樣本列表整理成批次。

在這種情況下，從 map-style 資料集載入大致相當於

for indices in batch_sampler:
    yield collate_fn([dataset[i] for i in indices])

從 iterable-style 資料集載入大致相當於

dataset_iter = iter(dataset)
for indices in batch_sampler:
    yield collate_fn([next(dataset_iter) for _ in indices])

可以使用自定義的 collate_fn 來自定義整理，例如，將序列資料填充到批次的最大長度。有關 collate_fn 的更多資訊，請參閱 此節。

禁用自動批次處理

在某些情況下，使用者可能希望在資料集程式碼中手動處理批次處理，或者只是載入個別樣本。例如，直接載入批次資料可能更便宜（例如，從資料庫批量讀取或讀取連續的記憶體塊），或者批次大小取決於資料，或者程式設計為處理個別樣本。在這些情況下，最好不要使用自動批次處理（其中 collate_fn 用於整理樣本），而是讓資料載入器直接傳回 dataset 物件的每個成員。

當 batch_size 和 batch_sampler 都是 None 時（batch_sampler 的預設值已經是 None），自動批次處理被禁用。從 dataset 獲得的每個樣本都會使用作為 collate_fn 參數傳遞的函式進行處理。

當自動批次處理被禁用時，預設的 collate_fn 只是將 NumPy 陣列轉換為 PyTorch 張量，並保持其他所有內容不變。

在這種情況下，從 map-style 資料集載入大致相當於

for index in sampler:
    yield collate_fn(dataset[index])

從 iterable-style 資料集載入大致相當於

for data in iter(dataset):
    yield collate_fn(data)

有關 collate_fn 的更多資訊，請參閱 此節。

使用 collate_fn

當啟用或禁用自動批次處理時，collate_fn 的使用略有不同。

當自動批次處理被禁用時，會為每個個別的資料樣本呼叫 collate_fn，並且輸出會從資料載入器迭代器產生。在這種情況下，預設的 collate_fn 只是將 NumPy 陣列轉換為 PyTorch 張量。

當啟用自動批次處理時，每次都會使用資料樣本的列表呼叫 collate_fn。預期它會將輸入樣本整理成批次，以便從資料載入器迭代器產生。本節的其餘部分描述了預設 collate_fn (default_collate()) 的行為。

例如，如果每個資料樣本由一個 3 通道圖像和一個整數類別標籤組成，即資料集的每個元素傳回一個元組 (image, class_index)，則預設的 collate_fn 將此類元組的列表整理成一個批次的圖像張量和一個批次的類別標籤張量的單個元組。特別是，預設的 collate_fn 具有以下屬性

• 它始終在前面新增一個維度作為批次維度。

• 它會自動將 NumPy 陣列和 Python 數值轉換為 PyTorch 張量。

• 它保留資料結構，例如，如果每個樣本是一個字典，它會輸出一個具有相同鍵集但批次的張量作為值（如果這些值無法轉換為張量，則為列表）的字典。list s、tuple s、namedtuple s 等也是如此。

使用者可以使用自定義的 collate_fn 來實現自定義批次處理，例如，沿著第一個以外的維度整理，填充各種長度的序列，或新增對自定義資料類型的支援。

如果您遇到 DataLoader 的輸出具有與您的預期不同的維度或類型的情況，您可能需要檢查您的 collate_fn。

單進程和多進程資料載入

預設情況下，DataLoader 使用單進程資料載入。

在一個 Python 進程中，全域直譯器鎖 (Global Interpreter Lock, GIL) 阻止了在多個執行緒上真正完全平行化 Python 程式碼。為了避免資料載入阻塞運算程式碼，PyTorch 提供了一個簡單的切換方式來執行多進程資料載入，只需將參數 num_workers 設定為一個正整數即可。

單進程資料載入（預設）

在此模式下，資料提取是在初始化 DataLoader 的同一個進程中完成的。因此，資料載入可能會阻塞運算。然而，當用於在進程之間共享資料的資源（例如，共享記憶體、檔案描述符）受到限制，或者當整個資料集很小並且可以完全載入到記憶體中時，可能更偏好此模式。此外，單進程載入通常會顯示更易讀的錯誤追蹤，因此對於除錯很有用。

多進程資料載入

將參數 num_workers 設定為正整數將會啟用多進程資料載入，並使用指定數量的載入器工作進程。

警告

經過幾次迭代後，loader worker 進程將會消耗與父進程相同數量的 CPU 記憶體，因為所有在 worker 進程中存取的 Python 物件都存在於父進程中。如果 Dataset 包含大量資料（例如，您在 Dataset 構造時載入一個非常大的檔名列表）和/或您使用大量的 workers，這可能會產生問題 (總體記憶體用量是 worker 數量 * 父進程大小)。最簡單的解決方法是用非引用計數表示（例如 Pandas、Numpy 或 PyArrow 物件）替換 Python 物件。查看 issue #13246 了解更多關於為何會發生這種情況以及如何解決這些問題的範例程式碼。

在此模式下，每次建立 DataLoader 的迭代器時（例如，當您呼叫 enumerate(dataloader) 時），會建立 num_workers 個工作進程。此時，dataset、collate_fn 和 worker_init_fn 會被傳遞給每個 worker，它們被用來初始化和提取資料。這意味著資料集存取及其內部 IO、轉換（包括 collate_fn）在 worker 進程中運行。

torch.utils.data.get_worker_info() 在 worker 進程中返回各種有用的資訊（包括 worker id、資料集副本、初始 seed 等），並在主進程中返回 None。使用者可以在資料集程式碼和/或 worker_init_fn 中使用此函數來單獨配置每個資料集副本，並確定程式碼是否在 worker 進程中運行。例如，這在分片資料集時特別有用。

對於 map-style 資料集，主進程使用 sampler 生成索引，並將它們發送到 workers。因此，任何 shuffle 隨機化都在主進程中完成，該進程透過分配要載入的索引來指導載入。

對於 iterable-style 資料集，由於每個 worker 進程都會獲得 dataset 物件的副本，因此 naive 的多進程載入通常會導致資料重複。使用 torch.utils.data.get_worker_info() 和/或 worker_init_fn，使用者可以獨立配置每個副本。（請參閱 IterableDataset 文件，了解如何實現此目的。）由於類似的原因，在多進程載入中，drop_last 參數會丟棄每個 worker 的 iterable-style 資料集副本的最後一個非完整批次。

一旦達到迭代的結尾，或者當迭代器被垃圾回收時，worker 就會關閉。

警告

通常不建議在多進程載入中返回 CUDA 張量，因為在使用 CUDA 和在多重處理中共享 CUDA 張量時存在許多微妙之處（參見 多重處理中的 CUDA）。相反，我們建議使用 自動記憶體釘選（即，設定 pin_memory=True），這可以快速將資料傳輸到啟用 CUDA 的 GPU。

特定於平台的行為

由於 workers 依賴於 Python multiprocessing，因此在 Windows 上與 Unix 上相比，worker 啟動行為有所不同。

• 在 Unix 上，fork() 是預設的 multiprocessing 啟動方法。 使用 fork()，子 worker 通常可以直接透過複製的位址空間存取 dataset 和 Python 參數函數。

• 在 Windows 或 MacOS 上，spawn() 是預設的 multiprocessing 啟動方法。 使用 spawn()，會啟動另一個直譯器，執行你的主程式碼，然後是接收 dataset、collate_fn 和其他引數的內部 worker 函式，這些引數透過 pickle 序列化傳輸。

這種獨立的序列化表示你應該採取兩個步驟，以確保在使用多進程資料載入時與 Windows 相容

• 將你的大部分主程式碼包在 if __name__ == '__main__': 區塊中，以確保在每個 worker 進程啟動時，它不會再次執行（很可能產生錯誤）。 你可以將你的 dataset 和 DataLoader 實例創建邏輯放在這裡，因為它不需要在 workers 中重新執行。

• 確保任何自訂的 collate_fn、worker_init_fn 或 dataset 程式碼都宣告為頂層定義，位於 __main__ 檢查之外。 這可確保它們在 worker 進程中可用。（這是必要的，因為函式僅作為參考而非 bytecode 進行 pickle 序列化。）

多進程資料載入中的隨機性

預設情況下，每個 worker 的 PyTorch 種子將設定為 base_seed + worker_id，其中 base_seed 是主進程使用其 RNG 產生的 long 值（因此，強制消耗一個 RNG 狀態）或指定的 generator。 但是，其他函式庫的種子可能會在初始化 worker 時重複，導致每個 worker 返回相同的隨機數。（請參閱 FAQ 中的 這一節。）

在 worker_init_fn 中，你可以使用 torch.utils.data.get_worker_info().seed 或 torch.initial_seed() 存取為每個 worker 設定的 PyTorch 種子，並在資料載入之前使用它來設定其他函式庫的種子。

記憶體固定

從固定（頁面鎖定）記憶體起始的 Host 到 GPU 複製速度更快。 有關何時以及如何通常使用固定記憶體的更多詳細資訊，請參閱 使用固定記憶體緩衝區。

對於資料載入，將 pin_memory=True 傳遞給 DataLoader 將會自動將取得的資料 Tensors 放入固定記憶體中，從而實現更快的資料傳輸到啟用 CUDA 的 GPU。

預設的記憶體固定邏輯僅識別 Tensors 以及包含 Tensors 的映射和可迭代對象。 預設情況下，如果固定邏輯看到一個自訂類型（如果你有一個返回自訂 batch 類型的 collate_fn，則會發生這種情況），或者如果你的 batch 中的每個元素都是自訂類型，則固定邏輯將無法識別它們，並且它將返回該 batch（或這些元素）而不固定記憶體。 若要為自訂 batch 或資料類型啟用記憶體固定，請在你的自訂類型上定義一個 pin_memory() 方法。

請參閱下面的範例。

範例

class SimpleCustomBatch:
    def __init__(self, data):
        transposed_data = list(zip(*data))
        self.inp = torch.stack(transposed_data[0], 0)
        self.tgt = torch.stack(transposed_data[1], 0)

    # custom memory pinning method on custom type
    def pin_memory(self):
        self.inp = self.inp.pin_memory()
        self.tgt = self.tgt.pin_memory()
        return self

def collate_wrapper(batch):
    return SimpleCustomBatch(batch)

inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
dataset = TensorDataset(inps, tgts)

loader = DataLoader(dataset, batch_size=2, collate_fn=collate_wrapper,
                    pin_memory=True)

for batch_ndx, sample in enumerate(loader):
    print(sample.inp.is_pinned())
    print(sample.tgt.is_pinned())

class torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=None, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None, generator=None, *, prefetch_factor=None, persistent_workers=False, pin_memory_device='', in_order=True)

資料載入器結合了資料集和取樣器，並提供在給定資料集上的可迭代物件。

DataLoader 支援單進程或多進程載入的映射樣式和可迭代樣式資料集，自訂載入順序以及可選的自動批次處理（collation）和記憶體釘選。

有關更多詳細資訊，請參閱 torch.utils.data 文件頁面。

參數

• dataset (Dataset) – 從中載入資料的資料集。

• batch_size (int, optional) – 每次載入的批次樣本數量（預設值：1）。

• shuffle (bool, optional) – 設定為 True 以在每個 epoch 重新洗牌資料（預設值：False）。

• sampler (Sampler or Iterable, optional) – 定義從資料集中抽取樣本的策略。 可以是任何實現了 __len__ 的 Iterable。 如果指定，則不能指定 shuffle。

• batch_sampler (Sampler or Iterable, optional) – 類似於 sampler，但一次返回一批索引。 與 batch_size，shuffle，sampler 和 drop_last 互斥。

• num_workers (int, optional) – 用於資料載入的子進程數量。 0 表示資料將在主進程中載入。（預設值：0）

• collate_fn (Callable, optional) – 合併樣本列表以形成 Tensor 的小批次。 當從映射樣式資料集使用批次載入時使用。

• pin_memory (bool, optional) – 如果為 True，資料載入器將在返回 Tensor 之前將其複製到 device/CUDA 釘選的記憶體中。 如果您的資料元素是自訂類型，或者您的 collate_fn 返回的批次是自訂類型，請參閱下面的範例。

• drop_last (bool, optional) – 如果資料集大小無法被批次大小整除，則設定為 True 以捨棄最後一個不完整的批次。 如果為 False 並且資料集的大小無法被批次大小整除，則最後一個批次將會更小。（預設值：False）

• timeout (numeric, optional) – 如果為正數，則為從 worker 收集批次的逾時值。 應始終為非負數。（預設值：0）

• worker_init_fn (Callable, optional) – 如果不是 None，則將在每個 worker 子進程上調用此函數，並將 worker ID（[0, num_workers - 1] 中的一個整數）作為輸入，在設定 seed 之後和資料載入之前。（預設值：None）

• multiprocessing_context (str or multiprocessing.context.BaseContext, optional) – 如果為 None，則將使用作業系統的預設 多進程上下文。（預設值：None）

• generator (torch.Generator, optional) – 如果不是 None，則 RandomSampler 將使用此 RNG 生成隨機索引，多進程將使用此 RNG 生成 worker 的 base_seed。（預設值：None）

• prefetch_factor (int, optional, keyword-only arg) – 每個 worker 預先載入的批次數量。 2 代表所有 worker 總共會預先載入 2 * num_workers 個批次。(預設值取決於 num_workers 的設定值。如果 num_workers=0，則預設值為 None。否則，如果 num_workers > 0，則預設值為 2)。

• persistent_workers (bool, optional) – 如果為 True，則資料載入器在資料集被消耗一次後不會關閉 worker 處理程序。這允許保持 workers 的 Dataset 實例存活。(預設值: False)

• pin_memory_device (str, optional) – 如果 pin_memory 為 True，則要 pin_memory 到的裝置。

• in_order (bool, optional) – 如果為 False，則資料載入器不會強制批次以先進先出的順序返回。僅在 num_workers > 0 時適用。(預設值: True)

警告

如果使用 spawn 啟動方法，則 worker_init_fn 不能是無法 pickled 的物件，例如 lambda 函式。 有關 PyTorch 中多處理的更多詳細資訊，請參閱多處理最佳實踐。

警告

len(dataloader) 啟發式方法基於使用的 sampler 的長度。 當 dataset 是 IterableDataset 時，它會改為根據 len(dataset) / batch_size 返回一個估計值，並根據 drop_last 進行適當的捨入，無論多處理載入配置如何。 這表示 PyTorch 可以做出的最佳猜測，因為 PyTorch 信任使用者 dataset 程式碼可以正確處理多處理載入以避免重複資料。

但是，如果分片導致多個 worker 具有不完整的最後一批，則此估計值仍然可能不準確，因為 (1) 否則完整的批次可以被分成多個批次，並且 (2) 當設置 drop_last 時，可能會丟棄超過一個批次價值的樣本。 不幸的是，PyTorch 無法普遍檢測到此類情況。

有關這兩種資料集類型以及 IterableDataset 如何與 多處理資料載入互動的更多詳細資訊，請參閱資料集類型。

警告

有關隨機種子相關問題，請參閱可重現性、我的資料載入器 worker 返回相同的隨機數以及多處理資料載入中的隨機性註解。

警告

將 in_order 設置為 False 會損害可重現性，並且在資料不平衡的情況下，可能導致將偏斜的資料分佈饋送到訓練器。

class torch.utils.data.Dataset

代表 Dataset 的抽象類別。

所有代表從鍵到資料樣本的映射的資料集都應該是它的子類別。 所有子類別都應該覆寫 __getitem__()，以支援獲取給定鍵的資料樣本。 子類別也可以選擇性地覆寫 __len__()，許多 Sampler 實作和 DataLoader 的預設選項都希望它返回資料集的大小。 子類別也可以選擇性地實作 __getitems__()，以加速批次樣本載入。 此方法接受批次樣本的索引列表並返回樣本列表。

注意

預設情況下，DataLoader 會建構一個產生整數索引的索引 sampler。 為了使其與具有非整數索引/鍵的 map-style 資料集一起使用，必須提供自訂 sampler。

class torch.utils.data.IterableDataset

可迭代的 Dataset。

所有代表資料樣本可迭代物件的資料集都應該是它的子類別。 當資料來自串流時，這種形式的資料集特別有用。

所有子類別都應該覆寫 __iter__()，它將返回此資料集中樣本的迭代器。

當子類別與 DataLoader 一起使用時，資料集中的每個項目都會從 DataLoader 迭代器中產生。 當 num_workers > 0 時，每個 worker process 都會有資料集物件的不同副本，因此通常希望獨立設定每個副本，以避免 worker 傳回重複的資料。 在 worker process 中呼叫 get_worker_info() 會傳回有關 worker 的資訊。 它可以被用在資料集的 __iter__() 方法中，或者用於 DataLoader 的 worker_init_fn 選項中，以修改每個副本的行為。

範例 1：在 __iter__() 中跨所有 worker 分割工作負載

>>> class MyIterableDataset(torch.utils.data.IterableDataset):
...     def __init__(self, start, end):
...         super(MyIterableDataset).__init__()
...         assert end > start, "this example code only works with end >= start"
...         self.start = start
...         self.end = end
...
...     def __iter__(self):
...         worker_info = torch.utils.data.get_worker_info()
...         if worker_info is None:  # single-process data loading, return the full iterator
...             iter_start = self.start
...             iter_end = self.end
...         else:  # in a worker process
...             # split workload
...             per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
...             worker_id = worker_info.id
...             iter_start = self.start + worker_id * per_worker
...             iter_end = min(iter_start + per_worker, self.end)
...         return iter(range(iter_start, iter_end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)

>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[tensor([3]), tensor([4]), tensor([5]), tensor([6])]

>>> # Mult-process loading with two worker processes
>>> # Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]

>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12)))
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]

範例 2：使用 worker_init_fn 跨所有 worker 分割工作負載

>>> class MyIterableDataset(torch.utils.data.IterableDataset):
...     def __init__(self, start, end):
...         super(MyIterableDataset).__init__()
...         assert end > start, "this example code only works with end >= start"
...         self.start = start
...         self.end = end
...
...     def __iter__(self):
...         return iter(range(self.start, self.end))
...
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
>>> ds = MyIterableDataset(start=3, end=7)

>>> # Single-process loading
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
[3, 4, 5, 6]
>>>
>>> # Directly doing multi-process loading yields duplicate data
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
[3, 3, 4, 4, 5, 5, 6, 6]

>>> # Define a `worker_init_fn` that configures each dataset copy differently
>>> def worker_init_fn(worker_id):
...     worker_info = torch.utils.data.get_worker_info()
...     dataset = worker_info.dataset  # the dataset copy in this worker process
...     overall_start = dataset.start
...     overall_end = dataset.end
...     # configure the dataset to only process the split workload
...     per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
...     worker_id = worker_info.id
...     dataset.start = overall_start + worker_id * per_worker
...     dataset.end = min(dataset.start + per_worker, overall_end)
...

>>> # Mult-process loading with the custom `worker_init_fn`
>>> # Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))
[3, 5, 4, 6]

>>> # With even more workers
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12, worker_init_fn=worker_init_fn)))
[3, 4, 5, 6]

class torch.utils.data.TensorDataset(*tensors)

封裝 tensors 的資料集。

每個樣本將會透過沿著第一個維度索引 tensors 來檢索。

參數

*tensors (Tensor) – 具有相同第一個維度大小的 tensors。

class torch.utils.data.StackDataset(*args, **kwargs)

作為多個資料集堆疊的資料集。

此類別適用於組裝複雜輸入資料的不同部分，這些資料以資料集的形式提供。

範例

>>> images = ImageDataset()
>>> texts = TextDataset()
>>> tuple_stack = StackDataset(images, texts)
>>> tuple_stack[0] == (images[0], texts[0])
>>> dict_stack = StackDataset(image=images, text=texts)
>>> dict_stack[0] == {'image': images[0], 'text': texts[0]}

參數

• *args (Dataset) – 作為元組傳回的堆疊資料集。

• **kwargs (Dataset) – 作為 dict 傳回的堆疊資料集。

class torch.utils.data.ConcatDataset(datasets)

作為多個資料集串聯的資料集。

此類別適用於組裝不同的現有資料集。

參數

datasets (sequence) – 要串聯的資料集清單

class torch.utils.data.ChainDataset(datasets)

用於鏈結多個 IterableDataset 的資料集。

此類別適用於組裝不同的現有資料集流。 鏈結操作是即時完成的，因此使用此類別串聯大規模資料集會很有效率。

參數

datasets (iterable of IterableDataset) – 要鏈結在一起的資料集

class torch.utils.data.Subset(dataset, indices)

位於指定索引處的資料集子集。

參數

• dataset (Dataset) – 整個資料集

• indices (sequence) – 為了子集選擇的整個集合中的索引

torch.utils.data._utils.collate.collate(batch, *, collate_fn_map=None)

通用的 collate 函式，用於處理每個 batch 中元素的集合類型。

此函式也會開啟函式登錄，以處理特定的元素類型。 default_collate_fn_map 針對 tensors、numpy 陣列、數字和字串提供預設的 collate 函式。

參數

• batch – 要 collate 的單個 batch

• collate_fn_map (Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]]) – 從元素類型到對應 collate 函式的可選字典映射。 如果此字典中不存在元素類型，則此函式將按插入順序逐一檢查字典的每個鍵，如果元素類型是該鍵的子類別，則會調用對應的 collate 函式。

範例

>>> def collate_tensor_fn(batch, *, collate_fn_map):
...     # Extend this function to handle batch of tensors
...     return torch.stack(batch, 0)
>>> def custom_collate(batch):
...     collate_map = {torch.Tensor: collate_tensor_fn}
...     return collate(batch, collate_fn_map=collate_map)
>>> # Extend `default_collate` by in-place modifying `default_collate_fn_map`
>>> default_collate_fn_map.update({torch.Tensor: collate_tensor_fn})

注意

每個 collate 函式都需要一個用於 batch 的位置引數和一個用於 collate 函式字典的關鍵字引數，如 collate_fn_map。

torch.utils.data.default_collate(batch)

接收一個批次 (batch) 的資料，並將批次中的元素放入一個張量 (tensor)，並增加一個額外的外部維度 - 批次大小 (batch size)。

確切的輸出類型可以是 torch.Tensor、Sequence 的 torch.Tensor、Collection 的 torch.Tensor，或保持不變，具體取決於輸入類型。當 batch_size 或 batch_sampler 在 DataLoader 中定義時，這會被用作預設的整理 (collation) 函式。

以下是基於批次中元素類型的通用輸入類型到輸出類型的映射：

• torch.Tensor -> torch.Tensor (具有新增的外部維度批次大小)

• NumPy Arrays -> torch.Tensor

• float -> torch.Tensor

• int -> torch.Tensor

• str -> str (不變)

• bytes -> bytes (不變)

• Mapping[K, V_i] -> Mapping[K, default_collate([V_1, V_2, …])]

• NamedTuple[V1_i, V2_i, …] -> NamedTuple[default_collate([V1_1, V1_2, …]), default_collate([V2_1, V2_2, …]), …]

• Sequence[V1_i, V2_i, …] -> Sequence[default_collate([V1_1, V1_2, …]), default_collate([V2_1, V2_2, …]), …]

參數

batch – 要 collate 的單個 batch

範例

>>> # Example with a batch of `int`s:
>>> default_collate([0, 1, 2, 3])
tensor([0, 1, 2, 3])
>>> # Example with a batch of `str`s:
>>> default_collate(['a', 'b', 'c'])
['a', 'b', 'c']
>>> # Example with `Map` inside the batch:
>>> default_collate([{'A': 0, 'B': 1}, {'A': 100, 'B': 100}])
{'A': tensor([  0, 100]), 'B': tensor([  1, 100])}
>>> # Example with `NamedTuple` inside the batch:
>>> Point = namedtuple('Point', ['x', 'y'])
>>> default_collate([Point(0, 0), Point(1, 1)])
Point(x=tensor([0, 1]), y=tensor([0, 1]))
>>> # Example with `Tuple` inside the batch:
>>> default_collate([(0, 1), (2, 3)])
[tensor([0, 2]), tensor([1, 3])]
>>> # Example with `List` inside the batch:
>>> default_collate([[0, 1], [2, 3]])
[tensor([0, 2]), tensor([1, 3])]
>>> # Two options to extend `default_collate` to handle specific type
>>> # Option 1: Write custom collate function and invoke `default_collate`
>>> def custom_collate(batch):
...     elem = batch[0]
...     if isinstance(elem, CustomType):  # Some custom condition
...         return ...
...     else:  # Fall back to `default_collate`
...         return default_collate(batch)
>>> # Option 2: In-place modify `default_collate_fn_map`
>>> def collate_customtype_fn(batch, *, collate_fn_map=None):
...     return ...
>>> default_collate_fn_map.update(CustomType, collate_customtype_fn)
>>> default_collate(batch)  # Handle `CustomType` automatically

torch.utils.data.default_convert(data)

將每個 NumPy 陣列元素轉換為 torch.Tensor。

如果輸入是 Sequence、Collection 或 Mapping，它會嘗試將裡面的每個元素轉換為 torch.Tensor。如果輸入不是 NumPy 陣列，則保持不變。 當 batch_sampler 和 batch_size 都沒有在 DataLoader 中定義時，這會被用作預設的整理 (collation) 函式。

一般的輸入類型到輸出類型的映射與 default_collate() 類似。 有關更多詳細訊息，請參閱那裡的描述。

參數

data – 要轉換的單個資料點

範例

>>> # Example with `int`
>>> default_convert(0)
0
>>> # Example with NumPy array
>>> default_convert(np.array([0, 1]))
tensor([0, 1])
>>> # Example with NamedTuple
>>> Point = namedtuple('Point', ['x', 'y'])
>>> default_convert(Point(0, 0))
Point(x=0, y=0)
>>> default_convert(Point(np.array(0), np.array(0)))
Point(x=tensor(0), y=tensor(0))
>>> # Example with List
>>> default_convert([np.array([0, 1]), np.array([2, 3])])
[tensor([0, 1]), tensor([2, 3])]

torch.utils.data.get_worker_info()

返回有關當前 DataLoader 迭代器工作進程的訊息。

在 worker 中呼叫時，這會返回一個保證具有以下屬性的物件：

• id: 當前 worker 的 ID。

• num_workers: worker 的總數。

• seed: 為當前 worker 設定的隨機種子。 此值由主進程 RNG 和 worker ID 決定。 有關更多詳細訊息，請參閱 DataLoader 的文件。

• dataset: **此**進程中資料集物件的副本。 請注意，這將是與主進程中不同的物件。

在主進程中呼叫時，這會返回 None。

注意

當在傳遞給 DataLoader 的 worker_init_fn 中使用時，此方法可用於以不同的方式設定每個 worker 進程，例如，使用 worker_id 來配置 dataset 物件以僅讀取分片資料集的特定部分，或使用 seed 來為資料集程式碼中使用的其他程式庫設定種子。

返回類型

Optional[WorkerInfo]

torch.utils.data.random_split(dataset, lengths, generator=<torch._C.Generator object>)

將資料集隨機分割為指定長度的不重疊新資料集。

如果給定一個總和為 1 的分數列表，則長度將自動計算為 floor(frac * len(dataset))，其中 frac 為提供的每個分數。

計算長度後，如果還有任何餘數，則將以循環方式將 1 個計數分配給長度，直到沒有餘數為止。

您可以選擇修正產生器以獲得可重現的結果，例如：

範例

>>> generator1 = torch.Generator().manual_seed(42)
>>> generator2 = torch.Generator().manual_seed(42)
>>> random_split(range(10), [3, 7], generator=generator1)
>>> random_split(range(30), [0.3, 0.3, 0.4], generator=generator2)

參數

• dataset (Dataset) – 要分割的資料集

• lengths (sequence) – 要產生的分割的長度或分數

• generator (Generator) – 用於隨機排列的產生器。

返回類型

List[Subset[_T]]

class torch.utils.data.Sampler(data_source=None)

所有 Sampler 的基底類別。

每個 Sampler 子類別都必須提供一個 __iter__() 方法，提供一種迭代資料集元素的索引或索引列表（批次）的方式，並且可以提供一個 __len__() 方法，該方法會回傳迭代器的長度。

參數

data_source (Dataset) – 此參數未使用，將在 2.2.0 中移除。您可能仍然有使用它的自定義實作。

範例

>>> class AccedingSequenceLengthSampler(Sampler[int]):
>>>     def __init__(self, data: List[str]) -> None:
>>>         self.data = data
>>>
>>>     def __len__(self) -> int:
>>>         return len(self.data)
>>>
>>>     def __iter__(self) -> Iterator[int]:
>>>         sizes = torch.tensor([len(x) for x in self.data])
>>>         yield from torch.argsort(sizes).tolist()
>>>
>>> class AccedingSequenceLengthBatchSampler(Sampler[List[int]]):
>>>     def __init__(self, data: List[str], batch_size: int) -> None:
>>>         self.data = data
>>>         self.batch_size = batch_size
>>>
>>>     def __len__(self) -> int:
>>>         return (len(self.data) + self.batch_size - 1) // self.batch_size
>>>
>>>     def __iter__(self) -> Iterator[List[int]]:
>>>         sizes = torch.tensor([len(x) for x in self.data])
>>>         for batch in torch.chunk(torch.argsort(sizes), len(self)):
>>>             yield batch.tolist()

注意

DataLoader 並非嚴格要求 __len__() 方法，但在任何涉及 DataLoader 長度的計算中，都應包含此方法。

class torch.utils.data.SequentialSampler(data_source)

依序取樣元素，始終保持相同的順序。

參數

data_source (Dataset) – 要从中取样的資料集

class torch.utils.data.RandomSampler(data_source, replacement=False, num_samples=None, generator=None)

隨機取樣元素。如果沒有放回，則從打亂的資料集中取樣。

如果有放回，則使用者可以指定 num_samples 來提取樣本。

參數

• data_source (Dataset) – 要从中取样的資料集

• replacement (bool) – 如果 True，則按需放回取樣，預設值=``False``

• num_samples (int) – 要提取的樣本數，預設值=`len(dataset)`。

• generator (Generator) – 用於取樣的 Generator。

class torch.utils.data.SubsetRandomSampler(indices, generator=None)

從給定的索引列表中隨機取樣元素，沒有放回。

參數

• indices (sequence) – 一個索引序列

• generator (Generator) – 用於取樣的 Generator。

class torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True, generator=None)

使用給定的機率（權重）從 [0,..,len(weights)-1] 中取樣元素。

參數

• weights (sequence) – 一個權重序列，不一定總和為 1

• num_samples (int) – 要提取的樣本數

• replacement (bool) – 如果 True，則放回取樣。 如果沒有，則在沒有放回的情況下提取它們，這意味著當為一行提取樣本索引時，不能再為該行提取該索引。

• generator (Generator) – 用於取樣的 Generator。

範例

>>> list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True))
[4, 4, 1, 4, 5]
>>> list(WeightedRandomSampler([0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False))
[0, 1, 4, 3, 2]

class torch.utils.data.BatchSampler(sampler, batch_size, drop_last)

封裝另一個取樣器以產生索引的小批量。

參數

• sampler (Sampler or Iterable) – 基礎取樣器。 可以是任何可迭代的物件

• batch_size (int) – 小批量的大小。

• drop_last (bool) – 如果 True，當最後一個批次的大小小於 batch_size 時，取樣器會捨棄最後一個批次。

範例

>>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False))
[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
>>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True))
[[0, 1, 2], [3, 4, 5], [6, 7, 8]]

class torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)

限制資料載入到資料集子集的取樣器。

它與 torch.nn.parallel.DistributedDataParallel 結合使用特別有用。 在這種情況下，每個行程都可以將 DistributedSampler 實例作為 DataLoader 取樣器傳遞，並載入專屬於它的原始資料集的子集。

注意

假設資料集的大小是恆定的，並且它的任何實例始終以相同的順序傳回相同的元素。

參數

• dataset (Dataset) – 用於取樣的資料集。

• num_replicas (int, optional) – 參與分散式訓練的行程數。 預設情況下，world_size 是從目前的分散式群組中檢索的。

• rank (int, optional) – 目前行程在 num_replicas 中的排名。 預設情況下，rank 是從目前的分散式群組中檢索的。

• shuffle (bool, optional) – 如果 True (預設)，取樣器會打亂索引。

• seed (int, optional) – 如果 shuffle=True，則用於打亂取樣器的隨機種子。 這個數字在分散式群組中的所有行程中應該相同。 預設值：0。

• drop_last (bool, optional) – 如果 True，則取樣器將捨棄資料的尾部，使其可以均勻地分配到副本數量中。 如果 False，則取樣器將新增額外的索引，使資料可以均勻地分配到副本中。 預設值：False。

警告

在分散式模式下，在每次 epoch 的開頭、**在**建立 DataLoader 迭代器**之前**，呼叫 set_epoch() 方法對於使洗牌在多個 epoch 中正常工作是必要的。 否則，將始終使用相同的順序。

範例

>>> sampler = DistributedSampler(dataset) if is_distributed else None
>>> loader = DataLoader(dataset, shuffle=(sampler is None),
...                     sampler=sampler)
>>> for epoch in range(start_epoch, n_epochs):
...     if is_distributed:
...         sampler.set_epoch(epoch)
...     train(loader)