據(jù)加載實(shí)戰(zhàn)：torchvision數(shù)據(jù)集與DataLoader高效使用指南)
1. 為什么值得把torchvision的數(shù)據(jù)集和DataLoader吃透剛接觸深度學(xué)習(xí)那會(huì)兒我最常干的事就是自己寫(xiě)一個(gè)Dataset類(lèi)然后手動(dòng)切batch、手動(dòng)打亂、手動(dòng)拼tensor代碼又長(zhǎng)又容易出錯(cuò)。后來(lái)被一位做視覺(jué)方向的朋友點(diǎn)醒torchvision里現(xiàn)成的數(shù)據(jù)集和DataLoader這套組合其實(shí)已經(jīng)把數(shù)據(jù)加載這條鏈路上80%的臟活累活都干完了你只需要搞清楚它怎么用、什么時(shí)候該改、改哪里。這句話我記到現(xiàn)在。torchvision數(shù)據(jù)集以及dataloader使用說(shuō)白了就是解決一個(gè)核心問(wèn)題如何把磁盤(pán)上的圖片、標(biāo)簽高效、穩(wěn)定、可復(fù)現(xiàn)地喂給模型訓(xùn)練循環(huán)。它適合所有用PyTorch做計(jì)算機(jī)視覺(jué)任務(wù)的人不管你是剛跑通第一個(gè)MNIST的新手還是已經(jīng)在做多機(jī)多卡訓(xùn)練的老手這套東西你每天都在碰。新手容易在num_workers、transform順序、collate_fn上翻車(chē)?yán)鲜謩t更關(guān)心采樣策略、內(nèi)存占用和跨進(jìn)程的坑。這篇文章我打算從設(shè)計(jì)思路講到實(shí)操細(xì)節(jié)把我在實(shí)際項(xiàng)目里踩過(guò)的坑和總結(jié)的技巧都攤開(kāi)說(shuō)盡量讓你看完就能直接抄作業(yè)。torchvision的數(shù)據(jù)集模塊本質(zhì)上提供了兩層?xùn)|西一層是現(xiàn)成的數(shù)據(jù)集類(lèi)比如CIFAR10、ImageFolder、COCO幫你處理下載、解壓、路徑索引、標(biāo)簽映射這些瑣事另一層是通用的數(shù)據(jù)變換和加載基礎(chǔ)設(shè)施配合torch.utils.data.DataLoader完成批處理、打亂、并行加載。理解這兩層的分工是后面所有操作的前提。很多人一上來(lái)就急著改源碼其實(shí)大部分需求用現(xiàn)成的組合就能滿足改之前先想清楚是不是真的需要?jiǎng)拥讓印?. 整體設(shè)計(jì)思路拆解Dataset和DataLoader到底怎么分工2.1 一個(gè)負(fù)責(zé)“取一條”一個(gè)負(fù)責(zé)“湊一批”我習(xí)慣用餐廳點(diǎn)菜來(lái)類(lèi)比。Dataset是后廚的備菜員你告訴它“我要第i道菜”它負(fù)責(zé)把這道菜一張圖一個(gè)標(biāo)簽準(zhǔn)備好遞給你至于這道菜怎么切、怎么擺盤(pán)transform它按你給的規(guī)矩來(lái)。DataLoader是傳菜的服務(wù)員它負(fù)責(zé)一次端走一桌的菜一個(gè)batch決定上菜順序shuffle、同時(shí)開(kāi)幾個(gè)窗口出菜num_workers、一桌坐幾個(gè)人batch_size。這個(gè)分工非常清晰Dataset只關(guān)心單條樣本的正確性DataLoader只關(guān)心批量組織的效率。這個(gè)設(shè)計(jì)帶來(lái)的直接好處是解耦。你想換數(shù)據(jù)集只需要換Dataset實(shí)現(xiàn)DataLoader那套參數(shù)原封不動(dòng)你想調(diào)batch策略只動(dòng)DataLoader數(shù)據(jù)集代碼一行不改。我在實(shí)際項(xiàng)目里經(jīng)常遇到這種情況同一個(gè)ImageFolder訓(xùn)練時(shí)用shuffleTrue、batch_size64驗(yàn)證時(shí)用shuffleFalse、batch_size128數(shù)據(jù)集對(duì)象復(fù)用只換DataLoader配置干凈利落。2.2 為什么要有__getitem__和__len__這兩個(gè)魔法方法自定義Dataset時(shí)你幾乎只需要實(shí)現(xiàn)兩個(gè)方法__len__返回樣本總數(shù)__getitem__根據(jù)索引返回一條樣本。PyTorch的采樣器Sampler依賴(lài)__len__來(lái)決定索引范圍DataLoader依賴(lài)__getitem__來(lái)真正取數(shù)據(jù)。這個(gè)約定看起來(lái)簡(jiǎn)單但它是整個(gè)并行加載機(jī)制的基礎(chǔ)——因?yàn)橛辛嗣鞔_的索引接口多個(gè)worker進(jìn)程才能各自獨(dú)立地按索引去取數(shù)據(jù)而不需要共享任何狀態(tài)。這里有個(gè)容易被忽略的點(diǎn)__getitem__里做的事情越少越好。我見(jiàn)過(guò)有人在里面做特別重的預(yù)處理比如整圖解碼加復(fù)雜增強(qiáng)結(jié)果num_workers開(kāi)大了反而更慢因?yàn)槊總€(gè)worker都在重復(fù)做重活。合理的做法是把輕量的索引、路徑拼接放在__getitem__把重的變換交給transform管道讓torchvision的變換在worker進(jìn)程里并行執(zhí)行。這個(gè)邊界劃清楚性能差別很明顯。2.3 采樣器、批處理、collate三件套的協(xié)作關(guān)系DataLoader內(nèi)部其實(shí)串了三個(gè)環(huán)節(jié)Sampler決定這一輪按什么順序取哪些索引BatchSampler把這些索引切成一個(gè)個(gè)batchcollate_fn把一個(gè)batch的樣本列表拼成一個(gè)tensor。默認(rèn)情況下shuffleTrue會(huì)用RandomSamplerbatch_size決定切分collate_fn用默認(rèn)的default_collate把list of tuple拼成tuple of tensor。理解這條鏈路你才知道該在哪個(gè)環(huán)節(jié)動(dòng)手。比如做目標(biāo)檢測(cè)時(shí)每張圖的標(biāo)注框數(shù)量不一樣默認(rèn)collate會(huì)報(bào)錯(cuò)因?yàn)閠ensor形狀對(duì)不齊。這時(shí)候你要么自定義collate_fn返回list而不是stack要么用torchvision檢測(cè)數(shù)據(jù)集自帶的collate。再比如做對(duì)比學(xué)習(xí)一個(gè)batch需要成對(duì)的樣本你就得自定義BatchSampler保證正負(fù)樣本配對(duì)出現(xiàn)。這些需求都不是改Dataset能解決的必須落到采樣或collate層。我一開(kāi)始總想在一個(gè)地方解決所有問(wèn)題后來(lái)發(fā)現(xiàn)按這條鏈路分層去改代碼清晰得多。3. 核心細(xì)節(jié)解析與實(shí)操要點(diǎn)從現(xiàn)成數(shù)據(jù)集到自定義Dataset3.1 用現(xiàn)成數(shù)據(jù)集快速起步以CIFAR10和ImageFolder為例torchvision內(nèi)置數(shù)據(jù)集最大的價(jià)值是開(kāi)箱即用且可復(fù)現(xiàn)。以CIFAR10為例你只需要指定root、train、download、transform四個(gè)參數(shù)它自動(dòng)幫你下載、校驗(yàn)、解壓、建立索引。第一次運(yùn)行會(huì)下載到本地之后直接讀緩存。這里的關(guān)鍵是transform參數(shù)它決定了每條樣本被取出來(lái)時(shí)做什么變換。我通常把訓(xùn)練和驗(yàn)證的transform分開(kāi)定義訓(xùn)練加隨機(jī)裁剪和翻轉(zhuǎn)驗(yàn)證只做resize和歸一化這是保證評(píng)估穩(wěn)定的基本操作。import torchvision from torchvision import transforms train_tf 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_tf 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]), ]) train_set torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_tf) val_set torchvision.datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformval_tf)ImageFolder則是處理自有圖片數(shù)據(jù)的主力。它要求目錄結(jié)構(gòu)是root/類(lèi)別名/圖片文件自動(dòng)把子目錄名映射成類(lèi)別索引。這個(gè)約定非常省事你只要把數(shù)據(jù)按類(lèi)別分好文件夾剩下的索引、標(biāo)簽映射它全包了。我實(shí)際項(xiàng)目里90%的分類(lèi)任務(wù)都是先用ImageFolder跑通baseline再考慮要不要換更復(fù)雜的自定義數(shù)據(jù)集。注意ImageFolder默認(rèn)按文件夾名的字母序分配類(lèi)別索引如果你對(duì)類(lèi)別順序有要求得自己傳class_to_idx或者事后核對(duì)。3.2 transform管道的順序陷阱為什么ToTensor必須靠后這是新手最容易翻車(chē)的地方之一。transforms.Compose是按列表順序依次執(zhí)行的而很多變換對(duì)輸入類(lèi)型有要求。ToTensor()把PIL Image或numpy數(shù)組轉(zhuǎn)成[0,1]范圍的float tensor并調(diào)整維度順序?yàn)镃,H,WNormalize()則要求輸入已經(jīng)是tensor。所以正確順序永遠(yuǎn)是幾何/顏色變換PIL階段→ ToTensor → Normalize。我見(jiàn)過(guò)有人把Normalize寫(xiě)在ToTensor前面直接報(bào)類(lèi)型錯(cuò)誤排查半天。還有一個(gè)隱蔽的坑RandomHorizontalFlip這類(lèi)隨機(jī)變換在訓(xùn)練時(shí)每次__getitem__都會(huì)重新采樣這意味著同一個(gè)索引在不同epoch取到的樣本可能不同這正是數(shù)據(jù)增強(qiáng)想要的效果。但如果你在驗(yàn)證集上誤用了隨機(jī)變換評(píng)估結(jié)果就會(huì)抖動(dòng)看起來(lái)像模型不穩(wěn)定其實(shí)是數(shù)據(jù)在變。我的習(xí)慣是訓(xùn)練和驗(yàn)證的transform嚴(yán)格分開(kāi)定義驗(yàn)證集只用確定性的resize、crop、normalize絕不摻隨機(jī)操作。3.3 自定義Dataset的正確姿勢(shì)繼承、實(shí)現(xiàn)、別亂加狀態(tài)當(dāng)現(xiàn)成數(shù)據(jù)集不滿足需求時(shí)比如多模態(tài)輸入、特殊標(biāo)注格式就得自己寫(xiě)Dataset。標(biāo)準(zhǔn)寫(xiě)法是繼承torch.utils.data.Dataset在__init__里把樣本路徑列表、標(biāo)簽列表、transform都存成實(shí)例屬性__len__返回列表長(zhǎng)度__getitem__按索引取路徑、讀圖、應(yīng)用transform、返回樣本和標(biāo)簽。這里有個(gè)重要原則__init__里不要做重活比如不要在里面把所有圖片都讀進(jìn)內(nèi)存除非你確實(shí)需要且內(nèi)存扛得住。把讀圖放到__getitem__讓worker進(jìn)程并行去讀才是正確姿勢(shì)。from torch.utils.data import Dataset from PIL import Image class MyDataset(Dataset): def __init__(self, samples, transformNone): self.samples samples # [(img_path, label), ...] self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label注意Image.open是惰性的真正解碼發(fā)生在讀取像素時(shí)所以convert(RGB)這一步很重要能統(tǒng)一處理灰度圖、帶alpha通道的PNG避免后面transform報(bào)通道數(shù)不匹配。這個(gè)細(xì)節(jié)我在處理混合來(lái)源圖片時(shí)吃過(guò)虧加上convert之后世界清凈了。3.4 num_workers、pin_memory、persistent_workers這幾個(gè)參數(shù)怎么定DataLoader的參數(shù)里對(duì)性能影響最大的就是num_workers。它決定開(kāi)幾個(gè)子進(jìn)程并行取數(shù)據(jù)。設(shè)成0表示在主進(jìn)程里取簡(jiǎn)單但會(huì)阻塞訓(xùn)練設(shè)成正數(shù)則開(kāi)子進(jìn)程。經(jīng)驗(yàn)值是不超過(guò)CPU物理核心數(shù)一般從4或8起步邊跑邊看GPU利用率。如果GPU利用率上不去、經(jīng)常等數(shù)據(jù)就加worker如果CPU已經(jīng)打滿、內(nèi)存暴漲就減worker。盲目設(shè)很大反而會(huì)因?yàn)檫M(jìn)程調(diào)度和內(nèi)存拷貝拖慢整體。pin_memoryTrue會(huì)把batch數(shù)據(jù)放到鎖頁(yè)內(nèi)存加速?gòu)腃PU到GPU的拷貝用GPU訓(xùn)練時(shí)基本必開(kāi)。persistent_workersTrue讓worker進(jìn)程在epoch之間不銷(xiāo)毀重建減少反復(fù)啟動(dòng)的開(kāi)銷(xiāo)數(shù)據(jù)量大、epoch多的時(shí)候收益明顯。這幾個(gè)參數(shù)組合起來(lái)我常用的訓(xùn)練配置是num_workers8, pin_memoryTrue, persistent_workersTrue配合batch_size64左右在單卡上跑得比較順。注意Windows上num_workers0有時(shí)會(huì)遇到多進(jìn)程啟動(dòng)問(wèn)題如果報(bào)奇怪的pickle錯(cuò)誤先把num_workers設(shè)成0確認(rèn)邏輯沒(méi)問(wèn)題再逐步調(diào)大排查。4. 實(shí)操過(guò)程與核心環(huán)節(jié)實(shí)現(xiàn)搭一條能跑的訓(xùn)練數(shù)據(jù)流4.1 從零搭一條完整的數(shù)據(jù)管道我把完整流程拆成五步定義transform、構(gòu)建Dataset、包DataLoader、寫(xiě)訓(xùn)練循環(huán)取batch、驗(yàn)證集單獨(dú)配置。下面這段是可以直接跑的骨架以ImageFolder為例。from torch.utils.data import DataLoader import torchvision train_set torchvision.datasets.ImageFolder(./data/train, transformtrain_tf) val_set torchvision.datasets.ImageFolder(./data/val, transformval_tf) train_loader DataLoader(train_set, batch_size64, shuffleTrue, num_workers8, pin_memoryTrue, persistent_workersTrue, drop_lastTrue) val_loader DataLoader(val_set, batch_size128, shuffleFalse, num_workers8, pin_memoryTrue) for epoch in range(num_epochs): for images, labels in train_loader: images images.cuda(non_blockingTrue) labels labels.cuda(non_blockingTrue) # 前向、反向、更新這里drop_lastTrue是為了保證每個(gè)batch大小一致避免最后一個(gè)不滿的batch影響B(tài)atchNorm的統(tǒng)計(jì)訓(xùn)練時(shí)我一般開(kāi)著驗(yàn)證集則保留全部樣本保證評(píng)估覆蓋完整。non_blockingTrue配合pin_memory能讓拷貝和計(jì)算重疊是個(gè)小但有效的優(yōu)化。4.2 batch_size和num_workers的聯(lián)合調(diào)參過(guò)程這兩個(gè)參數(shù)不是獨(dú)立的。我通常的做法是先固定num_workers4把batch_size從32往上試觀察顯存占用和單步耗時(shí)找到顯存放得下的最大值附近然后固定這個(gè)batch_size把num_workers從4加到8、16看GPU利用率和吞吐變化。如果加worker后吞吐不再提升甚至下降說(shuō)明數(shù)據(jù)加載已經(jīng)不是瓶頸停在這個(gè)值就好。舉個(gè)我實(shí)際調(diào)過(guò)的例子某次訓(xùn)練單步耗時(shí)約120msGPU利用率只有50%左右明顯在等數(shù)據(jù)。把num_workers從4提到8單步降到90ms利用率升到75%再提到16單步反而回到100ms因?yàn)檫M(jìn)程切換和內(nèi)存拷貝開(kāi)銷(xiāo)上來(lái)了。最后定在8。這個(gè)過(guò)程沒(méi)有公式只能實(shí)測(cè)但方向很明確看GPU是不是在等等就加worker加到頭就停。4.3 自定義collate_fn處理變長(zhǎng)樣本做檢測(cè)或序列任務(wù)時(shí)一個(gè)batch里樣本形狀不一致默認(rèn)collate會(huì)失敗。解決辦法是傳自定義collate_fn把樣本列表原樣打包成list返回讓模型內(nèi)部自己處理對(duì)齊。def my_collate(batch): images [item[0] for item in batch] targets [item[1] for item in batch] return images, targets loader DataLoader(dataset, batch_size8, collate_fnmy_collate)代價(jià)是batch不再是規(guī)整的tensor后續(xù)要自己stack或pad。我的經(jīng)驗(yàn)是能pad成統(tǒng)一形狀就盡量pad因?yàn)橐?guī)整tensor對(duì)GPU更友好實(shí)在不行再用list但要在模型里第一時(shí)間處理掉別讓不規(guī)則結(jié)構(gòu)在計(jì)算圖里傳太深。4.4 用Subset和ConcatDataset做數(shù)據(jù)集切分與合并實(shí)際項(xiàng)目里經(jīng)常需要從一個(gè)大訓(xùn)練集里切出小份做快速實(shí)驗(yàn)或者把多個(gè)來(lái)源的數(shù)據(jù)合并訓(xùn)練。torch.utils.data.Subset按索引列表取子集ConcatDataset把多個(gè)數(shù)據(jù)集首尾拼接。這兩個(gè)工具配合random_split能覆蓋大部分切分需求。from torch.utils.data import random_split, ConcatDataset train_part, val_part random_split(full_set, [0.9, 0.1], generatortorch.Generator().manual_seed(42)) combined ConcatDataset([set_a, set_b])random_split一定要傳固定種子的generator否則每次切分結(jié)果不同實(shí)驗(yàn)沒(méi)法復(fù)現(xiàn)。這個(gè)坑我踩過(guò)不止一次后來(lái)養(yǎng)成習(xí)慣凡是涉及隨機(jī)切分的地方都顯式給種子。5. 常見(jiàn)問(wèn)題與排查技巧實(shí)錄5.1 常見(jiàn)報(bào)錯(cuò)速查表現(xiàn)象可能原因排查與解決TypeError: default_collate: batch must contain tensors樣本返回類(lèi)型不統(tǒng)一或含非tensor檢查_(kāi)_getitem__返回值統(tǒng)一轉(zhuǎn)tensor或自定義collateRuntimeError: DataLoader worker exited unexpectedlyworker內(nèi)異常被吞、內(nèi)存不足先設(shè)num_workers0復(fù)現(xiàn)真實(shí)報(bào)錯(cuò)再查內(nèi)存訓(xùn)練loss劇烈抖動(dòng)驗(yàn)證集誤用隨機(jī)transform檢查val的transform是否含隨機(jī)操作GPU利用率低、訓(xùn)練慢數(shù)據(jù)加載成瓶頸增大num_workers、開(kāi)pin_memory顯存OOM但batch不大圖片分辨率過(guò)高或transform緩存降分辨率、檢查是否在__init__里緩存了整圖多進(jìn)程下卡死Windows多進(jìn)程或共享狀態(tài)問(wèn)題設(shè)num_workers0驗(yàn)證避免在Dataset里存不可pickle對(duì)象5.2 幾個(gè)我踩過(guò)的坑和獨(dú)家心得第一個(gè)坑是在Dataset里存了打開(kāi)的文件句柄或數(shù)據(jù)庫(kù)連接。這些對(duì)象不可picklenum_workers0時(shí)子進(jìn)程無(wú)法繼承直接報(bào)錯(cuò)。正確做法是在__getitem__里按需打開(kāi)、用完即關(guān)或者用worker_init_fn在每個(gè)worker里單獨(dú)初始化。第二個(gè)坑是transform里用了帶狀態(tài)的隨機(jī)數(shù)生成器。多worker下每個(gè)進(jìn)程的隨機(jī)種子默認(rèn)不同這本身沒(méi)問(wèn)題但如果你依賴(lài)某個(gè)全局隨機(jī)狀態(tài)做可復(fù)現(xiàn)增強(qiáng)就會(huì)亂。我的做法是增強(qiáng)全部用torchvision自帶的隨機(jī)變換它們內(nèi)部處理好了種子不自己造輪子。第三個(gè)心得是先用小數(shù)據(jù)跑通再上全量。我習(xí)慣先取100條樣本、num_workers0跑兩個(gè)epoch確認(rèn)數(shù)據(jù)流、transform、collate全對(duì)再切到全量和多worker。這樣出問(wèn)題時(shí)排查范圍小得多比一上來(lái)就全量多進(jìn)程調(diào)試省時(shí)間。第四個(gè)心得是驗(yàn)證集的DataLoader不要開(kāi)shuffle也不要drop_last。評(píng)估要覆蓋全部樣本且順序穩(wěn)定否則指標(biāo)沒(méi)法橫向比較。這個(gè)看似基礎(chǔ)但我在趕進(jìn)度時(shí)犯過(guò)導(dǎo)致兩次評(píng)估結(jié)果對(duì)不上白查了半天模型。5.3 內(nèi)存與性能的平衡技巧數(shù)據(jù)量大時(shí)num_workers開(kāi)大容易把內(nèi)存吃滿因?yàn)槊總€(gè)worker都會(huì)緩存一部分?jǐn)?shù)據(jù)。我的做法是監(jiān)控內(nèi)存占用找到吞吐和內(nèi)存的平衡點(diǎn)。另外如果圖片很大而模型輸入很小可以在__getitem__里先做一次快速降采樣再進(jìn)transform能顯著減少解碼和傳輸開(kāi)銷(xiāo)。還有個(gè)小技巧把常用的歸一化參數(shù)、均值方差提前算好寫(xiě)死別在每次transform里現(xiàn)算雖然省得不多但積少成多。6. 采樣策略進(jìn)階讓DataLoader按你的意圖取數(shù)據(jù)6.1 WeightedRandomSampler解決類(lèi)別不平衡分類(lèi)任務(wù)里類(lèi)別不平衡很常見(jiàn)少數(shù)類(lèi)樣本少模型容易偏向多數(shù)類(lèi)。WeightedRandomSampler給每個(gè)樣本一個(gè)采樣權(quán)重讓少數(shù)類(lèi)被抽到的概率更高。權(quán)重通常設(shè)成類(lèi)別頻率的倒數(shù)按樣本所屬類(lèi)別賦權(quán)。from torch.utils.data import WeightedRandomSampler import numpy as np labels [label for _, label in train_set.samples] class_counts np.bincount(labels) class_weights 1.0 / class_counts sample_weights [class_weights[l] for l in labels] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) loader DataLoader(train_set, batch_size64, samplersampler, num_workers8)注意用了sampler就不能再設(shè)shuffleTrue兩者沖突。replacementTrue表示有放回采樣這是加權(quán)采樣的常規(guī)做法。我實(shí)測(cè)下來(lái)這個(gè)改動(dòng)對(duì)少數(shù)類(lèi)的召回提升很明顯代價(jià)是每個(gè)epoch看到的樣本組合不同訓(xùn)練曲線會(huì)稍微抖一點(diǎn)屬于正?，F(xiàn)象。6.2 自定義BatchSampler實(shí)現(xiàn)配對(duì)采樣做對(duì)比學(xué)習(xí)或?qū)\生網(wǎng)絡(luò)時(shí)一個(gè)batch需要成對(duì)的正負(fù)樣本。這時(shí)候自定義BatchSampler最直接它接收一個(gè)基礎(chǔ)sampler按你的規(guī)則把索引組織成batch列表。from torch.utils.data import Sampler class PairBatchSampler(Sampler): def __init__(self, labels, batch_size, seed42): self.labels labels self.batch_size batch_size self.rng np.random.default_rng(seed) def __iter__(self): # 簡(jiǎn)化示意按類(lèi)別分組后配對(duì)取樣 ... def __len__(self): return len(self.labels) // self.batch_size這類(lèi)自定義采樣器邏輯不復(fù)雜但要注意__len__返回的batch數(shù)量要和實(shí)際迭代次數(shù)一致否則訓(xùn)練循環(huán)的進(jìn)度估計(jì)會(huì)錯(cuò)。我一般寫(xiě)完先用小數(shù)據(jù)打印幾個(gè)batch的索引確認(rèn)配對(duì)邏輯對(duì)再上全量。6.3 分布式訓(xùn)練下的DistributedSampler要點(diǎn)多卡訓(xùn)練時(shí)每個(gè)進(jìn)程要用DistributedSampler保證各卡看到不重疊的數(shù)據(jù)分片且每個(gè)epoch重新打亂。關(guān)鍵參數(shù)是num_replicas總進(jìn)程數(shù)和rank當(dāng)前進(jìn)程編號(hào)以及每個(gè)epoch開(kāi)始前調(diào)用sampler.set_epoch(epoch)否則各epoch的分片順序不變等于沒(méi)打亂。這個(gè)set_epoch我漏調(diào)過(guò)一次訓(xùn)練效果莫名變差查了很久才發(fā)現(xiàn)是分片沒(méi)刷新。7. 我個(gè)人的幾條實(shí)操建議數(shù)據(jù)加載這塊我的核心體會(huì)是先把單條樣本的正確性驗(yàn)證到位再談并行和性能。很多人一上來(lái)就調(diào)num_workers、追吞吐結(jié)果底層transform有bug多進(jìn)程下報(bào)錯(cuò)信息還被吞掉排查成本翻倍。我現(xiàn)在的固定流程是小樣本、單進(jìn)程、打印幾條樣本的形狀和標(biāo)簽確認(rèn)無(wú)誤后再逐步放大。另外別迷信默認(rèn)參數(shù)。DataLoader的默認(rèn)值是為了通用性不是為你的任務(wù)優(yōu)化過(guò)的。batch_size、num_workers、pin_memory、drop_last這幾個(gè)值得花半小時(shí)實(shí)測(cè)調(diào)一遍收益比調(diào)模型結(jié)構(gòu)來(lái)得快。最后再分享一個(gè)小習(xí)慣把數(shù)據(jù)管道的配置transform、batch、worker數(shù)和模型配置分開(kāi)寫(xiě)實(shí)驗(yàn)記錄里單獨(dú)存一份復(fù)現(xiàn)的時(shí)候能省很多事。這套東西用熟了你會(huì)發(fā)現(xiàn)數(shù)據(jù)加載不再是訓(xùn)練的瓶頸而是你可以放心交給它的基礎(chǔ)設(shè)施。 PyTorch數(shù)據(jù)加載實(shí)戰(zhàn)：torchvision數(shù)據(jù)集與DataLoader高效使用指南
1. 為什么值得把torchvision的數(shù)據(jù)集和DataLoader吃透剛接觸深度學(xué)習(xí)那會(huì)兒我最常干的事就是自己寫(xiě)一個(gè)Dataset類(lèi)然后手動(dòng)切batch、手動(dòng)打亂、手動(dòng)拼tensor代碼又長(zhǎng)又容易出錯(cuò)。后來(lái)被一位做視覺(jué)方向的朋友點(diǎn)醒torchvision里現(xiàn)成的數(shù)據(jù)集和DataLoader這套組合其實(shí)已經(jīng)把數(shù)據(jù)加載這條鏈路上80%的臟活累活都干完了你只需要搞清楚它怎么用、什么時(shí)候該改、改哪里。這句話我記到現(xiàn)在。torchvision數(shù)據(jù)集以及dataloader使用說(shuō)白了就是解決一個(gè)核心問(wèn)題如何把磁盤(pán)上的圖片、標(biāo)簽高效、穩(wěn)定、可復(fù)現(xiàn)地喂給模型訓(xùn)練循環(huán)。它適合所有用PyTorch做計(jì)算機(jī)視覺(jué)任務(wù)的人不管你是剛跑通第一個(gè)MNIST的新手還是已經(jīng)在做多機(jī)多卡訓(xùn)練的老手這套東西你每天都在碰。新手容易在num_workers、transform順序、collate_fn上翻車(chē)?yán)鲜謩t更關(guān)心采樣策略、內(nèi)存占用和跨進(jìn)程的坑。這篇文章我打算從設(shè)計(jì)思路講到實(shí)操細(xì)節(jié)把我在實(shí)際項(xiàng)目里踩過(guò)的坑和總結(jié)的技巧都攤開(kāi)說(shuō)盡量讓你看完就能直接抄作業(yè)。torchvision的數(shù)據(jù)集模塊本質(zhì)上提供了兩層?xùn)|西一層是現(xiàn)成的數(shù)據(jù)集類(lèi)比如CIFAR10、ImageFolder、COCO幫你處理下載、解壓、路徑索引、標(biāo)簽映射這些瑣事另一層是通用的數(shù)據(jù)變換和加載基礎(chǔ)設(shè)施配合torch.utils.data.DataLoader完成批處理、打亂、并行加載。理解這兩層的分工是后面所有操作的前提。很多人一上來(lái)就急著改源碼其實(shí)大部分需求用現(xiàn)成的組合就能滿足改之前先想清楚是不是真的需要?jiǎng)拥讓印?. 整體設(shè)計(jì)思路拆解Dataset和DataLoader到底怎么分工2.1 一個(gè)負(fù)責(zé)“取一條”一個(gè)負(fù)責(zé)“湊一批”我習(xí)慣用餐廳點(diǎn)菜來(lái)類(lèi)比。Dataset是后廚的備菜員你告訴它“我要第i道菜”它負(fù)責(zé)把這道菜一張圖一個(gè)標(biāo)簽準(zhǔn)備好遞給你至于這道菜怎么切、怎么擺盤(pán)transform它按你給的規(guī)矩來(lái)。DataLoader是傳菜的服務(wù)員它負(fù)責(zé)一次端走一桌的菜一個(gè)batch決定上菜順序shuffle、同時(shí)開(kāi)幾個(gè)窗口出菜num_workers、一桌坐幾個(gè)人batch_size。這個(gè)分工非常清晰Dataset只關(guān)心單條樣本的正確性DataLoader只關(guān)心批量組織的效率。這個(gè)設(shè)計(jì)帶來(lái)的直接好處是解耦。你想換數(shù)據(jù)集只需要換Dataset實(shí)現(xiàn)DataLoader那套參數(shù)原封不動(dòng)你想調(diào)batch策略只動(dòng)DataLoader數(shù)據(jù)集代碼一行不改。我在實(shí)際項(xiàng)目里經(jīng)常遇到這種情況同一個(gè)ImageFolder訓(xùn)練時(shí)用shuffleTrue、batch_size64驗(yàn)證時(shí)用shuffleFalse、batch_size128數(shù)據(jù)集對(duì)象復(fù)用只換DataLoader配置干凈利落。2.2 為什么要有__getitem__和__len__這兩個(gè)魔法方法自定義Dataset時(shí)你幾乎只需要實(shí)現(xiàn)兩個(gè)方法__len__返回樣本總數(shù)__getitem__根據(jù)索引返回一條樣本。PyTorch的采樣器Sampler依賴(lài)__len__來(lái)決定索引范圍DataLoader依賴(lài)__getitem__來(lái)真正取數(shù)據(jù)。這個(gè)約定看起來(lái)簡(jiǎn)單但它是整個(gè)并行加載機(jī)制的基礎(chǔ)——因?yàn)橛辛嗣鞔_的索引接口多個(gè)worker進(jìn)程才能各自獨(dú)立地按索引去取數(shù)據(jù)而不需要共享任何狀態(tài)。這里有個(gè)容易被忽略的點(diǎn)__getitem__里做的事情越少越好。我見(jiàn)過(guò)有人在里面做特別重的預(yù)處理比如整圖解碼加復(fù)雜增強(qiáng)結(jié)果num_workers開(kāi)大了反而更慢因?yàn)槊總€(gè)worker都在重復(fù)做重活。合理的做法是把輕量的索引、路徑拼接放在__getitem__把重的變換交給transform管道讓torchvision的變換在worker進(jìn)程里并行執(zhí)行。這個(gè)邊界劃清楚性能差別很明顯。2.3 采樣器、批處理、collate三件套的協(xié)作關(guān)系DataLoader內(nèi)部其實(shí)串了三個(gè)環(huán)節(jié)Sampler決定這一輪按什么順序取哪些索引BatchSampler把這些索引切成一個(gè)個(gè)batchcollate_fn把一個(gè)batch的樣本列表拼成一個(gè)tensor。默認(rèn)情況下shuffleTrue會(huì)用RandomSamplerbatch_size決定切分collate_fn用默認(rèn)的default_collate把list of tuple拼成tuple of tensor。理解這條鏈路你才知道該在哪個(gè)環(huán)節(jié)動(dòng)手。比如做目標(biāo)檢測(cè)時(shí)每張圖的標(biāo)注框數(shù)量不一樣默認(rèn)collate會(huì)報(bào)錯(cuò)因?yàn)閠ensor形狀對(duì)不齊。這時(shí)候你要么自定義collate_fn返回list而不是stack要么用torchvision檢測(cè)數(shù)據(jù)集自帶的collate。再比如做對(duì)比學(xué)習(xí)一個(gè)batch需要成對(duì)的樣本你就得自定義BatchSampler保證正負(fù)樣本配對(duì)出現(xiàn)。這些需求都不是改Dataset能解決的必須落到采樣或collate層。我一開(kāi)始總想在一個(gè)地方解決所有問(wèn)題后來(lái)發(fā)現(xiàn)按這條鏈路分層去改代碼清晰得多。3. 核心細(xì)節(jié)解析與實(shí)操要點(diǎn)從現(xiàn)成數(shù)據(jù)集到自定義Dataset3.1 用現(xiàn)成數(shù)據(jù)集快速起步以CIFAR10和ImageFolder為例torchvision內(nèi)置數(shù)據(jù)集最大的價(jià)值是開(kāi)箱即用且可復(fù)現(xiàn)。以CIFAR10為例你只需要指定root、train、download、transform四個(gè)參數(shù)它自動(dòng)幫你下載、校驗(yàn)、解壓、建立索引。第一次運(yùn)行會(huì)下載到本地之后直接讀緩存。這里的關(guān)鍵是transform參數(shù)它決定了每條樣本被取出來(lái)時(shí)做什么變換。我通常把訓(xùn)練和驗(yàn)證的transform分開(kāi)定義訓(xùn)練加隨機(jī)裁剪和翻轉(zhuǎn)驗(yàn)證只做resize和歸一化這是保證評(píng)估穩(wěn)定的基本操作。import torchvision from torchvision import transforms train_tf 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_tf 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]), ]) train_set torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_tf) val_set torchvision.datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformval_tf)ImageFolder則是處理自有圖片數(shù)據(jù)的主力。它要求目錄結(jié)構(gòu)是root/類(lèi)別名/圖片文件自動(dòng)把子目錄名映射成類(lèi)別索引。這個(gè)約定非常省事你只要把數(shù)據(jù)按類(lèi)別分好文件夾剩下的索引、標(biāo)簽映射它全包了。我實(shí)際項(xiàng)目里90%的分類(lèi)任務(wù)都是先用ImageFolder跑通baseline再考慮要不要換更復(fù)雜的自定義數(shù)據(jù)集。注意ImageFolder默認(rèn)按文件夾名的字母序分配類(lèi)別索引如果你對(duì)類(lèi)別順序有要求得自己傳class_to_idx或者事后核對(duì)。3.2 transform管道的順序陷阱為什么ToTensor必須靠后這是新手最容易翻車(chē)的地方之一。transforms.Compose是按列表順序依次執(zhí)行的而很多變換對(duì)輸入類(lèi)型有要求。ToTensor()把PIL Image或numpy數(shù)組轉(zhuǎn)成[0,1]范圍的float tensor并調(diào)整維度順序?yàn)镃,H,WNormalize()則要求輸入已經(jīng)是tensor。所以正確順序永遠(yuǎn)是幾何/顏色變換PIL階段→ ToTensor → Normalize。我見(jiàn)過(guò)有人把Normalize寫(xiě)在ToTensor前面直接報(bào)類(lèi)型錯(cuò)誤排查半天。還有一個(gè)隱蔽的坑RandomHorizontalFlip這類(lèi)隨機(jī)變換在訓(xùn)練時(shí)每次__getitem__都會(huì)重新采樣這意味著同一個(gè)索引在不同epoch取到的樣本可能不同這正是數(shù)據(jù)增強(qiáng)想要的效果。但如果你在驗(yàn)證集上誤用了隨機(jī)變換評(píng)估結(jié)果就會(huì)抖動(dòng)看起來(lái)像模型不穩(wěn)定其實(shí)是數(shù)據(jù)在變。我的習(xí)慣是訓(xùn)練和驗(yàn)證的transform嚴(yán)格分開(kāi)定義驗(yàn)證集只用確定性的resize、crop、normalize絕不摻隨機(jī)操作。3.3 自定義Dataset的正確姿勢(shì)繼承、實(shí)現(xiàn)、別亂加狀態(tài)當(dāng)現(xiàn)成數(shù)據(jù)集不滿足需求時(shí)比如多模態(tài)輸入、特殊標(biāo)注格式就得自己寫(xiě)Dataset。標(biāo)準(zhǔn)寫(xiě)法是繼承torch.utils.data.Dataset在__init__里把樣本路徑列表、標(biāo)簽列表、transform都存成實(shí)例屬性__len__返回列表長(zhǎng)度__getitem__按索引取路徑、讀圖、應(yīng)用transform、返回樣本和標(biāo)簽。這里有個(gè)重要原則__init__里不要做重活比如不要在里面把所有圖片都讀進(jìn)內(nèi)存除非你確實(shí)需要且內(nèi)存扛得住。把讀圖放到__getitem__讓worker進(jìn)程并行去讀才是正確姿勢(shì)。from torch.utils.data import Dataset from PIL import Image class MyDataset(Dataset): def __init__(self, samples, transformNone): self.samples samples # [(img_path, label), ...] self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label注意Image.open是惰性的真正解碼發(fā)生在讀取像素時(shí)所以convert(RGB)這一步很重要能統(tǒng)一處理灰度圖、帶alpha通道的PNG避免后面transform報(bào)通道數(shù)不匹配。這個(gè)細(xì)節(jié)我在處理混合來(lái)源圖片時(shí)吃過(guò)虧加上convert之后世界清凈了。3.4 num_workers、pin_memory、persistent_workers這幾個(gè)參數(shù)怎么定DataLoader的參數(shù)里對(duì)性能影響最大的就是num_workers。它決定開(kāi)幾個(gè)子進(jìn)程并行取數(shù)據(jù)。設(shè)成0表示在主進(jìn)程里取簡(jiǎn)單但會(huì)阻塞訓(xùn)練設(shè)成正數(shù)則開(kāi)子進(jìn)程。經(jīng)驗(yàn)值是不超過(guò)CPU物理核心數(shù)一般從4或8起步邊跑邊看GPU利用率。如果GPU利用率上不去、經(jīng)常等數(shù)據(jù)就加worker如果CPU已經(jīng)打滿、內(nèi)存暴漲就減worker。盲目設(shè)很大反而會(huì)因?yàn)檫M(jìn)程調(diào)度和內(nèi)存拷貝拖慢整體。pin_memoryTrue會(huì)把batch數(shù)據(jù)放到鎖頁(yè)內(nèi)存加速?gòu)腃PU到GPU的拷貝用GPU訓(xùn)練時(shí)基本必開(kāi)。persistent_workersTrue讓worker進(jìn)程在epoch之間不銷(xiāo)毀重建減少反復(fù)啟動(dòng)的開(kāi)銷(xiāo)數(shù)據(jù)量大、epoch多的時(shí)候收益明顯。這幾個(gè)參數(shù)組合起來(lái)我常用的訓(xùn)練配置是num_workers8, pin_memoryTrue, persistent_workersTrue配合batch_size64左右在單卡上跑得比較順。注意Windows上num_workers0有時(shí)會(huì)遇到多進(jìn)程啟動(dòng)問(wèn)題如果報(bào)奇怪的pickle錯(cuò)誤先把num_workers設(shè)成0確認(rèn)邏輯沒(méi)問(wèn)題再逐步調(diào)大排查。4. 實(shí)操過(guò)程與核心環(huán)節(jié)實(shí)現(xiàn)搭一條能跑的訓(xùn)練數(shù)據(jù)流4.1 從零搭一條完整的數(shù)據(jù)管道我把完整流程拆成五步定義transform、構(gòu)建Dataset、包DataLoader、寫(xiě)訓(xùn)練循環(huán)取batch、驗(yàn)證集單獨(dú)配置。下面這段是可以直接跑的骨架以ImageFolder為例。from torch.utils.data import DataLoader import torchvision train_set torchvision.datasets.ImageFolder(./data/train, transformtrain_tf) val_set torchvision.datasets.ImageFolder(./data/val, transformval_tf) train_loader DataLoader(train_set, batch_size64, shuffleTrue, num_workers8, pin_memoryTrue, persistent_workersTrue, drop_lastTrue) val_loader DataLoader(val_set, batch_size128, shuffleFalse, num_workers8, pin_memoryTrue) for epoch in range(num_epochs): for images, labels in train_loader: images images.cuda(non_blockingTrue) labels labels.cuda(non_blockingTrue) # 前向、反向、更新這里drop_lastTrue是為了保證每個(gè)batch大小一致避免最后一個(gè)不滿的batch影響B(tài)atchNorm的統(tǒng)計(jì)訓(xùn)練時(shí)我一般開(kāi)著驗(yàn)證集則保留全部樣本保證評(píng)估覆蓋完整。non_blockingTrue配合pin_memory能讓拷貝和計(jì)算重疊是個(gè)小但有效的優(yōu)化。4.2 batch_size和num_workers的聯(lián)合調(diào)參過(guò)程這兩個(gè)參數(shù)不是獨(dú)立的。我通常的做法是先固定num_workers4把batch_size從32往上試觀察顯存占用和單步耗時(shí)找到顯存放得下的最大值附近然后固定這個(gè)batch_size把num_workers從4加到8、16看GPU利用率和吞吐變化。如果加worker后吞吐不再提升甚至下降說(shuō)明數(shù)據(jù)加載已經(jīng)不是瓶頸停在這個(gè)值就好。舉個(gè)我實(shí)際調(diào)過(guò)的例子某次訓(xùn)練單步耗時(shí)約120msGPU利用率只有50%左右明顯在等數(shù)據(jù)。把num_workers從4提到8單步降到90ms利用率升到75%再提到16單步反而回到100ms因?yàn)檫M(jìn)程切換和內(nèi)存拷貝開(kāi)銷(xiāo)上來(lái)了。最后定在8。這個(gè)過(guò)程沒(méi)有公式只能實(shí)測(cè)但方向很明確看GPU是不是在等等就加worker加到頭就停。4.3 自定義collate_fn處理變長(zhǎng)樣本做檢測(cè)或序列任務(wù)時(shí)一個(gè)batch里樣本形狀不一致默認(rèn)collate會(huì)失敗。解決辦法是傳自定義collate_fn把樣本列表原樣打包成list返回讓模型內(nèi)部自己處理對(duì)齊。def my_collate(batch): images [item[0] for item in batch] targets [item[1] for item in batch] return images, targets loader DataLoader(dataset, batch_size8, collate_fnmy_collate)代價(jià)是batch不再是規(guī)整的tensor后續(xù)要自己stack或pad。我的經(jīng)驗(yàn)是能pad成統(tǒng)一形狀就盡量pad因?yàn)橐?guī)整tensor對(duì)GPU更友好實(shí)在不行再用list但要在模型里第一時(shí)間處理掉別讓不規(guī)則結(jié)構(gòu)在計(jì)算圖里傳太深。4.4 用Subset和ConcatDataset做數(shù)據(jù)集切分與合并實(shí)際項(xiàng)目里經(jīng)常需要從一個(gè)大訓(xùn)練集里切出小份做快速實(shí)驗(yàn)或者把多個(gè)來(lái)源的數(shù)據(jù)合并訓(xùn)練。torch.utils.data.Subset按索引列表取子集ConcatDataset把多個(gè)數(shù)據(jù)集首尾拼接。這兩個(gè)工具配合random_split能覆蓋大部分切分需求。from torch.utils.data import random_split, ConcatDataset train_part, val_part random_split(full_set, [0.9, 0.1], generatortorch.Generator().manual_seed(42)) combined ConcatDataset([set_a, set_b])random_split一定要傳固定種子的generator否則每次切分結(jié)果不同實(shí)驗(yàn)沒(méi)法復(fù)現(xiàn)。這個(gè)坑我踩過(guò)不止一次后來(lái)養(yǎng)成習(xí)慣凡是涉及隨機(jī)切分的地方都顯式給種子。5. 常見(jiàn)問(wèn)題與排查技巧實(shí)錄5.1 常見(jiàn)報(bào)錯(cuò)速查表現(xiàn)象可能原因排查與解決TypeError: default_collate: batch must contain tensors樣本返回類(lèi)型不統(tǒng)一或含非tensor檢查_(kāi)_getitem__返回值統(tǒng)一轉(zhuǎn)tensor或自定義collateRuntimeError: DataLoader worker exited unexpectedlyworker內(nèi)異常被吞、內(nèi)存不足先設(shè)num_workers0復(fù)現(xiàn)真實(shí)報(bào)錯(cuò)再查內(nèi)存訓(xùn)練loss劇烈抖動(dòng)驗(yàn)證集誤用隨機(jī)transform檢查val的transform是否含隨機(jī)操作GPU利用率低、訓(xùn)練慢數(shù)據(jù)加載成瓶頸增大num_workers、開(kāi)pin_memory顯存OOM但batch不大圖片分辨率過(guò)高或transform緩存降分辨率、檢查是否在__init__里緩存了整圖多進(jìn)程下卡死Windows多進(jìn)程或共享狀態(tài)問(wèn)題設(shè)num_workers0驗(yàn)證避免在Dataset里存不可pickle對(duì)象5.2 幾個(gè)我踩過(guò)的坑和獨(dú)家心得第一個(gè)坑是在Dataset里存了打開(kāi)的文件句柄或數(shù)據(jù)庫(kù)連接。這些對(duì)象不可picklenum_workers0時(shí)子進(jìn)程無(wú)法繼承直接報(bào)錯(cuò)。正確做法是在__getitem__里按需打開(kāi)、用完即關(guān)或者用worker_init_fn在每個(gè)worker里單獨(dú)初始化。第二個(gè)坑是transform里用了帶狀態(tài)的隨機(jī)數(shù)生成器。多worker下每個(gè)進(jìn)程的隨機(jī)種子默認(rèn)不同這本身沒(méi)問(wèn)題但如果你依賴(lài)某個(gè)全局隨機(jī)狀態(tài)做可復(fù)現(xiàn)增強(qiáng)就會(huì)亂。我的做法是增強(qiáng)全部用torchvision自帶的隨機(jī)變換它們內(nèi)部處理好了種子不自己造輪子。第三個(gè)心得是先用小數(shù)據(jù)跑通再上全量。我習(xí)慣先取100條樣本、num_workers0跑兩個(gè)epoch確認(rèn)數(shù)據(jù)流、transform、collate全對(duì)再切到全量和多worker。這樣出問(wèn)題時(shí)排查范圍小得多比一上來(lái)就全量多進(jìn)程調(diào)試省時(shí)間。第四個(gè)心得是驗(yàn)證集的DataLoader不要開(kāi)shuffle也不要drop_last。評(píng)估要覆蓋全部樣本且順序穩(wěn)定否則指標(biāo)沒(méi)法橫向比較。這個(gè)看似基礎(chǔ)但我在趕進(jìn)度時(shí)犯過(guò)導(dǎo)致兩次評(píng)估結(jié)果對(duì)不上白查了半天模型。5.3 內(nèi)存與性能的平衡技巧數(shù)據(jù)量大時(shí)num_workers開(kāi)大容易把內(nèi)存吃滿因?yàn)槊總€(gè)worker都會(huì)緩存一部分?jǐn)?shù)據(jù)。我的做法是監(jiān)控內(nèi)存占用找到吞吐和內(nèi)存的平衡點(diǎn)。另外如果圖片很大而模型輸入很小可以在__getitem__里先做一次快速降采樣再進(jìn)transform能顯著減少解碼和傳輸開(kāi)銷(xiāo)。還有個(gè)小技巧把常用的歸一化參數(shù)、均值方差提前算好寫(xiě)死別在每次transform里現(xiàn)算雖然省得不多但積少成多。6. 采樣策略進(jìn)階讓DataLoader按你的意圖取數(shù)據(jù)6.1 WeightedRandomSampler解決類(lèi)別不平衡分類(lèi)任務(wù)里類(lèi)別不平衡很常見(jiàn)少數(shù)類(lèi)樣本少模型容易偏向多數(shù)類(lèi)。WeightedRandomSampler給每個(gè)樣本一個(gè)采樣權(quán)重讓少數(shù)類(lèi)被抽到的概率更高。權(quán)重通常設(shè)成類(lèi)別頻率的倒數(shù)按樣本所屬類(lèi)別賦權(quán)。from torch.utils.data import WeightedRandomSampler import numpy as np labels [label for _, label in train_set.samples] class_counts np.bincount(labels) class_weights 1.0 / class_counts sample_weights [class_weights[l] for l in labels] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) loader DataLoader(train_set, batch_size64, samplersampler, num_workers8)注意用了sampler就不能再設(shè)shuffleTrue兩者沖突。replacementTrue表示有放回采樣這是加權(quán)采樣的常規(guī)做法。我實(shí)測(cè)下來(lái)這個(gè)改動(dòng)對(duì)少數(shù)類(lèi)的召回提升很明顯代價(jià)是每個(gè)epoch看到的樣本組合不同訓(xùn)練曲線會(huì)稍微抖一點(diǎn)屬于正?，F(xiàn)象。6.2 自定義BatchSampler實(shí)現(xiàn)配對(duì)采樣做對(duì)比學(xué)習(xí)或?qū)\生網(wǎng)絡(luò)時(shí)一個(gè)batch需要成對(duì)的正負(fù)樣本。這時(shí)候自定義BatchSampler最直接它接收一個(gè)基礎(chǔ)sampler按你的規(guī)則把索引組織成batch列表。from torch.utils.data import Sampler class PairBatchSampler(Sampler): def __init__(self, labels, batch_size, seed42): self.labels labels self.batch_size batch_size self.rng np.random.default_rng(seed) def __iter__(self): # 簡(jiǎn)化示意按類(lèi)別分組后配對(duì)取樣 ... def __len__(self): return len(self.labels) // self.batch_size這類(lèi)自定義采樣器邏輯不復(fù)雜但要注意__len__返回的batch數(shù)量要和實(shí)際迭代次數(shù)一致否則訓(xùn)練循環(huán)的進(jìn)度估計(jì)會(huì)錯(cuò)。我一般寫(xiě)完先用小數(shù)據(jù)打印幾個(gè)batch的索引確認(rèn)配對(duì)邏輯對(duì)再上全量。6.3 分布式訓(xùn)練下的DistributedSampler要點(diǎn)多卡訓(xùn)練時(shí)每個(gè)進(jìn)程要用DistributedSampler保證各卡看到不重疊的數(shù)據(jù)分片且每個(gè)epoch重新打亂。關(guān)鍵參數(shù)是num_replicas總進(jìn)程數(shù)和rank當(dāng)前進(jìn)程編號(hào)以及每個(gè)epoch開(kāi)始前調(diào)用sampler.set_epoch(epoch)否則各epoch的分片順序不變等于沒(méi)打亂。這個(gè)set_epoch我漏調(diào)過(guò)一次訓(xùn)練效果莫名變差查了很久才發(fā)現(xiàn)是分片沒(méi)刷新。7. 我個(gè)人的幾條實(shí)操建議數(shù)據(jù)加載這塊我的核心體會(huì)是先把單條樣本的正確性驗(yàn)證到位再談并行和性能。很多人一上來(lái)就調(diào)num_workers、追吞吐結(jié)果底層transform有bug多進(jìn)程下報(bào)錯(cuò)信息還被吞掉排查成本翻倍。我現(xiàn)在的固定流程是小樣本、單進(jìn)程、打印幾條樣本的形狀和標(biāo)簽確認(rèn)無(wú)誤后再逐步放大。另外別迷信默認(rèn)參數(shù)。DataLoader的默認(rèn)值是為了通用性不是為你的任務(wù)優(yōu)化過(guò)的。batch_size、num_workers、pin_memory、drop_last這幾個(gè)值得花半小時(shí)實(shí)測(cè)調(diào)一遍收益比調(diào)模型結(jié)構(gòu)來(lái)得快。最后再分享一個(gè)小習(xí)慣把數(shù)據(jù)管道的配置transform、batch、worker數(shù)和模型配置分開(kāi)寫(xiě)實(shí)驗(yàn)記錄里單獨(dú)存一份復(fù)現(xiàn)的時(shí)候能省很多事。這套東西用熟了你會(huì)發(fā)現(xiàn)數(shù)據(jù)加載不再是訓(xùn)練的瓶頸而是你可以放心交給它的基礎(chǔ)設(shè)施。