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