據(jù)管道全鏈路優(yōu)化：從Dataset到DataLoader的實(shí)踐)
深度學(xué)習(xí)項(xiàng)目做到后期大家基本都有一個(gè)共識(shí)模型結(jié)構(gòu)再花哨優(yōu)化器再先進(jìn)只要數(shù)據(jù)管道卡脖子GPU就只能干瞪眼。我見(jiàn)過(guò)太多團(tuán)隊(duì)把時(shí)間耗在調(diào)網(wǎng)絡(luò)結(jié)構(gòu)上結(jié)果一跑訓(xùn)練nvidia-smi里 GPU 利用率只有 30%剩下 70% 時(shí)間全在等數(shù)據(jù)。PyTorch 的數(shù)據(jù)引擎——從自定義Dataset封裝、DataLoader參數(shù)配置到多源數(shù)據(jù)融合——看著不起眼卻是決定訓(xùn)練吞吐量上限的核心環(huán)節(jié)。這篇文章不聊模型只聊數(shù)據(jù)側(cè)的全鏈路工程實(shí)踐怎么把自定義數(shù)據(jù)封裝得既優(yōu)雅又高效怎么配置 DataLoader 才能把 GPU 喂飽以及多源異構(gòu)數(shù)據(jù)圖像、視頻、高光譜文件、傳感器時(shí)序怎么在同一個(gè)管道里平穩(wěn)融合。內(nèi)容基于我個(gè)人在多個(gè)項(xiàng)目里的實(shí)戰(zhàn)經(jīng)驗(yàn)適合已經(jīng)跑通 PyTorch 基礎(chǔ)訓(xùn)練流程、想進(jìn)一步優(yōu)化訓(xùn)練效率和數(shù)據(jù)集擴(kuò)展性的同學(xué)參考。1. 為什么性能瓶頸常常藏在你忽略的數(shù)據(jù)管道里1.1 GPU 利用率低的常見(jiàn)元兇先拋一個(gè)反直覺(jué)的結(jié)論大多數(shù)訓(xùn)練任務(wù)跑得慢不是模型算不動(dòng)而是數(shù)據(jù)送不到。幾年前我接手一個(gè)視頻分類項(xiàng)目模型是 3D CNN顯存占用正常但訓(xùn)練一個(gè) epoch 要 7 個(gè)多小時(shí)。后來(lái)用最簡(jiǎn)單的計(jì)時(shí)法測(cè)了一下發(fā)現(xiàn)每個(gè) step 里next(iter(loader))平均耗時(shí)接近 600 毫秒而 GPU 前向加反向只需要 200 毫秒。換句話說(shuō)有將近三分之二的時(shí)間GPU 都在空等。這種問(wèn)題在圖像分類里可能還不明顯因?yàn)閱螐垐D片讀取也就幾十毫秒但一旦涉及視頻幀序列、高光譜 HDR 文件、點(diǎn)云這類大塊頭數(shù)據(jù)IO 耗時(shí)會(huì)被急劇放大。很多人的第一反應(yīng)是換更好的硬盤但 SSD 和 HDD 的差異只能解決一部分問(wèn)題——如果瓶頸在 CPU 預(yù)處理解碼、歸一化、數(shù)據(jù)增強(qiáng)換硬盤根本沒(méi)意義。1.2 IO、CPU 與 GPU 之間的三角協(xié)作關(guān)系理解 PyTorch 數(shù)據(jù)引擎首先要建立生產(chǎn)-消費(fèi)的流水線思維。GPU 是消費(fèi)方它只認(rèn) CUDA 張量數(shù)據(jù)管道是生產(chǎn)方負(fù)責(zé)從磁盤讀原始數(shù)據(jù)、做預(yù)處理、組成 batch、搬到顯存。這三者之間的協(xié)作關(guān)系可以用一個(gè)例子類比GPU 是流水線上的裝配工CPU 是物料搬運(yùn)工磁盤是倉(cāng)庫(kù)。裝配工手速再快如果搬運(yùn)工一次只能搬一件、而且搬一件要跑一趟倉(cāng)庫(kù)整條流水線就廢了。PyTorch 的DataLoader之所以引入num_workers、prefetch_factor、pin_memory這些參數(shù)本質(zhì)上就是給這條流水線加多搬運(yùn)工和緩沖區(qū)。理解了這一層你就能明白調(diào)數(shù)據(jù)管道的本質(zhì)是在有限的 CPU 核數(shù)和內(nèi)存帶寬下盡可能讓數(shù)據(jù)準(zhǔn)備的速度追平甚至超過(guò) GPU 的消費(fèi)速度。1.3 先學(xué)會(huì)用兩分鐘定位瓶頸在動(dòng)手優(yōu)化之前先把問(wèn)題定位清楚。我的習(xí)慣是三步走訓(xùn)練時(shí)另開(kāi)一個(gè)終端跑watch -n 1 nvidia-smi看 GPU 利用率是否經(jīng)常掉到 80% 以下在訓(xùn)練循環(huán)里單獨(dú)計(jì)時(shí)start time.time(); data next(iter(loader)); print(time.time() - start)連續(xù)打幾十個(gè) step看數(shù)據(jù)加載耗時(shí)是否穩(wěn)定用 PyTorch 自帶的torch.utils.bottleneck跑一次它能粗粒度地區(qū)分 CPU、GPU、數(shù)據(jù)加載各自的耗時(shí)占比。這三步做完基本就能判斷是磁盤 IO 的鍋、CPU 預(yù)處理的鍋還是參數(shù)配置不當(dāng)?shù)腻?。下面各?jié)的內(nèi)容都是在定位完瓶頸之后才動(dòng)手的具體優(yōu)化手段。2. 自定義 Dataset三個(gè)核心方法與一個(gè)隱藏契約2.1__init__、__len__、__getitem__的標(biāo)準(zhǔn)寫(xiě)法和職責(zé)邊界自定義數(shù)據(jù)封裝聽(tīng)起來(lái)是老生常談但很多人對(duì)三個(gè)方法的職責(zé)邊界理解是模糊的。我的原則是__init__只負(fù)責(zé)記賬不負(fù)責(zé)干活。它應(yīng)該把所有樣本的索引、路徑、標(biāo)簽元數(shù)據(jù)預(yù)先整理成一個(gè)列表或 DataFrame而不是在這個(gè)階段就去讀數(shù)據(jù)。原因很簡(jiǎn)單__init__只在創(chuàng)建 Dataset 時(shí)執(zhí)行一次如果在這里做耗時(shí)操作不僅拖慢啟動(dòng)還會(huì)在訓(xùn)練過(guò)程中毫無(wú)意義地占用內(nèi)存。__len__返回樣本總數(shù)必須是一個(gè)確定且穩(wěn)定的數(shù)值。這個(gè)值不僅用于for epoch的輪數(shù)計(jì)算還直接影響Sampler的索引分配和分布式訓(xùn)練的DistributedSampler切分邏輯所以不能有任何歧義。__getitem__是真正的生產(chǎn)車間接收一個(gè)整數(shù)索引返回一個(gè)樣本。這里才是讀文件、解碼、預(yù)處理、數(shù)據(jù)增強(qiáng)該發(fā)生的地方。一個(gè)標(biāo)準(zhǔn)的圖像分類 Dataset 大致長(zhǎng)這樣class ImageDataset(Dataset): def __init__(self, file_list, label_list, transformNone): self.file_list file_list self.label_list label_list self.transform transform # 這里可以做一次數(shù)據(jù)校驗(yàn)比如檢查文件是否都存在 # 但不要讀文件內(nèi)容 def __len__(self): return len(self.file_list) def __getitem__(self, idx): img Image.open(self.file_list[idx]).convert(RGB) if self.transform: img self.transform(img) label self.label_list[idx] return img, label2.2 緩存策略什么數(shù)據(jù)適合在__init__里就加載進(jìn)內(nèi)存前面我說(shuō)__init__不要干活但有一個(gè)例外當(dāng)你的數(shù)據(jù)集總量小到能裝進(jìn)內(nèi)存時(shí)直接在__init__里把所有數(shù)據(jù)讀進(jìn)來(lái)然后__getitem__只做內(nèi)存切片效果會(huì)好到爆炸。比如 CIFAR-10、MNIST 這類數(shù)據(jù)集幾十萬(wàn)個(gè)樣本也就幾個(gè) GB完全可以全量駐留內(nèi)存把磁盤 IO 徹底消除。判斷標(biāo)準(zhǔn)很簡(jiǎn)單數(shù)據(jù)集壓縮后的大小 × 解壓后的膨脹系數(shù)如果不超過(guò)可用內(nèi)存的 50%就值得做全量緩存。需要注意的是解壓膨脹系數(shù)往往很驚人一個(gè) 100MB 的 JPEG 壓縮包解壓成 RGB 張量數(shù)組可能占 300MB 以上。所以穩(wěn)妥的做法是緩存原始字節(jié)而不是緩存解碼后的張量class CachedImageDataset(Dataset): def __init__(self, file_list, label_list, transformNone): self.raw_bytes [] for path in file_list: with open(path, rb) as f: self.raw_bytes.append(f.read()) self.label_list label_list self.transform transform def __getitem__(self, idx): img Image.open(io.BytesIO(self.raw_bytes[idx])).convert(RGB) ...這樣既省掉了重復(fù)的磁盤尋址又不會(huì)一次性把解碼后的矩陣全部堆在內(nèi)存里。實(shí)測(cè)中這種方式對(duì)幾萬(wàn)張小圖的數(shù)據(jù)集能帶來(lái) 5 到 10 倍的加載提速。2.3 分布式訓(xùn)練下 Dataset 的數(shù)據(jù)串?dāng)_風(fēng)險(xiǎn)單機(jī)單卡時(shí)Dataset 的索引隨便寫(xiě)都沒(méi)問(wèn)題。但一旦切到DistributedDataParallel事情就變了每個(gè)進(jìn)程會(huì)通過(guò)DistributedSampler拿到數(shù)據(jù)集的一個(gè)分片。如果你在 Dataset 里用了全局的隨機(jī)狀態(tài)比如給樣本打亂順序、隨機(jī)采樣又沒(méi)有做種子隔離就會(huì)出現(xiàn)不同進(jìn)程讀到重復(fù)樣本、或者某個(gè)進(jìn)程漏掉樣本的問(wèn)題。一個(gè)典型的錯(cuò)誤是在 Dataset 內(nèi)部維護(hù)一個(gè)隨機(jī)索引列表每次__getitem__的時(shí)候動(dòng)態(tài) shuffle。這會(huì)讓DistributedSampler的分片邏輯完全失效因?yàn)槟銓?shí)際返回的樣本和 sampler 給定的索引對(duì)不上。正確的做法是Dataset 保持純粹——給什么索引就返回什么樣本隨機(jī)性交給 DataLoader 的shuffleTrue或 sampler 來(lái)處理。這就是我說(shuō)的隱藏契約Dataset 的索引語(yǔ)義必須是確定的、可重入的任何隨機(jī)化都必須在外部完成。2.4 幾個(gè)常見(jiàn)的反模式排查過(guò)不少同事的代碼以下幾個(gè)坑出現(xiàn)頻率最高在__len__里做耗時(shí)統(tǒng)計(jì)比如每次調(diào)用都重新數(shù)一遍目錄下的文件這會(huì)讓每次 epoch 統(tǒng)計(jì)都變慢。__getitem__里修改全局變量多 worker 場(chǎng)景下主進(jìn)程和 worker 進(jìn)程是fork出來(lái)的全局變量的修改可能會(huì)造成數(shù)據(jù)錯(cuò)亂。返回 Python 對(duì)象而不是張量比如直接返回PIL.Image對(duì)象collate_fn會(huì)拿它沒(méi)辦法最后要么報(bào)錯(cuò)要么被迫走default_collate的慢路徑。這些都是寫(xiě)完 Dataset 之后自查的重點(diǎn)項(xiàng)。3. DataLoader 參數(shù)調(diào)優(yōu)worker 數(shù)量、prefetch 與 pin_memory 的配合邏輯3.1num_workers不是越大越好很多人以為num_workers是數(shù)據(jù)加載速度倍增器調(diào)得越大越好結(jié)果一開(kāi) 16 個(gè) workerCPU 直接被打滿GPU 利用率反而下降。原因在于worker 進(jìn)程之間需要把數(shù)據(jù)通過(guò) IPC進(jìn)程間通信傳回主進(jìn)程這個(gè)傳輸本身有開(kāi)銷同時(shí)每個(gè) worker 都有自己的數(shù)據(jù)預(yù)處理副本如果 CPU 核數(shù)不夠就會(huì)在進(jìn)程調(diào)度上產(chǎn)生大量競(jìng)爭(zhēng)。我的調(diào)參經(jīng)驗(yàn)是num_workers設(shè)置在 CPU 物理核數(shù)的一半到四分之三之間通常比較穩(wěn)。比如 8 核 CPU 就設(shè) 4 到 616 核就設(shè) 8 到 12。但這個(gè)數(shù)字和具體的數(shù)據(jù)格式強(qiáng)相關(guān)——如果每個(gè)樣本的解碼耗時(shí)很長(zhǎng)比如高光譜文件worker 多點(diǎn)反而有優(yōu)勢(shì)因?yàn)槠款i在 CPU 計(jì)算而不是 IPC 傳輸如果樣本本來(lái)就小worker 多了純粹是內(nèi)耗。判斷當(dāng)前 worker 數(shù)是否合適的辦法觀察訓(xùn)練日志里每個(gè) step 的耗時(shí)逐漸增大num_workers如果耗時(shí)不再下降甚至反彈說(shuō)明已經(jīng)到臨界點(diǎn)了。3.2pin_memory和prefetch_factor到底在解決什么問(wèn)題pin_memoryTrue的作用是把數(shù)據(jù)鎖頁(yè)到不可換出的內(nèi)存區(qū)域。理解這個(gè)問(wèn)題需要一點(diǎn) CUDA 基礎(chǔ)GPU 從 CPU 側(cè)拷貝數(shù)據(jù)時(shí)通常是異步 DMA 傳輸而 DMA 要求源地址在物理內(nèi)存中固定不變。如果數(shù)據(jù)在普通內(nèi)存里操作系統(tǒng)可能會(huì)把它換出到交換分區(qū)導(dǎo)致拷貝前需要先做一次同步頁(yè)表鎖定過(guò)程變慢。所以pin_memory本質(zhì)上是用鎖頁(yè)內(nèi)存換取了GPU 拷貝更快的收益。prefetch_factorPyTorch 1.7 之后可用則是控制每個(gè) worker 最多能預(yù)取多少個(gè) batch 放在隊(duì)列里。默認(rèn)值是 2對(duì)于普通場(chǎng)景夠用如果數(shù)據(jù)加載時(shí)間波動(dòng)比較大——比如有的視頻序列特別長(zhǎng)、有的特別短——可以適當(dāng)調(diào)高到 4 甚至 8給流水線多留一些余量避免偶爾的長(zhǎng)樣本加載把 GPU 餓死。這兩個(gè)參數(shù)配合使用的完整模板DataLoader( dataset, batch_size64, num_workers8, pin_memoryTrue, prefetch_factor4, persistent_workersTrue, # 多輪 epoch 時(shí)復(fù)用 worker避免重復(fù)創(chuàng)建 )persistent_workersTrue這個(gè)參數(shù)也值得單獨(dú)說(shuō)一下它讓 worker 在 epoch 結(jié)束后不退出下一次 epoch 繼續(xù)復(fù)用。省掉了反復(fù) fork worker 進(jìn)程的開(kāi)銷對(duì)訓(xùn)練幾十個(gè) epoch 的場(chǎng)景幫助很大。但注意它要求 Dataset 是可重入的不能依賴在__getitem__里做不可重復(fù)的操作。3.3 自定義collate_fn的場(chǎng)景變長(zhǎng)序列與多源數(shù)據(jù)拼接collate_fn是 DataLoader 中另一個(gè)被低估的組件。默認(rèn)的default_collate做的事情是把 batch 中每個(gè)樣本的張量堆疊成一個(gè)更大的張量。這個(gè)操作要求所有樣本的 shape 完全一致一旦遇到變長(zhǎng)序列、不同尺寸的圖像、或者樣本本身就是混合結(jié)構(gòu)它就會(huì)直接拋異常。遇到變長(zhǎng)數(shù)據(jù)最常見(jiàn)的處理思路是pad 到 batch 內(nèi)最大長(zhǎng)度。我自己在 NLP 和視頻特征序列項(xiàng)目里都這么干過(guò)def collate_variable_length(batch): images, seqs, labels zip(*batch) images torch.stack(images) max_len max([s.size(0) for s in seqs]) seqs_padded torch.zeros(len(seqs), max_len, seqs[0].size(1)) seq_lens torch.tensor([s.size(0) for s in seqs]) for i, s in enumerate(seqs): seqs_padded[i, :s.size(0)] s return images, seqs_padded, seq_lens, torch.tensor(labels)這里返回seq_lens是為了讓模型在讀取時(shí)能用mask回避 padding 帶來(lái)的無(wú)效計(jì)算。如果你用的是 Transformer 結(jié)構(gòu)這個(gè) mask 是必須的。3.4 通過(guò)Sampler控制采樣邏輯而不是改 Dataset很多人遇到某些類樣本太多、想均衡采樣的需求時(shí)第一反應(yīng)是往 Dataset 里加邏輯。但更干凈的做法是實(shí)現(xiàn)自定義Sampler。PyTorch 提供了WeightedRandomSampler可以按權(quán)重控制每個(gè)樣本被抽到的概率。對(duì)于類別不平衡的多源融合數(shù)據(jù)這個(gè)方案非常實(shí)用weights [1.0 / class_sample_count[label] for label in all_labels] sampler WeightedRandomSampler(weights, num_sampleslen(weights), replacementTrue)這里的關(guān)鍵是按類別頻率的倒數(shù)設(shè)權(quán)重讓少數(shù)類樣本在每輪 epoch 里被抽到的概率更高。我曾在類別比例 100:1 的多源數(shù)據(jù)集上用這個(gè)方案把少數(shù)類 F1 從 0.3 提到了 0.58效果比在 loss 里加權(quán)重更直接。4. 多源異構(gòu)數(shù)據(jù)的融合從硬拼到軟融的工程選型4.1 什么是多源融合為什么它不是簡(jiǎn)單的 concat多源融合multi-source fusion在機(jī)器人定位、遙感分析、多模態(tài)識(shí)別等場(chǎng)景里非常常見(jiàn)。它的難點(diǎn)不只是把數(shù)據(jù)拼起來(lái)而是不同類型的數(shù)據(jù)有完全不同的采樣頻率、數(shù)據(jù)維度、語(yǔ)義粒度和缺失模式。比如一個(gè)機(jī)器人定位項(xiàng)目里圖像是 30 FPSIMU 是 200 HzGPS 是 1 Hz三個(gè)源在時(shí)間軸上天然不對(duì)齊。如果只是簡(jiǎn)單地在樣本里 concat模型學(xué)到的往往是強(qiáng)勢(shì)模態(tài)的信息弱勢(shì)模態(tài)會(huì)被淹沒(méi)。所以在工程上我傾向于把多源融合分成三個(gè)層次來(lái)設(shè)計(jì)數(shù)據(jù)層融合在 Dataset 層面把多個(gè)源組織成統(tǒng)一樣本結(jié)構(gòu)特征層融合各模態(tài)先獨(dú)立編碼成特征向量再在模型內(nèi)部融合決策層融合各模態(tài)分別出預(yù)測(cè)結(jié)果再用加權(quán)投票或門控網(wǎng)絡(luò)綜合。數(shù)據(jù)引擎主要負(fù)責(zé)的是第一層——把多源數(shù)據(jù)送進(jìn)模型之前先解決好怎么組織、怎么對(duì)齊、怎么平衡這三個(gè)問(wèn)題。4.2 用ConcatDataset和自定義FusionDataset組織多源數(shù)據(jù)如果你的多源數(shù)據(jù)是不同來(lái)源但相同標(biāo)簽空間的樣本比如一部分是高光譜圖像、一部分是 RGB 圖像label 都是地物類別——最簡(jiǎn)單的方式是分別建 Dataset再用torch.utils.data.ConcatDataset合并。它內(nèi)部維護(hù)了各子數(shù)據(jù)集的樣本偏移量__len__會(huì)返回總和、__getitem__會(huì)根據(jù)索引自動(dòng)路由到正確的子集。優(yōu)點(diǎn)是零成本、天然支持缺點(diǎn)是無(wú)法在采樣時(shí)對(duì)源做權(quán)重控制——你只能用WeightedRandomSampler在全局樣本層面調(diào)整。如果你的多源數(shù)據(jù)是同一個(gè)樣本的不同模態(tài)比如一個(gè)樣本同時(shí)包含視頻幀和對(duì)應(yīng)的 IMU 時(shí)序那就需要自定義一個(gè)FusionDataset讓__getitem__返回一個(gè)字典結(jié)構(gòu)class FusionDataset(Dataset): def __init__(self, video_files, imu_files, labels): self.video_files video_files self.imu_files imu_files self.labels labels def __getitem__(self, idx): frames self._load_video_frames(self.video_files[idx]) # [T, H, W, C] imu_seq self._load_imu_sequence(self.imu_files[idx]) # [T, D] # 在這里做時(shí)間對(duì)齊、插值或分段采樣 frames, imu_seq self._align(frames, imu_seq) return {frames: frames, imu: imu_seq, label: self.labels[idx]}返回字典的好處是后續(xù)collate_fn可以精準(zhǔn)地按 key 處理不同張量——視頻幀要 stack 成五維張量IMU 序列要 pad 到等長(zhǎng)label 單獨(dú)收集。這種結(jié)構(gòu)比返回 tuple 清晰得多也更容易擴(kuò)展新的數(shù)據(jù)源。4.3 時(shí)間對(duì)齊不同采樣頻率數(shù)據(jù)的關(guān)鍵難題及處理策略時(shí)間對(duì)齊是多源融合里最容易被忽視、也最容易出錯(cuò)的地方。以視頻和 IMU 為例視頻每幀有固定的時(shí)間戳IMU 的每個(gè)采樣也有時(shí)間戳但兩者的時(shí)間戳起點(diǎn)往往不同采樣頻率也不同。我通常采用以低頻源為錨點(diǎn)的策略以視頻幀的時(shí)間戳作為標(biāo)準(zhǔn)對(duì) IMU 序列做線性插值取出與每幀時(shí)間戳對(duì)齊的 IMU 值def align_imu_to_frames(frame_timestamps, imu_timestamps, imu_values): aligned np.zeros((len(frame_timestamps), imu_values.shape[1])) for i, t in enumerate(frame_timestamps): idx np.searchsorted(imu_timestamps, t) if idx 0: aligned[i] imu_values[0] elif idx len(imu_timestamps): aligned[i] imu_values[-1] else: t0, t1 imu_timestamps[idx-1], imu_timestamps[idx] alpha (t - t0) / (t1 - t0 1e-8) aligned[i] imu_values[idx-1] * (1 - alpha) imu_values[idx] * alpha return aligned注意這里對(duì)邊界情況做了鉗制——時(shí)間戳落在 IMU 采樣范圍之外時(shí)用首尾值填充而不是報(bào)錯(cuò)。理由是傳感器數(shù)據(jù)在啟動(dòng)和停止階段經(jīng)常有缺失你需要在工程上容忍這些不完美。4.4 多源不平衡如何讓最后加進(jìn)來(lái)的那個(gè)源不白加多源融合項(xiàng)目最常見(jiàn)的結(jié)局是加了新數(shù)據(jù)源模型效果沒(méi)變甚至變差了。原因多數(shù)不是模型不行而是采樣策略把大量便宜的源比如圖像反復(fù)喂給模型而昂貴的源比如激光雷達(dá)點(diǎn)云出現(xiàn)的頻率太低模型根本學(xué)不到它的規(guī)律。解決思路有兩個(gè)層面。一是前面提到的WeightedRandomSampler在樣本層面調(diào)節(jié)各類別/各源的采樣概率二是引入源感知的 batch 組成。我正在用的一個(gè)方案是FusionDataset返回時(shí)額外給出source_id然后在collate_fn里統(tǒng)計(jì)當(dāng)前 batch 中各源的比例如果某個(gè)源占比低于預(yù)設(shè)閾值就用replacementTrue的方式強(qiáng)制多抽幾個(gè)該源的樣本補(bǔ)進(jìn)來(lái)。雖然這樣做會(huì)讓 batch 的分布和自然分布有偏差但在多源融合場(chǎng)景里保持每個(gè) batch 都能讓模型看到所有模態(tài)比嚴(yán)格遵循自然分布更重要。5. 高光譜、視頻等大文件數(shù)據(jù)的管道工程化處理5.1 高光譜 HDR/SPE 文件的讀取姿勢(shì)高光譜數(shù)據(jù)是數(shù)據(jù)引擎的極端壓力測(cè)試。它不像普通圖片那樣有統(tǒng)一的解碼庫(kù)HDR 和 SPE 文件往往來(lái)自特定儀器廠商格式五花八門。網(wǎng)上搜pytorch 處理高光譜 hdr 文件和 spe 文件能找到一堆零散腳本但真正可復(fù)用的經(jīng)驗(yàn)是下面幾條區(qū)分文件頭和數(shù)據(jù)區(qū)HDR/SPE 文件頭部包含波長(zhǎng)信息、采集參數(shù)等元數(shù)據(jù)數(shù)據(jù)區(qū)是原始光譜矩陣。讀取時(shí)通常用numpy.fromfile跳過(guò)文件頭再按(bands, height, width)的維度 reshape。不同廠商的字節(jié)序可能不同一般用f4或f4顯式指定大端/小端避免解析出亂碼。使用np.memmap做內(nèi)存映射整幅高光譜圖像可能達(dá)到幾百 MB 甚至幾個(gè) GB__getitem__里直接np.load會(huì)把內(nèi)存撐爆。用np.memmap可以按需讀取目標(biāo)波段或目標(biāo)區(qū)域def read_hdr_band(file_path, offset, band_index, height, width): mmap np.memmap(file_path, dtypef4, moder, offsetoffset, shape(bands, height, width)) band_data mmap[band_index].copy() # copy 出來(lái)再返回避免 memmap 句柄泄漏 return band_data波段采樣高光譜波段可能有幾百個(gè)但相鄰波段相關(guān)性極高。做分類任務(wù)時(shí)沒(méi)必要把所有波段都喂給模型通過(guò)主成分分析或互信息篩選出幾十個(gè)代表波段數(shù)據(jù)量直接降一個(gè)量級(jí)訓(xùn)練速度和效果反而更好。5.2 視頻動(dòng)作分類的數(shù)據(jù)管道設(shè)計(jì)視頻分類任務(wù)類似 UCF101 數(shù)據(jù)集實(shí)戰(zhàn)的數(shù)據(jù)管道比圖像復(fù)雜在多了一維時(shí)間。一個(gè) 10 秒的視頻均勻抽 16 幀每幀 224×224解碼成本是單張圖片的十幾倍。而且視頻解碼器OpenCV 的VideoCapture或者PyAV是出了名的又慢又不穩(wěn)定處理不好就是訓(xùn)練最大短板。我的做法是離線抽幀在線采樣訓(xùn)練前先把視頻均勻抽幀保存為單張 JPEG文件名帶上時(shí)間戳訓(xùn)練時(shí)Dataset只維護(hù)視頻 ID 幀索引列表__getitem__隨機(jī)選一個(gè)起始幀連續(xù)取 N 幀做數(shù)據(jù)增強(qiáng)。這樣做有幾個(gè)好處一是解碼成本被一次性攤銷二是隨機(jī)片段的采樣策略天然實(shí)現(xiàn)了時(shí)序數(shù)據(jù)增強(qiáng)三是因?yàn)閹呀?jīng)是單張圖片可以沿用圖像格式的成熟緩存策略。一個(gè)額外的坑是OpenCV 的VideoCapture在線讀取視頻時(shí)seek 精度不穩(wěn)定某些編碼格式下set(CAP_PROP_POS_FRAMES)跳轉(zhuǎn)到的幀不是目標(biāo)幀。離線抽幀徹底規(guī)避了這個(gè)問(wèn)題。5.3 雙隊(duì)列架構(gòu)讓數(shù)據(jù)準(zhǔn)備和模型訓(xùn)練互不阻塞當(dāng)單進(jìn)程的DataLoader無(wú)論如何也喂不飽 GPU 時(shí)就該考慮更極致的方案了。我在處理超大數(shù)據(jù)集時(shí)用過(guò)雙隊(duì)列架構(gòu)一個(gè)生產(chǎn)者進(jìn)程負(fù)責(zé)從磁盤批量讀取原始數(shù)據(jù)放進(jìn)內(nèi)存隊(duì)列另一個(gè)進(jìn)程池負(fù)責(zé)從隊(duì)列取數(shù)據(jù)做預(yù)處理和增強(qiáng)再通過(guò)DataLoader的 worker 消費(fèi)。這個(gè)架構(gòu)本質(zhì)上是在 DataLoader 前面再加一級(jí)緩存層。queue multiprocessing.Queue(maxsize8) # producer: 從磁盤讀原始字節(jié) def producer(file_batch): for path in file_batch: with open(path, rb) as f: queue.put(f.read())不過(guò)說(shuō)實(shí)話這個(gè)方案工程復(fù)雜度較高多數(shù)場(chǎng)景用不上。只有當(dāng)num_workers已經(jīng)拉滿、CPU 占用也跑滿了、但 GPU 還是在等數(shù)據(jù)時(shí)才值得考慮引入。普通項(xiàng)目先把 DataLoader 的參數(shù)調(diào)優(yōu)做扎實(shí)收益已經(jīng)非常可觀。5.4 避免重復(fù) IO緩存增強(qiáng) vs 緩存原始數(shù)據(jù)最后聊一個(gè)容易被忽略的優(yōu)化點(diǎn)數(shù)據(jù)增強(qiáng)應(yīng)該放在哪個(gè)環(huán)節(jié)。很多教程把 transform 放在__getitem__里但這意味著每次 epoch 都要重新解碼原始數(shù)據(jù)、重新做增強(qiáng)。如果你的數(shù)據(jù)量不大但訓(xùn)練輪數(shù)很多可以考慮把已增強(qiáng)的結(jié)果緩存到磁盤或內(nèi)存用法是一輪 epoch 生成一次增強(qiáng)副本后面幾輪直接復(fù)用。不同方案對(duì)比方案內(nèi)存占用IO 開(kāi)銷數(shù)據(jù)多樣性適用場(chǎng)景每次完整解碼增強(qiáng)低高高數(shù)據(jù)集大、訓(xùn)練輪數(shù)少緩存原始字節(jié)中中高每次重新增強(qiáng)數(shù)據(jù)集可容納進(jìn)內(nèi)存緩存增強(qiáng)后張量高低低每輪相同數(shù)據(jù)集小、訓(xùn)練輪數(shù)多memmap 按需讀取低低中單樣本文件極大高光譜/視頻注意緩存增強(qiáng)后張量會(huì)導(dǎo)致數(shù)據(jù)多樣性下降容易過(guò)擬合所以通常只在強(qiáng)約束下用比如做知識(shí)蒸餾的固定輸入。我的習(xí)慣是優(yōu)先保證每輪能重新增強(qiáng)只有確認(rèn)內(nèi)存夠用、且增強(qiáng)不是關(guān)鍵因素時(shí)才退回到緩存方案。6. 從環(huán)境到代碼數(shù)據(jù)管道的性能剖析與隱性故障排查6.1 數(shù)據(jù)管道性能剖析的完整排查鏈路前面說(shuō)過(guò)定位瓶頸的三步法這里展開(kāi)講具體操作。我先跑一個(gè)縮短版訓(xùn)練比如 50 個(gè) step在這個(gè)訓(xùn)練里同時(shí)記錄 GPU 利用率和數(shù)據(jù)加載耗時(shí)。拿到數(shù)據(jù)后按下面的順序排查磁盤 IO 是否打滿用iostat -x 1看%util如果長(zhǎng)期接近 100%說(shuō)明瓶頸在磁盤優(yōu)先上緩存或memmap。CPU 是否打滿用htop看各核占用如果所有核都 100% 而 GPU 空閑說(shuō)明預(yù)處理太慢需要優(yōu)化解碼邏輯或增加 worker如果核數(shù)有富余。IPC 傳輸是否成為瓶頸如果 worker 很多、CPU 沒(méi)打滿但 GPU 還是餓很可能數(shù)據(jù)在從 worker 傳回主進(jìn)程的過(guò)程中耗時(shí)過(guò)大。這時(shí)可以試試減小傳輸數(shù)據(jù)量比如在 worker 里提前把圖像轉(zhuǎn)成小尺寸張量或減少num_workers。對(duì)這個(gè)流程我曾吃過(guò)一次大虧有個(gè)項(xiàng)目數(shù)據(jù)在機(jī)械硬盤上num_workers8CPU 完全沒(méi)有打滿但每個(gè) step 的加載時(shí)間穩(wěn)定在 1.2 秒。用iostat一查磁盤%util95% 以上——多個(gè) worker 同時(shí)隨機(jī)讀不同位置機(jī)械硬盤的尋道被反復(fù)放大8 個(gè) worker 反而比 2 個(gè)更慢。換成 SSD 后問(wèn)題直接消失。這說(shuō)明任何調(diào)優(yōu)手段都必須先確認(rèn)瓶頸在哪一層不能盲目堆參數(shù)。6.2 幾個(gè)隱蔽性極強(qiáng)的隱性錯(cuò)誤有些問(wèn)題不會(huì)立刻報(bào)錯(cuò)但會(huì)悄悄污染訓(xùn)練結(jié)果排查起來(lái)非常頭疼。我把自己踩過(guò)、也幫別人排查過(guò)的列幾個(gè)重點(diǎn)Worker 隨機(jī)種子未隔離多個(gè) worker 用同一個(gè)隨機(jī)種子做數(shù)據(jù)增強(qiáng)時(shí)一個(gè) batch 里會(huì)出現(xiàn)多張一模一樣的增強(qiáng)圖模型學(xué)到的特征分布會(huì)被扭曲。解決方式是在 Dataset 里依靠 PyTorch 對(duì)每個(gè) worker 自動(dòng)設(shè)置不同隨機(jī)種子PyTorch 1.x 之后會(huì)為每個(gè) worker 設(shè)置獨(dú)立基礎(chǔ)種子但如果你在外部用np.random.seed手動(dòng)固定了全局種子這個(gè)機(jī)制會(huì)被繞過(guò)。注意不要在數(shù)據(jù)讀取代碼里手動(dòng)重設(shè)全局種子。訓(xùn)練集和驗(yàn)證集共享緩存如果 Dataset 的緩存字典用的是全局變量訓(xùn)練和驗(yàn)證兩個(gè) loader 指向同一個(gè) Dataset 時(shí)驗(yàn)證階段可能讀到訓(xùn)練時(shí)緩存的數(shù)據(jù)。訓(xùn)練和驗(yàn)證必須使用獨(dú)立的 Dataset 實(shí)例或者明確區(qū)分緩存命名空間。文件句柄泄漏使用Image.open、np.memmap、open()時(shí)忘記關(guān)閉或copy長(zhǎng)期運(yùn)行會(huì)導(dǎo)致句柄數(shù)暴漲最后整個(gè)進(jìn)程崩潰。我之前處理高光譜數(shù)據(jù)時(shí)就是因?yàn)?memmap 句柄沒(méi)有主動(dòng)關(guān)閉訓(xùn)練到第 40 個(gè) epoch 時(shí)系統(tǒng)報(bào)Too many open files。在__getitem__里對(duì)memmap做.copy()并顯式del是簡(jiǎn)單可靠的規(guī)避方法。shuffle 和斷點(diǎn)續(xù)訓(xùn)的配合問(wèn)題如果訓(xùn)練中斷后從 checkpoint 恢復(fù)但 DataLoader 的shuffleTrue沒(méi)有固定generator的種子恢復(fù)后的數(shù)據(jù)順序會(huì)和中斷前完全不同學(xué)習(xí)率調(diào)度和最佳 checkpoint 的對(duì)比全部失效。斷點(diǎn)續(xù)訓(xùn)時(shí)必須固定DataLoader的generatorg torch.Generator() g.manual_seed(42) DataLoader(..., shuffleTrue, generatorg)驗(yàn)證集 loader 的 worker 數(shù)過(guò)高驗(yàn)證階段通常不需要那么多數(shù)據(jù)增強(qiáng)但很多人會(huì)把訓(xùn)練的超參原樣套給驗(yàn)證 loader導(dǎo)致驗(yàn)證每個(gè) epoch 耗時(shí)和訓(xùn)練幾乎一樣。驗(yàn)證 loader 一般設(shè)num_workers2或干脆num_workers0、shuffleFalse、prefetch_factor保持默認(rèn)即可。6.3 PyTorch 與 CUDA 環(huán)境配套對(duì)數(shù)據(jù)加載的影響環(huán)境配置雖然不直接屬于數(shù)據(jù)引擎編碼范疇但數(shù)據(jù)張量從 CPU 搬到 GPU 的過(guò)程強(qiáng)依賴 CUDA 版本和 PyTorch 的配套情況。熱詞里頻繁出現(xiàn)pytorch 安裝教程 gpu、cuda 版本、anaconda 配置 pytorch 環(huán)境這類搜索背后其實(shí)是同一個(gè)痛點(diǎn)裝了 PyTorch 的 CPU 版本pin_memory和.cuda()的加速路徑都不能生效數(shù)據(jù)搬移會(huì)經(jīng)過(guò)一條很慢的共享內(nèi)存路徑。我的建議是無(wú)論偏不偏重 GPU都用torch.cuda.is_available()在項(xiàng)目啟動(dòng)時(shí)就做一次顯式檢查并且打印當(dāng)前 CUDA 版本assert torch.cuda.is_available(), CUDA not available, data loading will fall back to CPU-only path print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))環(huán)境是否配套用一句命令就能確認(rèn)python -c import torch; print(torch.__version__, torch.version.cuda, torch.backends.cudnn.version())如果輸出里torch.version.cuda和你在系統(tǒng)里裝的驅(qū)動(dòng)版本差距過(guò)大pin_memory的異步傳輸特性可能無(wú)法正常發(fā)揮訓(xùn)練速度會(huì)比預(yù)期慢一截。另外Windows 上num_workers的 fork 機(jī)制和 Linux 不一樣Windows 默認(rèn)使用spawnworker 會(huì)在每個(gè)進(jìn)程里重新導(dǎo)入整個(gè)模塊如果 Dataset 里涉及大量全局初始化會(huì)顯著拖慢啟動(dòng)——所以生產(chǎn)環(huán)境我強(qiáng)烈建議用 Linux。6.4 一個(gè)完整的自定義 DataLoader 配置參考把前面所有討論落到一個(gè)可復(fù)用的配置模板上。假設(shè)項(xiàng)目是多源融合場(chǎng)景包含視頻片段和 IMU 時(shí)序數(shù)據(jù)量中等偏大transform torchvision.transforms.Compose([ torchvision.transforms.RandomResizedCrop(224), torchvision.transforms.RandomHorizontalFlip(), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) train_dataset FusionDataset(video_files, imu_files, labels, transformtransform) train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers8, # 物理核數(shù)的一半左右 pin_memoryTrue, # GPU 環(huán)境下必開(kāi) prefetch_factor4, # 數(shù)據(jù)波動(dòng)較大預(yù)取多一些 persistent_workersTrue, # 多輪 epoch 復(fù)用 worker collate_fncollate_variable_length, drop_lastTrue, # 避免最后不完整的 batch 影響 BN 統(tǒng)計(jì) )這些參數(shù)不是拍腦袋定的num_workers8來(lái)自 CPU 核心數(shù)測(cè)試pin_memoryTrue經(jīng)過(guò)對(duì)比實(shí)驗(yàn)確認(rèn)有 5% 左右的吞吐提升prefetch_factor4是因?yàn)橐曨l樣本加載時(shí)長(zhǎng)波動(dòng)大、需要更多緩沖。你完全可以基于自己機(jī)器的實(shí)測(cè)數(shù)據(jù)調(diào)整但每一行都應(yīng)該有我測(cè)過(guò)/我對(duì)比過(guò)的支撐。數(shù)據(jù)引擎的優(yōu)化沒(méi)有銀彈核心方法論就一句話讓每一層磁盤、CPU、內(nèi)存、GPU都別閑著也別堵著。先把瓶頸定位清楚再針對(duì)性地選方案——該加緩存加緩存該調(diào) worker 調(diào) worker該重新設(shè)計(jì)采樣器就重新設(shè)計(jì)。我個(gè)人的體會(huì)是80% 的項(xiàng)目靠自定義 Dataset DataLoader 參數(shù)調(diào)優(yōu)這兩板斧就能拿到顯著的吞吐提升剩下 20% 的高光譜、視頻、多源融合場(chǎng)景才需要?jiǎng)佑镁彺婕軜?gòu)、memmap 和自定義 sampler 這類進(jìn)階手段。最后再分享一個(gè)實(shí)測(cè)小技巧調(diào)參時(shí)不要同時(shí)改多個(gè)變量每次只動(dòng)一個(gè)參數(shù)、記錄一次吞吐量做成一個(gè)簡(jiǎn)單的對(duì)比表你會(huì)發(fā)現(xiàn)數(shù)據(jù)管道的最佳參數(shù)組合很快就浮出水面了比憑感覺(jué)亂試高效得多。 PyTorch數(shù)據(jù)管道全鏈路優(yōu)化：從Dataset到DataLoader的實(shí)踐
深度學(xué)習(xí)項(xiàng)目做到后期大家基本都有一個(gè)共識(shí)模型結(jié)構(gòu)再花哨優(yōu)化器再先進(jìn)只要數(shù)據(jù)管道卡脖子GPU就只能干瞪眼。我見(jiàn)過(guò)太多團(tuán)隊(duì)把時(shí)間耗在調(diào)網(wǎng)絡(luò)結(jié)構(gòu)上結(jié)果一跑訓(xùn)練nvidia-smi里 GPU 利用率只有 30%剩下 70% 時(shí)間全在等數(shù)據(jù)。PyTorch 的數(shù)據(jù)引擎——從自定義Dataset封裝、DataLoader參數(shù)配置到多源數(shù)據(jù)融合——看著不起眼卻是決定訓(xùn)練吞吐量上限的核心環(huán)節(jié)。這篇文章不聊模型只聊數(shù)據(jù)側(cè)的全鏈路工程實(shí)踐怎么把自定義數(shù)據(jù)封裝得既優(yōu)雅又高效怎么配置 DataLoader 才能把 GPU 喂飽以及多源異構(gòu)數(shù)據(jù)圖像、視頻、高光譜文件、傳感器時(shí)序怎么在同一個(gè)管道里平穩(wěn)融合。內(nèi)容基于我個(gè)人在多個(gè)項(xiàng)目里的實(shí)戰(zhàn)經(jīng)驗(yàn)適合已經(jīng)跑通 PyTorch 基礎(chǔ)訓(xùn)練流程、想進(jìn)一步優(yōu)化訓(xùn)練效率和數(shù)據(jù)集擴(kuò)展性的同學(xué)參考。1. 為什么性能瓶頸常常藏在你忽略的數(shù)據(jù)管道里1.1 GPU 利用率低的常見(jiàn)元兇先拋一個(gè)反直覺(jué)的結(jié)論大多數(shù)訓(xùn)練任務(wù)跑得慢不是模型算不動(dòng)而是數(shù)據(jù)送不到。幾年前我接手一個(gè)視頻分類項(xiàng)目模型是 3D CNN顯存占用正常但訓(xùn)練一個(gè) epoch 要 7 個(gè)多小時(shí)。后來(lái)用最簡(jiǎn)單的計(jì)時(shí)法測(cè)了一下發(fā)現(xiàn)每個(gè) step 里next(iter(loader))平均耗時(shí)接近 600 毫秒而 GPU 前向加反向只需要 200 毫秒。換句話說(shuō)有將近三分之二的時(shí)間GPU 都在空等。這種問(wèn)題在圖像分類里可能還不明顯因?yàn)閱螐垐D片讀取也就幾十毫秒但一旦涉及視頻幀序列、高光譜 HDR 文件、點(diǎn)云這類大塊頭數(shù)據(jù)IO 耗時(shí)會(huì)被急劇放大。很多人的第一反應(yīng)是換更好的硬盤但 SSD 和 HDD 的差異只能解決一部分問(wèn)題——如果瓶頸在 CPU 預(yù)處理解碼、歸一化、數(shù)據(jù)增強(qiáng)換硬盤根本沒(méi)意義。1.2 IO、CPU 與 GPU 之間的三角協(xié)作關(guān)系理解 PyTorch 數(shù)據(jù)引擎首先要建立生產(chǎn)-消費(fèi)的流水線思維。GPU 是消費(fèi)方它只認(rèn) CUDA 張量數(shù)據(jù)管道是生產(chǎn)方負(fù)責(zé)從磁盤讀原始數(shù)據(jù)、做預(yù)處理、組成 batch、搬到顯存。這三者之間的協(xié)作關(guān)系可以用一個(gè)例子類比GPU 是流水線上的裝配工CPU 是物料搬運(yùn)工磁盤是倉(cāng)庫(kù)。裝配工手速再快如果搬運(yùn)工一次只能搬一件、而且搬一件要跑一趟倉(cāng)庫(kù)整條流水線就廢了。PyTorch 的DataLoader之所以引入num_workers、prefetch_factor、pin_memory這些參數(shù)本質(zhì)上就是給這條流水線加多搬運(yùn)工和緩沖區(qū)。理解了這一層你就能明白調(diào)數(shù)據(jù)管道的本質(zhì)是在有限的 CPU 核數(shù)和內(nèi)存帶寬下盡可能讓數(shù)據(jù)準(zhǔn)備的速度追平甚至超過(guò) GPU 的消費(fèi)速度。1.3 先學(xué)會(huì)用兩分鐘定位瓶頸在動(dòng)手優(yōu)化之前先把問(wèn)題定位清楚。我的習(xí)慣是三步走訓(xùn)練時(shí)另開(kāi)一個(gè)終端跑watch -n 1 nvidia-smi看 GPU 利用率是否經(jīng)常掉到 80% 以下在訓(xùn)練循環(huán)里單獨(dú)計(jì)時(shí)start time.time(); data next(iter(loader)); print(time.time() - start)連續(xù)打幾十個(gè) step看數(shù)據(jù)加載耗時(shí)是否穩(wěn)定用 PyTorch 自帶的torch.utils.bottleneck跑一次它能粗粒度地區(qū)分 CPU、GPU、數(shù)據(jù)加載各自的耗時(shí)占比。這三步做完基本就能判斷是磁盤 IO 的鍋、CPU 預(yù)處理的鍋還是參數(shù)配置不當(dāng)?shù)腻?。下面各?jié)的內(nèi)容都是在定位完瓶頸之后才動(dòng)手的具體優(yōu)化手段。2. 自定義 Dataset三個(gè)核心方法與一個(gè)隱藏契約2.1__init__、__len__、__getitem__的標(biāo)準(zhǔn)寫(xiě)法和職責(zé)邊界自定義數(shù)據(jù)封裝聽(tīng)起來(lái)是老生常談但很多人對(duì)三個(gè)方法的職責(zé)邊界理解是模糊的。我的原則是__init__只負(fù)責(zé)記賬不負(fù)責(zé)干活。它應(yīng)該把所有樣本的索引、路徑、標(biāo)簽元數(shù)據(jù)預(yù)先整理成一個(gè)列表或 DataFrame而不是在這個(gè)階段就去讀數(shù)據(jù)。原因很簡(jiǎn)單__init__只在創(chuàng)建 Dataset 時(shí)執(zhí)行一次如果在這里做耗時(shí)操作不僅拖慢啟動(dòng)還會(huì)在訓(xùn)練過(guò)程中毫無(wú)意義地占用內(nèi)存。__len__返回樣本總數(shù)必須是一個(gè)確定且穩(wěn)定的數(shù)值。這個(gè)值不僅用于for epoch的輪數(shù)計(jì)算還直接影響Sampler的索引分配和分布式訓(xùn)練的DistributedSampler切分邏輯所以不能有任何歧義。__getitem__是真正的生產(chǎn)車間接收一個(gè)整數(shù)索引返回一個(gè)樣本。這里才是讀文件、解碼、預(yù)處理、數(shù)據(jù)增強(qiáng)該發(fā)生的地方。一個(gè)標(biāo)準(zhǔn)的圖像分類 Dataset 大致長(zhǎng)這樣class ImageDataset(Dataset): def __init__(self, file_list, label_list, transformNone): self.file_list file_list self.label_list label_list self.transform transform # 這里可以做一次數(shù)據(jù)校驗(yàn)比如檢查文件是否都存在 # 但不要讀文件內(nèi)容 def __len__(self): return len(self.file_list) def __getitem__(self, idx): img Image.open(self.file_list[idx]).convert(RGB) if self.transform: img self.transform(img) label self.label_list[idx] return img, label2.2 緩存策略什么數(shù)據(jù)適合在__init__里就加載進(jìn)內(nèi)存前面我說(shuō)__init__不要干活但有一個(gè)例外當(dāng)你的數(shù)據(jù)集總量小到能裝進(jìn)內(nèi)存時(shí)直接在__init__里把所有數(shù)據(jù)讀進(jìn)來(lái)然后__getitem__只做內(nèi)存切片效果會(huì)好到爆炸。比如 CIFAR-10、MNIST 這類數(shù)據(jù)集幾十萬(wàn)個(gè)樣本也就幾個(gè) GB完全可以全量駐留內(nèi)存把磁盤 IO 徹底消除。判斷標(biāo)準(zhǔn)很簡(jiǎn)單數(shù)據(jù)集壓縮后的大小 × 解壓后的膨脹系數(shù)如果不超過(guò)可用內(nèi)存的 50%就值得做全量緩存。需要注意的是解壓膨脹系數(shù)往往很驚人一個(gè) 100MB 的 JPEG 壓縮包解壓成 RGB 張量數(shù)組可能占 300MB 以上。所以穩(wěn)妥的做法是緩存原始字節(jié)而不是緩存解碼后的張量class CachedImageDataset(Dataset): def __init__(self, file_list, label_list, transformNone): self.raw_bytes [] for path in file_list: with open(path, rb) as f: self.raw_bytes.append(f.read()) self.label_list label_list self.transform transform def __getitem__(self, idx): img Image.open(io.BytesIO(self.raw_bytes[idx])).convert(RGB) ...這樣既省掉了重復(fù)的磁盤尋址又不會(huì)一次性把解碼后的矩陣全部堆在內(nèi)存里。實(shí)測(cè)中這種方式對(duì)幾萬(wàn)張小圖的數(shù)據(jù)集能帶來(lái) 5 到 10 倍的加載提速。2.3 分布式訓(xùn)練下 Dataset 的數(shù)據(jù)串?dāng)_風(fēng)險(xiǎn)單機(jī)單卡時(shí)Dataset 的索引隨便寫(xiě)都沒(méi)問(wèn)題。但一旦切到DistributedDataParallel事情就變了每個(gè)進(jìn)程會(huì)通過(guò)DistributedSampler拿到數(shù)據(jù)集的一個(gè)分片。如果你在 Dataset 里用了全局的隨機(jī)狀態(tài)比如給樣本打亂順序、隨機(jī)采樣又沒(méi)有做種子隔離就會(huì)出現(xiàn)不同進(jìn)程讀到重復(fù)樣本、或者某個(gè)進(jìn)程漏掉樣本的問(wèn)題。一個(gè)典型的錯(cuò)誤是在 Dataset 內(nèi)部維護(hù)一個(gè)隨機(jī)索引列表每次__getitem__的時(shí)候動(dòng)態(tài) shuffle。這會(huì)讓DistributedSampler的分片邏輯完全失效因?yàn)槟銓?shí)際返回的樣本和 sampler 給定的索引對(duì)不上。正確的做法是Dataset 保持純粹——給什么索引就返回什么樣本隨機(jī)性交給 DataLoader 的shuffleTrue或 sampler 來(lái)處理。這就是我說(shuō)的隱藏契約Dataset 的索引語(yǔ)義必須是確定的、可重入的任何隨機(jī)化都必須在外部完成。2.4 幾個(gè)常見(jiàn)的反模式排查過(guò)不少同事的代碼以下幾個(gè)坑出現(xiàn)頻率最高在__len__里做耗時(shí)統(tǒng)計(jì)比如每次調(diào)用都重新數(shù)一遍目錄下的文件這會(huì)讓每次 epoch 統(tǒng)計(jì)都變慢。__getitem__里修改全局變量多 worker 場(chǎng)景下主進(jìn)程和 worker 進(jìn)程是fork出來(lái)的全局變量的修改可能會(huì)造成數(shù)據(jù)錯(cuò)亂。返回 Python 對(duì)象而不是張量比如直接返回PIL.Image對(duì)象collate_fn會(huì)拿它沒(méi)辦法最后要么報(bào)錯(cuò)要么被迫走default_collate的慢路徑。這些都是寫(xiě)完 Dataset 之后自查的重點(diǎn)項(xiàng)。3. DataLoader 參數(shù)調(diào)優(yōu)worker 數(shù)量、prefetch 與 pin_memory 的配合邏輯3.1num_workers不是越大越好很多人以為num_workers是數(shù)據(jù)加載速度倍增器調(diào)得越大越好結(jié)果一開(kāi) 16 個(gè) workerCPU 直接被打滿GPU 利用率反而下降。原因在于worker 進(jìn)程之間需要把數(shù)據(jù)通過(guò) IPC進(jìn)程間通信傳回主進(jìn)程這個(gè)傳輸本身有開(kāi)銷同時(shí)每個(gè) worker 都有自己的數(shù)據(jù)預(yù)處理副本如果 CPU 核數(shù)不夠就會(huì)在進(jìn)程調(diào)度上產(chǎn)生大量競(jìng)爭(zhēng)。我的調(diào)參經(jīng)驗(yàn)是num_workers設(shè)置在 CPU 物理核數(shù)的一半到四分之三之間通常比較穩(wěn)。比如 8 核 CPU 就設(shè) 4 到 616 核就設(shè) 8 到 12。但這個(gè)數(shù)字和具體的數(shù)據(jù)格式強(qiáng)相關(guān)——如果每個(gè)樣本的解碼耗時(shí)很長(zhǎng)比如高光譜文件worker 多點(diǎn)反而有優(yōu)勢(shì)因?yàn)槠款i在 CPU 計(jì)算而不是 IPC 傳輸如果樣本本來(lái)就小worker 多了純粹是內(nèi)耗。判斷當(dāng)前 worker 數(shù)是否合適的辦法觀察訓(xùn)練日志里每個(gè) step 的耗時(shí)逐漸增大num_workers如果耗時(shí)不再下降甚至反彈說(shuō)明已經(jīng)到臨界點(diǎn)了。3.2pin_memory和prefetch_factor到底在解決什么問(wèn)題pin_memoryTrue的作用是把數(shù)據(jù)鎖頁(yè)到不可換出的內(nèi)存區(qū)域。理解這個(gè)問(wèn)題需要一點(diǎn) CUDA 基礎(chǔ)GPU 從 CPU 側(cè)拷貝數(shù)據(jù)時(shí)通常是異步 DMA 傳輸而 DMA 要求源地址在物理內(nèi)存中固定不變。如果數(shù)據(jù)在普通內(nèi)存里操作系統(tǒng)可能會(huì)把它換出到交換分區(qū)導(dǎo)致拷貝前需要先做一次同步頁(yè)表鎖定過(guò)程變慢。所以pin_memory本質(zhì)上是用鎖頁(yè)內(nèi)存換取了GPU 拷貝更快的收益。prefetch_factorPyTorch 1.7 之后可用則是控制每個(gè) worker 最多能預(yù)取多少個(gè) batch 放在隊(duì)列里。默認(rèn)值是 2對(duì)于普通場(chǎng)景夠用如果數(shù)據(jù)加載時(shí)間波動(dòng)比較大——比如有的視頻序列特別長(zhǎng)、有的特別短——可以適當(dāng)調(diào)高到 4 甚至 8給流水線多留一些余量避免偶爾的長(zhǎng)樣本加載把 GPU 餓死。這兩個(gè)參數(shù)配合使用的完整模板DataLoader( dataset, batch_size64, num_workers8, pin_memoryTrue, prefetch_factor4, persistent_workersTrue, # 多輪 epoch 時(shí)復(fù)用 worker避免重復(fù)創(chuàng)建 )persistent_workersTrue這個(gè)參數(shù)也值得單獨(dú)說(shuō)一下它讓 worker 在 epoch 結(jié)束后不退出下一次 epoch 繼續(xù)復(fù)用。省掉了反復(fù) fork worker 進(jìn)程的開(kāi)銷對(duì)訓(xùn)練幾十個(gè) epoch 的場(chǎng)景幫助很大。但注意它要求 Dataset 是可重入的不能依賴在__getitem__里做不可重復(fù)的操作。3.3 自定義collate_fn的場(chǎng)景變長(zhǎng)序列與多源數(shù)據(jù)拼接collate_fn是 DataLoader 中另一個(gè)被低估的組件。默認(rèn)的default_collate做的事情是把 batch 中每個(gè)樣本的張量堆疊成一個(gè)更大的張量。這個(gè)操作要求所有樣本的 shape 完全一致一旦遇到變長(zhǎng)序列、不同尺寸的圖像、或者樣本本身就是混合結(jié)構(gòu)它就會(huì)直接拋異常。遇到變長(zhǎng)數(shù)據(jù)最常見(jiàn)的處理思路是pad 到 batch 內(nèi)最大長(zhǎng)度。我自己在 NLP 和視頻特征序列項(xiàng)目里都這么干過(guò)def collate_variable_length(batch): images, seqs, labels zip(*batch) images torch.stack(images) max_len max([s.size(0) for s in seqs]) seqs_padded torch.zeros(len(seqs), max_len, seqs[0].size(1)) seq_lens torch.tensor([s.size(0) for s in seqs]) for i, s in enumerate(seqs): seqs_padded[i, :s.size(0)] s return images, seqs_padded, seq_lens, torch.tensor(labels)這里返回seq_lens是為了讓模型在讀取時(shí)能用mask回避 padding 帶來(lái)的無(wú)效計(jì)算。如果你用的是 Transformer 結(jié)構(gòu)這個(gè) mask 是必須的。3.4 通過(guò)Sampler控制采樣邏輯而不是改 Dataset很多人遇到某些類樣本太多、想均衡采樣的需求時(shí)第一反應(yīng)是往 Dataset 里加邏輯。但更干凈的做法是實(shí)現(xiàn)自定義Sampler。PyTorch 提供了WeightedRandomSampler可以按權(quán)重控制每個(gè)樣本被抽到的概率。對(duì)于類別不平衡的多源融合數(shù)據(jù)這個(gè)方案非常實(shí)用weights [1.0 / class_sample_count[label] for label in all_labels] sampler WeightedRandomSampler(weights, num_sampleslen(weights), replacementTrue)這里的關(guān)鍵是按類別頻率的倒數(shù)設(shè)權(quán)重讓少數(shù)類樣本在每輪 epoch 里被抽到的概率更高。我曾在類別比例 100:1 的多源數(shù)據(jù)集上用這個(gè)方案把少數(shù)類 F1 從 0.3 提到了 0.58效果比在 loss 里加權(quán)重更直接。4. 多源異構(gòu)數(shù)據(jù)的融合從硬拼到軟融的工程選型4.1 什么是多源融合為什么它不是簡(jiǎn)單的 concat多源融合multi-source fusion在機(jī)器人定位、遙感分析、多模態(tài)識(shí)別等場(chǎng)景里非常常見(jiàn)。它的難點(diǎn)不只是把數(shù)據(jù)拼起來(lái)而是不同類型的數(shù)據(jù)有完全不同的采樣頻率、數(shù)據(jù)維度、語(yǔ)義粒度和缺失模式。比如一個(gè)機(jī)器人定位項(xiàng)目里圖像是 30 FPSIMU 是 200 HzGPS 是 1 Hz三個(gè)源在時(shí)間軸上天然不對(duì)齊。如果只是簡(jiǎn)單地在樣本里 concat模型學(xué)到的往往是強(qiáng)勢(shì)模態(tài)的信息弱勢(shì)模態(tài)會(huì)被淹沒(méi)。所以在工程上我傾向于把多源融合分成三個(gè)層次來(lái)設(shè)計(jì)數(shù)據(jù)層融合在 Dataset 層面把多個(gè)源組織成統(tǒng)一樣本結(jié)構(gòu)特征層融合各模態(tài)先獨(dú)立編碼成特征向量再在模型內(nèi)部融合決策層融合各模態(tài)分別出預(yù)測(cè)結(jié)果再用加權(quán)投票或門控網(wǎng)絡(luò)綜合。數(shù)據(jù)引擎主要負(fù)責(zé)的是第一層——把多源數(shù)據(jù)送進(jìn)模型之前先解決好怎么組織、怎么對(duì)齊、怎么平衡這三個(gè)問(wèn)題。4.2 用ConcatDataset和自定義FusionDataset組織多源數(shù)據(jù)如果你的多源數(shù)據(jù)是不同來(lái)源但相同標(biāo)簽空間的樣本比如一部分是高光譜圖像、一部分是 RGB 圖像label 都是地物類別——最簡(jiǎn)單的方式是分別建 Dataset再用torch.utils.data.ConcatDataset合并。它內(nèi)部維護(hù)了各子數(shù)據(jù)集的樣本偏移量__len__會(huì)返回總和、__getitem__會(huì)根據(jù)索引自動(dòng)路由到正確的子集。優(yōu)點(diǎn)是零成本、天然支持缺點(diǎn)是無(wú)法在采樣時(shí)對(duì)源做權(quán)重控制——你只能用WeightedRandomSampler在全局樣本層面調(diào)整。如果你的多源數(shù)據(jù)是同一個(gè)樣本的不同模態(tài)比如一個(gè)樣本同時(shí)包含視頻幀和對(duì)應(yīng)的 IMU 時(shí)序那就需要自定義一個(gè)FusionDataset讓__getitem__返回一個(gè)字典結(jié)構(gòu)class FusionDataset(Dataset): def __init__(self, video_files, imu_files, labels): self.video_files video_files self.imu_files imu_files self.labels labels def __getitem__(self, idx): frames self._load_video_frames(self.video_files[idx]) # [T, H, W, C] imu_seq self._load_imu_sequence(self.imu_files[idx]) # [T, D] # 在這里做時(shí)間對(duì)齊、插值或分段采樣 frames, imu_seq self._align(frames, imu_seq) return {frames: frames, imu: imu_seq, label: self.labels[idx]}返回字典的好處是后續(xù)collate_fn可以精準(zhǔn)地按 key 處理不同張量——視頻幀要 stack 成五維張量IMU 序列要 pad 到等長(zhǎng)label 單獨(dú)收集。這種結(jié)構(gòu)比返回 tuple 清晰得多也更容易擴(kuò)展新的數(shù)據(jù)源。4.3 時(shí)間對(duì)齊不同采樣頻率數(shù)據(jù)的關(guān)鍵難題及處理策略時(shí)間對(duì)齊是多源融合里最容易被忽視、也最容易出錯(cuò)的地方。以視頻和 IMU 為例視頻每幀有固定的時(shí)間戳IMU 的每個(gè)采樣也有時(shí)間戳但兩者的時(shí)間戳起點(diǎn)往往不同采樣頻率也不同。我通常采用以低頻源為錨點(diǎn)的策略以視頻幀的時(shí)間戳作為標(biāo)準(zhǔn)對(duì) IMU 序列做線性插值取出與每幀時(shí)間戳對(duì)齊的 IMU 值def align_imu_to_frames(frame_timestamps, imu_timestamps, imu_values): aligned np.zeros((len(frame_timestamps), imu_values.shape[1])) for i, t in enumerate(frame_timestamps): idx np.searchsorted(imu_timestamps, t) if idx 0: aligned[i] imu_values[0] elif idx len(imu_timestamps): aligned[i] imu_values[-1] else: t0, t1 imu_timestamps[idx-1], imu_timestamps[idx] alpha (t - t0) / (t1 - t0 1e-8) aligned[i] imu_values[idx-1] * (1 - alpha) imu_values[idx] * alpha return aligned注意這里對(duì)邊界情況做了鉗制——時(shí)間戳落在 IMU 采樣范圍之外時(shí)用首尾值填充而不是報(bào)錯(cuò)。理由是傳感器數(shù)據(jù)在啟動(dòng)和停止階段經(jīng)常有缺失你需要在工程上容忍這些不完美。4.4 多源不平衡如何讓最后加進(jìn)來(lái)的那個(gè)源不白加多源融合項(xiàng)目最常見(jiàn)的結(jié)局是加了新數(shù)據(jù)源模型效果沒(méi)變甚至變差了。原因多數(shù)不是模型不行而是采樣策略把大量便宜的源比如圖像反復(fù)喂給模型而昂貴的源比如激光雷達(dá)點(diǎn)云出現(xiàn)的頻率太低模型根本學(xué)不到它的規(guī)律。解決思路有兩個(gè)層面。一是前面提到的WeightedRandomSampler在樣本層面調(diào)節(jié)各類別/各源的采樣概率二是引入源感知的 batch 組成。我正在用的一個(gè)方案是FusionDataset返回時(shí)額外給出source_id然后在collate_fn里統(tǒng)計(jì)當(dāng)前 batch 中各源的比例如果某個(gè)源占比低于預(yù)設(shè)閾值就用replacementTrue的方式強(qiáng)制多抽幾個(gè)該源的樣本補(bǔ)進(jìn)來(lái)。雖然這樣做會(huì)讓 batch 的分布和自然分布有偏差但在多源融合場(chǎng)景里保持每個(gè) batch 都能讓模型看到所有模態(tài)比嚴(yán)格遵循自然分布更重要。5. 高光譜、視頻等大文件數(shù)據(jù)的管道工程化處理5.1 高光譜 HDR/SPE 文件的讀取姿勢(shì)高光譜數(shù)據(jù)是數(shù)據(jù)引擎的極端壓力測(cè)試。它不像普通圖片那樣有統(tǒng)一的解碼庫(kù)HDR 和 SPE 文件往往來(lái)自特定儀器廠商格式五花八門。網(wǎng)上搜pytorch 處理高光譜 hdr 文件和 spe 文件能找到一堆零散腳本但真正可復(fù)用的經(jīng)驗(yàn)是下面幾條區(qū)分文件頭和數(shù)據(jù)區(qū)HDR/SPE 文件頭部包含波長(zhǎng)信息、采集參數(shù)等元數(shù)據(jù)數(shù)據(jù)區(qū)是原始光譜矩陣。讀取時(shí)通常用numpy.fromfile跳過(guò)文件頭再按(bands, height, width)的維度 reshape。不同廠商的字節(jié)序可能不同一般用f4或f4顯式指定大端/小端避免解析出亂碼。使用np.memmap做內(nèi)存映射整幅高光譜圖像可能達(dá)到幾百 MB 甚至幾個(gè) GB__getitem__里直接np.load會(huì)把內(nèi)存撐爆。用np.memmap可以按需讀取目標(biāo)波段或目標(biāo)區(qū)域def read_hdr_band(file_path, offset, band_index, height, width): mmap np.memmap(file_path, dtypef4, moder, offsetoffset, shape(bands, height, width)) band_data mmap[band_index].copy() # copy 出來(lái)再返回避免 memmap 句柄泄漏 return band_data波段采樣高光譜波段可能有幾百個(gè)但相鄰波段相關(guān)性極高。做分類任務(wù)時(shí)沒(méi)必要把所有波段都喂給模型通過(guò)主成分分析或互信息篩選出幾十個(gè)代表波段數(shù)據(jù)量直接降一個(gè)量級(jí)訓(xùn)練速度和效果反而更好。5.2 視頻動(dòng)作分類的數(shù)據(jù)管道設(shè)計(jì)視頻分類任務(wù)類似 UCF101 數(shù)據(jù)集實(shí)戰(zhàn)的數(shù)據(jù)管道比圖像復(fù)雜在多了一維時(shí)間。一個(gè) 10 秒的視頻均勻抽 16 幀每幀 224×224解碼成本是單張圖片的十幾倍。而且視頻解碼器OpenCV 的VideoCapture或者PyAV是出了名的又慢又不穩(wěn)定處理不好就是訓(xùn)練最大短板。我的做法是離線抽幀在線采樣訓(xùn)練前先把視頻均勻抽幀保存為單張 JPEG文件名帶上時(shí)間戳訓(xùn)練時(shí)Dataset只維護(hù)視頻 ID 幀索引列表__getitem__隨機(jī)選一個(gè)起始幀連續(xù)取 N 幀做數(shù)據(jù)增強(qiáng)。這樣做有幾個(gè)好處一是解碼成本被一次性攤銷二是隨機(jī)片段的采樣策略天然實(shí)現(xiàn)了時(shí)序數(shù)據(jù)增強(qiáng)三是因?yàn)閹呀?jīng)是單張圖片可以沿用圖像格式的成熟緩存策略。一個(gè)額外的坑是OpenCV 的VideoCapture在線讀取視頻時(shí)seek 精度不穩(wěn)定某些編碼格式下set(CAP_PROP_POS_FRAMES)跳轉(zhuǎn)到的幀不是目標(biāo)幀。離線抽幀徹底規(guī)避了這個(gè)問(wèn)題。5.3 雙隊(duì)列架構(gòu)讓數(shù)據(jù)準(zhǔn)備和模型訓(xùn)練互不阻塞當(dāng)單進(jìn)程的DataLoader無(wú)論如何也喂不飽 GPU 時(shí)就該考慮更極致的方案了。我在處理超大數(shù)據(jù)集時(shí)用過(guò)雙隊(duì)列架構(gòu)一個(gè)生產(chǎn)者進(jìn)程負(fù)責(zé)從磁盤批量讀取原始數(shù)據(jù)放進(jìn)內(nèi)存隊(duì)列另一個(gè)進(jìn)程池負(fù)責(zé)從隊(duì)列取數(shù)據(jù)做預(yù)處理和增強(qiáng)再通過(guò)DataLoader的 worker 消費(fèi)。這個(gè)架構(gòu)本質(zhì)上是在 DataLoader 前面再加一級(jí)緩存層。queue multiprocessing.Queue(maxsize8) # producer: 從磁盤讀原始字節(jié) def producer(file_batch): for path in file_batch: with open(path, rb) as f: queue.put(f.read())不過(guò)說(shuō)實(shí)話這個(gè)方案工程復(fù)雜度較高多數(shù)場(chǎng)景用不上。只有當(dāng)num_workers已經(jīng)拉滿、CPU 占用也跑滿了、但 GPU 還是在等數(shù)據(jù)時(shí)才值得考慮引入。普通項(xiàng)目先把 DataLoader 的參數(shù)調(diào)優(yōu)做扎實(shí)收益已經(jīng)非常可觀。5.4 避免重復(fù) IO緩存增強(qiáng) vs 緩存原始數(shù)據(jù)最后聊一個(gè)容易被忽略的優(yōu)化點(diǎn)數(shù)據(jù)增強(qiáng)應(yīng)該放在哪個(gè)環(huán)節(jié)。很多教程把 transform 放在__getitem__里但這意味著每次 epoch 都要重新解碼原始數(shù)據(jù)、重新做增強(qiáng)。如果你的數(shù)據(jù)量不大但訓(xùn)練輪數(shù)很多可以考慮把已增強(qiáng)的結(jié)果緩存到磁盤或內(nèi)存用法是一輪 epoch 生成一次增強(qiáng)副本后面幾輪直接復(fù)用。不同方案對(duì)比方案內(nèi)存占用IO 開(kāi)銷數(shù)據(jù)多樣性適用場(chǎng)景每次完整解碼增強(qiáng)低高高數(shù)據(jù)集大、訓(xùn)練輪數(shù)少緩存原始字節(jié)中中高每次重新增強(qiáng)數(shù)據(jù)集可容納進(jìn)內(nèi)存緩存增強(qiáng)后張量高低低每輪相同數(shù)據(jù)集小、訓(xùn)練輪數(shù)多memmap 按需讀取低低中單樣本文件極大高光譜/視頻注意緩存增強(qiáng)后張量會(huì)導(dǎo)致數(shù)據(jù)多樣性下降容易過(guò)擬合所以通常只在強(qiáng)約束下用比如做知識(shí)蒸餾的固定輸入。我的習(xí)慣是優(yōu)先保證每輪能重新增強(qiáng)只有確認(rèn)內(nèi)存夠用、且增強(qiáng)不是關(guān)鍵因素時(shí)才退回到緩存方案。6. 從環(huán)境到代碼數(shù)據(jù)管道的性能剖析與隱性故障排查6.1 數(shù)據(jù)管道性能剖析的完整排查鏈路前面說(shuō)過(guò)定位瓶頸的三步法這里展開(kāi)講具體操作。我先跑一個(gè)縮短版訓(xùn)練比如 50 個(gè) step在這個(gè)訓(xùn)練里同時(shí)記錄 GPU 利用率和數(shù)據(jù)加載耗時(shí)。拿到數(shù)據(jù)后按下面的順序排查磁盤 IO 是否打滿用iostat -x 1看%util如果長(zhǎng)期接近 100%說(shuō)明瓶頸在磁盤優(yōu)先上緩存或memmap。CPU 是否打滿用htop看各核占用如果所有核都 100% 而 GPU 空閑說(shuō)明預(yù)處理太慢需要優(yōu)化解碼邏輯或增加 worker如果核數(shù)有富余。IPC 傳輸是否成為瓶頸如果 worker 很多、CPU 沒(méi)打滿但 GPU 還是餓很可能數(shù)據(jù)在從 worker 傳回主進(jìn)程的過(guò)程中耗時(shí)過(guò)大。這時(shí)可以試試減小傳輸數(shù)據(jù)量比如在 worker 里提前把圖像轉(zhuǎn)成小尺寸張量或減少num_workers。對(duì)這個(gè)流程我曾吃過(guò)一次大虧有個(gè)項(xiàng)目數(shù)據(jù)在機(jī)械硬盤上num_workers8CPU 完全沒(méi)有打滿但每個(gè) step 的加載時(shí)間穩(wěn)定在 1.2 秒。用iostat一查磁盤%util95% 以上——多個(gè) worker 同時(shí)隨機(jī)讀不同位置機(jī)械硬盤的尋道被反復(fù)放大8 個(gè) worker 反而比 2 個(gè)更慢。換成 SSD 后問(wèn)題直接消失。這說(shuō)明任何調(diào)優(yōu)手段都必須先確認(rèn)瓶頸在哪一層不能盲目堆參數(shù)。6.2 幾個(gè)隱蔽性極強(qiáng)的隱性錯(cuò)誤有些問(wèn)題不會(huì)立刻報(bào)錯(cuò)但會(huì)悄悄污染訓(xùn)練結(jié)果排查起來(lái)非常頭疼。我把自己踩過(guò)、也幫別人排查過(guò)的列幾個(gè)重點(diǎn)Worker 隨機(jī)種子未隔離多個(gè) worker 用同一個(gè)隨機(jī)種子做數(shù)據(jù)增強(qiáng)時(shí)一個(gè) batch 里會(huì)出現(xiàn)多張一模一樣的增強(qiáng)圖模型學(xué)到的特征分布會(huì)被扭曲。解決方式是在 Dataset 里依靠 PyTorch 對(duì)每個(gè) worker 自動(dòng)設(shè)置不同隨機(jī)種子PyTorch 1.x 之后會(huì)為每個(gè) worker 設(shè)置獨(dú)立基礎(chǔ)種子但如果你在外部用np.random.seed手動(dòng)固定了全局種子這個(gè)機(jī)制會(huì)被繞過(guò)。注意不要在數(shù)據(jù)讀取代碼里手動(dòng)重設(shè)全局種子。訓(xùn)練集和驗(yàn)證集共享緩存如果 Dataset 的緩存字典用的是全局變量訓(xùn)練和驗(yàn)證兩個(gè) loader 指向同一個(gè) Dataset 時(shí)驗(yàn)證階段可能讀到訓(xùn)練時(shí)緩存的數(shù)據(jù)。訓(xùn)練和驗(yàn)證必須使用獨(dú)立的 Dataset 實(shí)例或者明確區(qū)分緩存命名空間。文件句柄泄漏使用Image.open、np.memmap、open()時(shí)忘記關(guān)閉或copy長(zhǎng)期運(yùn)行會(huì)導(dǎo)致句柄數(shù)暴漲最后整個(gè)進(jìn)程崩潰。我之前處理高光譜數(shù)據(jù)時(shí)就是因?yàn)?memmap 句柄沒(méi)有主動(dòng)關(guān)閉訓(xùn)練到第 40 個(gè) epoch 時(shí)系統(tǒng)報(bào)Too many open files。在__getitem__里對(duì)memmap做.copy()并顯式del是簡(jiǎn)單可靠的規(guī)避方法。shuffle 和斷點(diǎn)續(xù)訓(xùn)的配合問(wèn)題如果訓(xùn)練中斷后從 checkpoint 恢復(fù)但 DataLoader 的shuffleTrue沒(méi)有固定generator的種子恢復(fù)后的數(shù)據(jù)順序會(huì)和中斷前完全不同學(xué)習(xí)率調(diào)度和最佳 checkpoint 的對(duì)比全部失效。斷點(diǎn)續(xù)訓(xùn)時(shí)必須固定DataLoader的generatorg torch.Generator() g.manual_seed(42) DataLoader(..., shuffleTrue, generatorg)驗(yàn)證集 loader 的 worker 數(shù)過(guò)高驗(yàn)證階段通常不需要那么多數(shù)據(jù)增強(qiáng)但很多人會(huì)把訓(xùn)練的超參原樣套給驗(yàn)證 loader導(dǎo)致驗(yàn)證每個(gè) epoch 耗時(shí)和訓(xùn)練幾乎一樣。驗(yàn)證 loader 一般設(shè)num_workers2或干脆num_workers0、shuffleFalse、prefetch_factor保持默認(rèn)即可。6.3 PyTorch 與 CUDA 環(huán)境配套對(duì)數(shù)據(jù)加載的影響環(huán)境配置雖然不直接屬于數(shù)據(jù)引擎編碼范疇但數(shù)據(jù)張量從 CPU 搬到 GPU 的過(guò)程強(qiáng)依賴 CUDA 版本和 PyTorch 的配套情況。熱詞里頻繁出現(xiàn)pytorch 安裝教程 gpu、cuda 版本、anaconda 配置 pytorch 環(huán)境這類搜索背后其實(shí)是同一個(gè)痛點(diǎn)裝了 PyTorch 的 CPU 版本pin_memory和.cuda()的加速路徑都不能生效數(shù)據(jù)搬移會(huì)經(jīng)過(guò)一條很慢的共享內(nèi)存路徑。我的建議是無(wú)論偏不偏重 GPU都用torch.cuda.is_available()在項(xiàng)目啟動(dòng)時(shí)就做一次顯式檢查并且打印當(dāng)前 CUDA 版本assert torch.cuda.is_available(), CUDA not available, data loading will fall back to CPU-only path print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))環(huán)境是否配套用一句命令就能確認(rèn)python -c import torch; print(torch.__version__, torch.version.cuda, torch.backends.cudnn.version())如果輸出里torch.version.cuda和你在系統(tǒng)里裝的驅(qū)動(dòng)版本差距過(guò)大pin_memory的異步傳輸特性可能無(wú)法正常發(fā)揮訓(xùn)練速度會(huì)比預(yù)期慢一截。另外Windows 上num_workers的 fork 機(jī)制和 Linux 不一樣Windows 默認(rèn)使用spawnworker 會(huì)在每個(gè)進(jìn)程里重新導(dǎo)入整個(gè)模塊如果 Dataset 里涉及大量全局初始化會(huì)顯著拖慢啟動(dòng)——所以生產(chǎn)環(huán)境我強(qiáng)烈建議用 Linux。6.4 一個(gè)完整的自定義 DataLoader 配置參考把前面所有討論落到一個(gè)可復(fù)用的配置模板上。假設(shè)項(xiàng)目是多源融合場(chǎng)景包含視頻片段和 IMU 時(shí)序數(shù)據(jù)量中等偏大transform torchvision.transforms.Compose([ torchvision.transforms.RandomResizedCrop(224), torchvision.transforms.RandomHorizontalFlip(), torchvision.transforms.ToTensor(), torchvision.transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) train_dataset FusionDataset(video_files, imu_files, labels, transformtransform) train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers8, # 物理核數(shù)的一半左右 pin_memoryTrue, # GPU 環(huán)境下必開(kāi) prefetch_factor4, # 數(shù)據(jù)波動(dòng)較大預(yù)取多一些 persistent_workersTrue, # 多輪 epoch 復(fù)用 worker collate_fncollate_variable_length, drop_lastTrue, # 避免最后不完整的 batch 影響 BN 統(tǒng)計(jì) )這些參數(shù)不是拍腦袋定的num_workers8來(lái)自 CPU 核心數(shù)測(cè)試pin_memoryTrue經(jīng)過(guò)對(duì)比實(shí)驗(yàn)確認(rèn)有 5% 左右的吞吐提升prefetch_factor4是因?yàn)橐曨l樣本加載時(shí)長(zhǎng)波動(dòng)大、需要更多緩沖。你完全可以基于自己機(jī)器的實(shí)測(cè)數(shù)據(jù)調(diào)整但每一行都應(yīng)該有我測(cè)過(guò)/我對(duì)比過(guò)的支撐。數(shù)據(jù)引擎的優(yōu)化沒(méi)有銀彈核心方法論就一句話讓每一層磁盤、CPU、內(nèi)存、GPU都別閑著也別堵著。先把瓶頸定位清楚再針對(duì)性地選方案——該加緩存加緩存該調(diào) worker 調(diào) worker該重新設(shè)計(jì)采樣器就重新設(shè)計(jì)。我個(gè)人的體會(huì)是80% 的項(xiàng)目靠自定義 Dataset DataLoader 參數(shù)調(diào)優(yōu)這兩板斧就能拿到顯著的吞吐提升剩下 20% 的高光譜、視頻、多源融合場(chǎng)景才需要?jiǎng)佑镁彺婕軜?gòu)、memmap 和自定義 sampler 這類進(jìn)階手段。最后再分享一個(gè)實(shí)測(cè)小技巧調(diào)參時(shí)不要同時(shí)改多個(gè)變量每次只動(dòng)一個(gè)參數(shù)、記錄一次吞吐量做成一個(gè)簡(jiǎn)單的對(duì)比表你會(huì)發(fā)現(xiàn)數(shù)據(jù)管道的最佳參數(shù)組合很快就浮出水面了比憑感覺(jué)亂試高效得多。