1. 搞懂這三個組件數(shù)據(jù)管道才算真正入門剛接觸深度學(xué)習(xí)那會兒我寫訓(xùn)練腳本最怕的不是模型不收斂而是數(shù)據(jù)加載這一環(huán)。明明手里有幾千張圖、幾萬條文本卻不知道怎么把它們規(guī)規(guī)矩矩地喂給網(wǎng)絡(luò)。后來踩了不少坑才明白PyTorch 的數(shù)據(jù)加載體系其實就靠三個角色撐起來Dataset、TensorDataset和DataLoader。這三個東西配合好了訓(xùn)練循環(huán)干凈利落配合不好要么內(nèi)存爆掉要么訓(xùn)練速度慢得讓人想砸鍵盤。這篇文章我想從一個實際寫代碼的人的角度把這三個組件的用法、邊界和配合方式講透。不管你是剛學(xué) PyTorch 的新手還是已經(jīng)能跑通幾個 Demo 但總覺得數(shù)據(jù)管道寫得別扭的開發(fā)者都能從這里找到可以直接抄作業(yè)的寫法。核心關(guān)鍵詞就三個Dataset 負責(zé)定義數(shù)據(jù)長什么樣、怎么取TensorDataset 是 Dataset 的一個懶人封裝DataLoader 負責(zé)批量、打亂、并行加載。理解了這個分工后面所有的花式操作都是在這個骨架上長出來的。我見過太多人把這三個東西混著用比如在 DataLoader 里做歸一化、在 Dataset 里做 batch 拼接結(jié)果代碼又亂又慢。所以下面我會先講清楚每個組件該干什么、不該干什么再給出完整的實操流程和排查經(jīng)驗。你跟著走一遍基本就能把數(shù)據(jù)管道這塊的坑填平。2. 三個組件的分工與設(shè)計邏輯2.1 Dataset數(shù)據(jù)集的抽象接口Dataset 在 PyTorch 里是一個抽象類你自定義的數(shù)據(jù)集只要繼承它并實現(xiàn)兩個魔法方法就算完成__len__和__getitem__。前者告訴框架這個數(shù)據(jù)集有多少條樣本后者告訴框架第 i 條樣本長什么樣。為什么設(shè)計成這兩個方法因為訓(xùn)練時框架需要知道兩件事總共要迭代多少步以及給定一個索引怎么拿到數(shù)據(jù)。這種設(shè)計的好處是解耦——你的數(shù)據(jù)可以存在硬盤上、內(nèi)存里、數(shù)據(jù)庫里甚至實時生成只要你能按索引返回一條樣本Dataset 就不關(guān)心底層細節(jié)。我剛開始寫 Dataset 的時候犯過一個典型錯誤在__getitem__里做大量計算比如讀圖、解碼、增強全塞進去。結(jié)果 DataLoader 開多進程時每個 worker 都在重復(fù)做這些重活CPU 直接跑滿GPU 卻在等數(shù)據(jù)。后來才明白__getitem__應(yīng)該盡量輕重的預(yù)處理要么提前做好緩存要么放到 collate_fn 里批量做。一個標(biāo)準(zhǔn)的自定義 Dataset 長這樣from torch.utils.data import Dataset class MyDataset(Dataset): def __init__(self, data_list, label_list, transformNone): self.data_list data_list self.label_list label_list self.transform transform def __len__(self): return len(self.data_list) def __getitem__(self, idx): x self.data_list[idx] y self.label_list[idx] if self.transform: x self.transform(x) return x, y這里__init__只做輕量級的賦值真正的取數(shù)邏輯放在__getitem__。注意返回的必須是一個樣本不是一批樣本這是新手最容易搞混的地方。2.2 TensorDataset張量數(shù)據(jù)的快捷通道如果你的數(shù)據(jù)已經(jīng)是張量而且不需要復(fù)雜的預(yù)處理那 TensorDataset 就是為你準(zhǔn)備的。它把多個張量按第一維對齊自動幫你實現(xiàn)__len__和__getitem__。用法簡單到一行from torch.utils.data import TensorDataset dataset TensorDataset(features, labels)這行代碼等價于把 features 和 labels 按行打包取第 i 條時返回(features[i], labels[i])。它的價值在于省去了手寫 Dataset 的樣板代碼特別適合特征工程已經(jīng)做完、數(shù)據(jù)全在內(nèi)存里的場景。但 TensorDataset 有個硬性前提所有張量的第一維長度必須一致。我遇到過有人把長度不匹配的特征和標(biāo)簽硬塞進去結(jié)果訓(xùn)練時報索引越界排查半天才發(fā)現(xiàn)是數(shù)據(jù)對齊出了問題。所以用之前一定要確認features.shape[0] labels.shape[0]。另外TensorDataset 不支持自定義 transform。如果你需要在取數(shù)時做增強要么提前把增強后的數(shù)據(jù)存成張量要么老老實實寫自定義 Dataset。這一點很多人會忽略以為 TensorDataset 能包辦一切結(jié)果卡在數(shù)據(jù)增強上。2.3 DataLoader批量與并行的調(diào)度器DataLoader 是把 Dataset 變成可迭代批次的核心。它負責(zé)四件事批量組裝、打亂順序、多進程加載、內(nèi)存固定。這四個功能對應(yīng)四個關(guān)鍵參數(shù)batch_size、shuffle、num_workers、pin_memory。為什么需要 DataLoader因為訓(xùn)練時我們不會一條一條喂數(shù)據(jù)那樣 GPU 利用率極低。批量喂數(shù)據(jù)能讓矩陣運算吃滿硬件而打亂順序能防止模型學(xué)到樣本排列的虛假規(guī)律。多進程加載則是為了把數(shù)據(jù)準(zhǔn)備和模型計算重疊起來讓 GPU 不用干等。一個典型的 DataLoader 初始化from torch.utils.data import DataLoader loader DataLoader( dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue )這里每個參數(shù)都有講究。batch_size不是越大越好要看顯存和收斂效果num_workers設(shè)成 CPU 核心數(shù)左右比較合適設(shè)太大反而會因為進程切換開銷變慢pin_memoryTrue在 GPU 訓(xùn)練時能加速數(shù)據(jù)傳輸drop_lastTrue則在最后一批樣本數(shù)不足時丟棄避免 BatchNorm 在單樣本批次上出問題。2.4 三者的協(xié)作關(guān)系把這三個組件串起來看關(guān)系非常清晰Dataset 定義數(shù)據(jù)源TensorDataset 是 Dataset 的一種快捷實現(xiàn)DataLoader 消費 Dataset 并產(chǎn)出批次。數(shù)據(jù)流向是Dataset單條樣本→ DataLoader批次樣本→ 訓(xùn)練循環(huán)模型輸入理解這條鏈路后你就知道每個環(huán)節(jié)該放什么。數(shù)據(jù)清洗和增強放在 Dataset 的__getitem__或提前做好批量組裝和并行調(diào)度交給 DataLoader訓(xùn)練循環(huán)只負責(zé)拿批次、前向、反向、更新。職責(zé)分明代碼才好維護。3. 核心細節(jié)與實操要點拆解3.1 自定義 Dataset 的取數(shù)邏輯怎么寫才高效寫__getitem__時最關(guān)鍵的原則是只做必要的、輕量的操作。如果你的數(shù)據(jù)在硬盤上讀文件是不可避免的但解碼和增強可以考慮緩存。我常用的一個模式是第一次讀取時把解碼后的結(jié)果存到內(nèi)存字典里后續(xù)直接命中緩存。class CachedDataset(Dataset): def __init__(self, paths, labels): self.paths paths self.labels labels self.cache {} def __len__(self): return len(self.paths) def __getitem__(self, idx): if idx not in self.cache: img read_and_decode(self.paths[idx]) self.cache[idx] img return self.cache[idx], self.labels[idx]這個模式在數(shù)據(jù)集能放進內(nèi)存時非常有效但要注意內(nèi)存占用。如果數(shù)據(jù)量大到內(nèi)存放不下就得用 LRU 緩存或者干脆不緩存靠多進程并行來提速。另一個要點是返回類型要統(tǒng)一。如果有的樣本返回張量有的返回 PIL 圖像DataLoader 的默認 collate 函數(shù)會懵掉。所以要么在__getitem__里統(tǒng)一轉(zhuǎn)成張量要么自定義 collate_fn 來處理。3.2 TensorDataset 的適用邊界與轉(zhuǎn)換技巧TensorDataset 最適合的場景是數(shù)據(jù)已經(jīng)全部轉(zhuǎn)成張量且不需要在取數(shù)時做隨機變換。比如你做完特征工程后得到X_train和y_train兩個張量直接包起來就能用。但如果你手里是 NumPy 數(shù)組或 Python 列表需要先轉(zhuǎn)成張量import torch from torch.utils.data import TensorDataset X_tensor torch.tensor(X_train, dtypetorch.float32) y_tensor torch.tensor(y_train, dtypetorch.long) dataset TensorDataset(X_tensor, y_tensor)這里有個細節(jié)dtype 要選對。特征通常是 float32分類標(biāo)簽通常是 long即 int64。如果標(biāo)簽用了 float32交叉熵損失函數(shù)會報錯。我見過有人因為標(biāo)簽 dtype 不對訓(xùn)練時 loss 一直是 nan查了半天才發(fā)現(xiàn)是這里的問題。TensorDataset 還有一個限制它只能按索引返回對應(yīng)位置的張量切片不能做條件篩選或動態(tài)采樣。如果你需要這些功能就得回到自定義 Dataset。3.3 DataLoader 參數(shù)調(diào)優(yōu)的實戰(zhàn)經(jīng)驗DataLoader 的參數(shù)里num_workers和pin_memory是最容易調(diào)出問題的兩個。num_workers的默認值是 0意味著在主進程里加載數(shù)據(jù)。這樣雖然簡單但數(shù)據(jù)準(zhǔn)備和模型計算是串行的GPU 利用率上不去。設(shè)成大于 0 的值會開啟多進程加載但要注意Windows 上多進程需要把訓(xùn)練代碼放在if __name__ __main__:保護塊里否則會無限遞歸創(chuàng)建進程。每個 worker 會復(fù)制一份 Dataset如果 Dataset 里存了大數(shù)組內(nèi)存會成倍增長。worker 數(shù)量不是越多越好一般設(shè)成 CPU 物理核心數(shù)或者核心數(shù)減一留一個給主進程。pin_memoryTrue的作用是把數(shù)據(jù)固定到頁鎖定內(nèi)存這樣從 CPU 傳到 GPU 時可以用更快的 DMA 通道。這個參數(shù)只在 GPU 訓(xùn)練時有意義CPU 訓(xùn)練時設(shè)了也沒用。實測下來開啟后小批量數(shù)據(jù)傳輸能快 10% 到 30%批量越大收益越明顯。還有一個容易被忽略的參數(shù)是persistent_workers。當(dāng)num_workers 0時每個 epoch 結(jié)束后 worker 會被銷毀再重建如果數(shù)據(jù)集初始化很慢這個開銷很可觀。設(shè)persistent_workersTrue可以讓 worker 在 epoch 之間保持存活省去重建時間。3.4 collate_fn 的定制時機默認的 collate_fn 會把一個批次里的樣本堆疊成張量。但如果你的樣本長度不一比如文本序列或變長音頻默認行為就會報錯。這時候需要自定義 collate_fn 來做填充。from torch.nn.utils.rnn import pad_sequence def collate_fn(batch): sequences, labels zip(*batch) padded pad_sequence(sequences, batch_firstTrue, padding_value0) labels torch.tensor(labels) return padded, labels這個函數(shù)接收一個列表列表里是__getitem__返回的樣本。你需要自己決定怎么把變長數(shù)據(jù)對齊。注意填充值的選擇要和后續(xù)的 mask 邏輯配合否則填充部分會影響模型計算。4. 完整實操流程與關(guān)鍵環(huán)節(jié)4.1 從原始數(shù)據(jù)到可訓(xùn)練管道的完整步驟假設(shè)我們有一個圖像分類任務(wù)數(shù)據(jù)是硬盤上的圖片文件和對應(yīng)的標(biāo)簽。完整流程分五步第一步整理數(shù)據(jù)索引。把文件路徑和標(biāo)簽讀成兩個列表確保一一對應(yīng)。這一步可以在__init__里做也可以提前存成 CSV 再讀。第二步定義 Dataset。繼承 Dataset實現(xiàn)__len__和__getitem__。在__getitem__里讀圖、做基礎(chǔ)變換、返回張量和標(biāo)簽。from PIL import Image from torchvision import transforms class ImageDataset(Dataset): def __init__(self, paths, labels, transformNone): self.paths paths self.labels labels self.transform transform def __len__(self): return len(self.paths) def __getitem__(self, idx): img Image.open(self.paths[idx]).convert(RGB) if self.transform: img self.transform(img) return img, self.labels[idx]第三步定義變換。訓(xùn)練集和驗證集的變換通常不同。訓(xùn)練集加隨機裁剪、翻轉(zhuǎn)驗證集只做中心裁剪和歸一化。train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])第四步構(gòu)建 DataLoader。訓(xùn)練集打亂驗證集不打亂。訓(xùn)練集開多進程和 pin_memory驗證集可以適當(dāng)降低 worker 數(shù)。train_loader DataLoader( ImageDataset(train_paths, train_labels, train_transform), batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( ImageDataset(val_paths, val_labels, val_transform), batch_size32, shuffleFalse, num_workers2, pin_memoryTrue )第五步寫訓(xùn)練循環(huán)。循環(huán)里只做模型相關(guān)操作數(shù)據(jù)加載交給 DataLoader。for epoch in range(num_epochs): model.train() for images, labels in train_loader: images images.cuda(non_blockingTrue) labels labels.cuda(non_blockingTrue) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step()注意non_blockingTrue配合pin_memoryTrue使用能讓數(shù)據(jù)傳輸異步進行進一步壓榨性能。4.2 參數(shù)計算batch_size 與 num_workers 怎么定batch_size的選擇要平衡顯存和收斂。顯存夠的話從 32 或 64 起步觀察 GPU 利用率和 loss 曲線。如果 loss 震蕩厲害適當(dāng)減小如果訓(xùn)練太慢且顯存有余適當(dāng)增大。經(jīng)驗公式是batch_size 乘以單樣本顯存占用不超過總顯存的 80%。num_workers的計算更直接看 CPU 核心數(shù)。用os.cpu_count()拿到邏輯核心數(shù)然后取一半到全部之間。比如 8 核機器設(shè) 4 到 8 都合理。設(shè)完后用nvidia-smi看 GPU 利用率如果利用率上不去且 CPU 沒跑滿可以再加如果 CPU 跑滿但 GPU 還是等說明數(shù)據(jù)加載是瓶頸需要優(yōu)化__getitem__或加緩存。4.3 實操現(xiàn)場一個文本分類的完整例子文本任務(wù)和圖像任務(wù)在數(shù)據(jù)管道上略有不同主要是變長序列的處理。假設(shè)我們有一個情感分類任務(wù)每條樣本是一段文本和一個標(biāo)簽。import torch from torch.utils.data import Dataset, DataLoader from torch.nn.utils.rnn import pad_sequence class TextDataset(Dataset): def __init__(self, texts, labels, vocab, max_len128): self.texts texts self.labels labels self.vocab vocab self.max_len max_len def __len__(self): return len(self.texts) def __getitem__(self, idx): tokens self.texts[idx].split()[:self.max_len] ids [self.vocab.get(t, self.vocab[unk]) for t in tokens] return torch.tensor(ids, dtypetorch.long), self.labels[idx] def collate_fn(batch): sequences, labels zip(*batch) padded pad_sequence(sequences, batch_firstTrue, padding_value0) labels torch.tensor(labels, dtypetorch.long) return padded, labels dataset TextDataset(texts, labels, vocab) loader DataLoader(dataset, batch_size16, shuffleTrue, num_workers2, collate_fncollate_fn)這個例子里__getitem__返回變長張量collate_fn 負責(zé)填充。填充值 0 對應(yīng)詞表中的 padding 符號后續(xù)模型里要用 mask 把填充位置屏蔽掉。4.4 驗證數(shù)據(jù)管道是否正確的三個檢查點寫完管道后別急著開訓(xùn)。先做三個檢查檢查一取一個批次看形狀。打印images.shape和labels.shape確認維度符合模型輸入要求。圖像通常是[B, C, H, W]文本通常是[B, L]。檢查二看標(biāo)簽分布。用torch.bincount(labels)看一個批次里的類別分布確認沒有全是一類的情況。如果打亂后還是偏斜嚴重說明數(shù)據(jù)本身不平衡需要采樣策略。檢查三測加載速度。用time.time()包住一個 epoch 的迭代看純數(shù)據(jù)加載耗時。如果加載時間超過模型計算時間說明管道是瓶頸需要優(yōu)化。5. 常見問題與排查技巧實錄5.1 多進程加載報錯與卡死Windows 上最常見的報錯是RuntimeError: An attempt has been made to start a new process before the current process has finished its bootstrapping phase。解決辦法是把訓(xùn)練入口放在if __name__ __main__:下面。這個坑我踩過不止一次尤其是在 Jupyter Notebook 里直接跑多進程代碼幾乎必掛。另一個問題是 worker 卡死。如果__getitem__里有死循環(huán)或者等待鎖的操作worker 會一直掛起主進程也跟著等。排查方法是把num_workers設(shè)成 0看是否正常。如果正常再逐個排查__getitem__里的阻塞操作。5.2 內(nèi)存暴漲的幾種原因內(nèi)存暴漲通常有三個來源Dataset 里存了大數(shù)組每個 worker 復(fù)制一份。解決辦法是把大數(shù)據(jù)放到共享內(nèi)存或者用num_workers0。緩存無上限。前面提到的緩存模式如果不加限制會把整個數(shù)據(jù)集都塞進內(nèi)存。可以用functools.lru_cache限制緩存大小。pin_memory 配合大 batch。頁鎖定內(nèi)存不會被換出占用的是物理內(nèi)存。如果 batch 很大pin memory 會吃掉不少內(nèi)存?？梢赃m當(dāng)減小 batch 或關(guān)閉 pin_memory。5.3 數(shù)據(jù)加載慢的優(yōu)化路徑數(shù)據(jù)加載慢的排查順序是先看__getitem__里有沒有重操作再看num_workers是否合理最后看磁盤 IO 是否瓶頸。如果__getitem__里有圖像解碼考慮提前把圖片轉(zhuǎn)成二進制格式如 LMDB 或 HDF5減少小文件讀取開銷。如果磁盤是機械硬盤隨機讀大量小文件會很慢換成 SSD 或者把數(shù)據(jù)打包成大文件能顯著改善。還有一個技巧是預(yù)取。DataLoader 本身有預(yù)取機制但可以通過prefetch_factor參數(shù)調(diào)整每個 worker 預(yù)取的批次數(shù)。默認是 2設(shè)大一點能更好地掩蓋 IO 延遲但會占用更多內(nèi)存。5.4 常見問題速查表問題現(xiàn)象可能原因解決辦法訓(xùn)練時 loss 為 nan標(biāo)簽 dtype 不對標(biāo)簽轉(zhuǎn)成 torch.longGPU 利用率低num_workers 太小或為 0增大 num_workers開啟 pin_memory內(nèi)存持續(xù)增長緩存無上限或 worker 復(fù)制大數(shù)組限制緩存大小減少 worker 數(shù)Windows 多進程報錯缺少 main 保護訓(xùn)練代碼放入if __name__ __main__:變長序列報錯默認 collate 無法堆疊自定義 collate_fn 做填充最后一批報錯batch 只剩一條設(shè)置 drop_lastTrue驗證集結(jié)果抖動驗證集開了 shuffle驗證集 shuffleFalse5.5 幾個我踩過的坑和對應(yīng)技巧第一個坑是在 Dataset 里做歸一化。歸一化應(yīng)該用transforms.Normalize或者提前算好均值和方差不要在__getitem__里現(xiàn)算?，F(xiàn)算不僅慢而且多進程下每個 worker 算的結(jié)果可能不一致。第二個坑是忘了設(shè) shuffle。訓(xùn)練集不 shuffle 的話模型會學(xué)到樣本順序的規(guī)律泛化能力大打折扣。我見過有人訓(xùn)練集和驗證集都忘了 shuffle結(jié)果驗證集準(zhǔn)確率虛高上線后一塌糊涂。第三個坑是num_workers 設(shè)太大。有次我設(shè)了 16 個 worker結(jié)果 CPU 頻繁切換進程加載速度反而比 4 個 worker 還慢。后來學(xué)乖了從 4 開始試逐步加找到性價比最高的點。第四個坑是TensorDataset 里混了不同 dtype。特征 float32、標(biāo)簽 long 是標(biāo)配但有人把標(biāo)簽也弄成 float32訓(xùn)練時交叉熵直接報錯。記住分類任務(wù)的標(biāo)簽必須是 long。6. 進階用法與性能壓榨6.1 用 Sampler 控制采樣策略DataLoader 的shuffle參數(shù)本質(zhì)上是在用 RandomSampler。如果你需要更復(fù)雜的采樣比如類別平衡采樣或分層采樣可以自定義 Sampler。from torch.utils.data import WeightedRandomSampler class_counts [100, 900] weights [1.0 / c for c in class_counts] sample_weights [weights[label] for label in labels] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights)) loader DataLoader(dataset, batch_size32, samplersampler)注意用了 sampler 就不能再設(shè) shuffleTrue兩者互斥。這個技巧在不平衡數(shù)據(jù)集上非常有用能讓模型見到更多少數(shù)類樣本。6.2 IterableDataset 處理流式數(shù)據(jù)如果數(shù)據(jù)是流式的比如從數(shù)據(jù)庫游標(biāo)或日志文件實時讀取用 map-style 的 Dataset 就不合適了。這時候可以用 IterableDataset它實現(xiàn)的是__iter__而不是__getitem__。from torch.utils.data import IterableDataset class StreamDataset(IterableDataset): def __init__(self, file_path): self.file_path file_path def __iter__(self): with open(self.file_path) as f: for line in f: yield parse_line(line)IterableDataset 在多進程下需要小心因為每個 worker 會獨立迭代可能導(dǎo)致數(shù)據(jù)重復(fù)或遺漏。通常需要配合worker_init_fn來劃分每個 worker 負責(zé)的數(shù)據(jù)分片。6.3 把數(shù)據(jù)管道和訓(xùn)練循環(huán)解耦一個我強烈推薦的做法是把數(shù)據(jù)管道的構(gòu)建封裝成函數(shù)訓(xùn)練腳本只調(diào)用函數(shù)拿 loader。這樣數(shù)據(jù)邏輯和模型邏輯分開換數(shù)據(jù)集時不用動訓(xùn)練代碼。def build_loaders(config): train_dataset ImageDataset(config.train_paths, config.train_labels, train_transform) val_dataset ImageDataset(config.val_paths, config.val_labels, val_transform) train_loader DataLoader(train_dataset, batch_sizeconfig.batch_size, shuffleTrue, num_workersconfig.num_workers, pin_memoryTrue, drop_lastTrue) val_loader DataLoader(val_dataset, batch_sizeconfig.batch_size, shuffleFalse, num_workersconfig.num_workers, pin_memoryTrue) return train_loader, val_loader這個模式在項目變大后優(yōu)勢明顯配置改一改就能跑不同的實驗不用在訓(xùn)練腳本里到處找數(shù)據(jù)加載的代碼。6.4 性能對比不同配置下的加載速度我在一臺 8 核 CPU、單卡 GPU 的機器上做過一組對比測試數(shù)據(jù)集是 5 萬張 224x224 的圖片batch_size 固定 32只改 num_workers 和 pin_memorynum_workerspin_memory每 epoch 加載耗時GPU 利用率0False85 秒45%4False32 秒78%4True28 秒85%8True26 秒88%16True31 秒82%從表里能看出num_workers 從 0 到 4 提升最大pin_memory 開啟后又有明顯改善但 worker 加到 16 反而因為進程開銷變慢了。所以調(diào)參要適可而止不是越大越好。6.5 一個容易被忽略的細節(jié)隨機種子多進程加載時每個 worker 的隨機狀態(tài)是獨立的。如果你在__getitem__里做隨機增強不同 worker 可能產(chǎn)生相同的隨機序列導(dǎo)致增強效果打折扣。解決辦法是在worker_init_fn里給每個 worker 設(shè)不同的種子。def worker_init_fn(worker_id): np.random.seed(torch.initial_seed() % 2**32 worker_id) loader DataLoader(dataset, num_workers4, worker_init_fnworker_init_fn)這個細節(jié)在論文復(fù)現(xiàn)時特別重要因為增強的隨機性直接影響最終精度。我吃過這個虧同樣的代碼兩次跑出不同的結(jié)果查了很久才發(fā)現(xiàn)是 worker 種子的問題。7. 我個人在實際操作中的體會寫了這么多最后分享幾點我自己的真實感受。Dataset、TensorDataset、DataLoader 這三個東西看起來簡單但真正用好需要理解它們各自的責(zé)任邊界。我的經(jīng)驗是Dataset 只負責(zé)取一條樣本DataLoader 只負責(zé)組裝批次中間不要越界。一旦你在 Dataset 里做 batch 操作或者在 DataLoader 里做數(shù)據(jù)變換代碼就會變得難以維護。TensorDataset 是個好東西但它只適合數(shù)據(jù)已經(jīng)規(guī)整、不需要動態(tài)變換的場景。一旦涉及數(shù)據(jù)增強或變長序列老老實實寫自定義 Dataset 和 collate_fn別想著偷懶。還有一點數(shù)據(jù)管道的性能優(yōu)化要放在項目早期做。我見過太多項目前期用num_workers0跑通就行后期數(shù)據(jù)量上來后訓(xùn)練慢得無法忍受再回頭改管道牽一發(fā)動全身。所以從一開始就把 num_workers、pin_memory、collate_fn 這些配好后面省心得多。如果你現(xiàn)在正在寫數(shù)據(jù)管道建議先拿一個小數(shù)據(jù)集把整個流程跑通確認形狀、dtype、批次組裝都正確再換大數(shù)據(jù)集。這樣排查問題容易得多不會一上來就被各種報錯淹沒。