據(jù)PyTorch加載實戰(zhàn):從Dataset設(shè)計到Dataloader避坑指南)
簡介面向PyTorch開發(fā)者與高光譜圖像處理研究者提供一套基于DataLoader加載高光譜數(shù)據(jù)的完整代碼實現(xiàn)解決HSI光譜通道眾多、數(shù)據(jù)量大所帶來的加載效率與內(nèi)存壓力問題。壓縮包為7z格式共7個文件包含3個Python腳本dataloader.py、train.py、utils.py、2個mat數(shù)據(jù)文件及2個pyc緩存文件整體僅5.69MB結(jié)構(gòu)清晰便于快速移植和二次開發(fā)。腳本從自定義Dataset類入手演示了高光譜數(shù)據(jù)的讀取、歸一化、隨機打亂和批量送入流程并覆蓋num_workers多線程加速、pin_memory緩存、自定義collate_fn合并多維樣本等關(guān)鍵技巧同時提供IndianPines數(shù)據(jù)集與訓(xùn)練腳本支持本地直接運行方便對比不同batch_size、預(yù)處理和采樣策略對訓(xùn)練效果的影響。與普通RGB數(shù)據(jù)不同高光譜圖像單樣本維度更高這套代碼中的配置方式能有效規(guī)避內(nèi)存不足和加載瓶頸幫助用戶搭起可復(fù)用的數(shù)據(jù)管道。已有612人學(xué)習(xí)下載適合剛接觸高光譜與PyTorch的學(xué)生、科研人員也可作為現(xiàn)有項目數(shù)據(jù)加載模塊的優(yōu)化參考。1. 高光譜數(shù)據(jù)用 PyTorch Dataloader 加載先解決“一維數(shù)組當圖像”的問題做高光譜分類的同學(xué)八成都有過這種經(jīng)歷從 .mat 或 .h5 里導(dǎo)出的數(shù)據(jù)明明是三維立方體——長、寬、幾十上百個波段但打印張量形狀時卻看到[1, 102, 512]甚至[102, 512]直接拿給 CNN 用第一層卷積就炸了維度。更隱蔽的問題是高光譜數(shù)據(jù)通常不是一張大圖而是成千上萬個像素級樣本每個樣本是一個(波段數(shù),)的向量標簽是類別編號。這種數(shù)據(jù)沒法像 ImageFolder 那樣直接喂給 PyTorch 的默認 Dataloader必須自己寫 Dataset 類來封裝切片、歸一化、樣本配對。這篇筆記要解決的就是這件事把高光譜原始數(shù)據(jù)轉(zhuǎn)換成 PyTorch Dataloader 能直接吃的格式并給出完整的 Dataset 實現(xiàn)、參數(shù)設(shè)置和我在實際跑實驗時踩過的幾個坑。適合剛接觸高光譜分類、被數(shù)據(jù)加載卡住的研究生也適合已經(jīng)寫了簡陋加載代碼但想理清批處理邏輯的開發(fā)者。這里不廢話直接上可復(fù)現(xiàn)的代碼路徑。2. 高光譜數(shù)據(jù)的組織方式與 Dataset 設(shè)計從三維立方體到像素級樣本2.1 高光譜數(shù)據(jù)與普通圖像的三個關(guān)鍵差異高光譜數(shù)據(jù)通常是一個三維數(shù)組形狀為(height, width, bands)比如著名的模擬項目X數(shù)據(jù)集是(145, 145, 200)某跨平臺系統(tǒng)常用數(shù)據(jù)集是(512, 217, 103)。它和 RGB 圖像最大的區(qū)別在于波段數(shù)不是 3 而是幾十上百而且波段之間往往有較強的相關(guān)性甚至存在噪聲波段。普通圖像加載后可以直接做歸一化(H, W, 3)但高光譜數(shù)據(jù)如果直接整圖喂進網(wǎng)絡(luò)會帶來兩個問題一是顯存爆炸(512, 217, 103)的 float32 張量一次前向就要占用約 45 MB還沒算中間激活二是分類任務(wù)通常不是對整圖做語義分割而是對中心像素打標簽——每個像素對應(yīng)一個地物類別訓(xùn)練樣本是像素及其鄰域。所以常見的做法不是把整幅高光譜圖當作一個樣本而是把圖像切分成像素級樣本。每個樣本以某個像素為中心取一個固定大小的鄰域窗口比如window_size11即每個樣本的形狀是(11, 11, bands)。如果只取單個像素就是(bands,)的向量但那樣丟失了空間上下文信息分類精度會大打折扣。通常用鄰域窗口來構(gòu)造樣本這樣既能利用空間特征又能讓樣本數(shù)量達到幾萬個足夠訓(xùn)練小型網(wǎng)絡(luò)。在 Dataset 設(shè)計上PyTorch 要求我們實現(xiàn)__len__和__getitem__兩個方法。高光譜 Dataset 的核心邏輯是先在初始化階段加載原始數(shù)據(jù)、歸一化、確定有效樣本的索引列表然后在__getitem__里根據(jù)索引計算出中心像素坐標裁剪鄰域窗口返回(window, label)對。這里要注意裁剪窗口必然涉及到邊界處理。窗口越界時不能簡單丟棄因為數(shù)據(jù)集邊緣像素數(shù)量不少全丟了會損失大量樣本。我一般會用pad方式處理邊界或者只取那些窗口完整的像素作為有效樣本。2.2 一個可直接用的高光譜 Dataset 實現(xiàn)下面這個代碼是從我實際跑過的模擬項目X分類實驗里抽出來的去掉了與模型相關(guān)的部分只保留數(shù)據(jù)加載模塊。文件結(jié)構(gòu)上我把 Dataset 寫在hyperspectral_dataset.py里方便后續(xù)復(fù)用。import numpy as np import torch from torch.utils.data import Dataset from sklearn.preprocessing import StandardScaler import h5py class HyperspectralDataset(Dataset): def __init__(self, data_path, label_path, window_size11, use_pcaTrue, n_components30, trainFalse): data_path: .h5 或 .npy 文件路徑數(shù)據(jù)形狀為 (height, width, bands) label_path: .npy 或 .mat 標簽文件形狀為 (height, width) window_size: 鄰域窗口大小建議奇數(shù)如 11, 13, 15 use_pca: 是否對波段降維True 會大幅減少顯存占用 train: 是否返回原始索引便于后續(xù)劃分訓(xùn)練/驗證/測試 # 讀取高光譜數(shù)據(jù) if data_path.endswith(.h5): with h5py.File(data_path, r) as f: # 根據(jù)實際 key 調(diào)整常見的有 data 或 X self.data f[data][:] else: self.data np.load(data_path) self.labels np.load(label_path) self.window_size window_size self.pad window_size // 2 self.use_pca use_pca # 數(shù)據(jù)歸一化按波段做標準化避免某些波段數(shù)值量級過大 scaler StandardScaler() h, w, b self.data.shape # 把每個波段拉平逐波段標準化 data_flatten self.data.reshape(-1, b) data_flatten scaler.fit_transform(data_flatten) self.data data_flatten.reshape(h, w, b) # 可選 PCA 降維保留主要波段信息 if use_pca: from sklearn.decomposition import PCA pca PCA(n_componentsn_components) h, w, b self.data.shape data_flatten self.data.reshape(-1, b) data_pca pca.fit_transform(data_flatten) self.data data_pca.reshape(h, w, n_components) self.bands n_components else: self.bands b # 統(tǒng)計有效樣本坐標 # 假設(shè)標簽中 0 表示背景需要去掉具體類別范圍根據(jù)數(shù)據(jù)調(diào)整 self.valid_indices [] for i in range(self.pad, self.labels.shape[0] - self.pad): for j in range(self.pad, self.labels.shape[1] - self.pad): if self.labels[i, j] 0: # 非背景 self.valid_indices.append((i, j)) self.train train def __len__(self): return len(self.valid_indices) def __getitem__(self, idx): i, j self.valid_indices[idx] # 裁剪鄰域窗口形狀: (window_size, window_size, bands) window self.data[i - self.pad:i self.pad 1, j - self.pad:j self.pad 1] # 轉(zhuǎn)成 (bands, window_size, window_size) 以符合 PyTorch NCHW window np.transpose(window, (2, 0, 1)) # 類別標簽 label self.labels[i, j] - 1 # 若標簽從 1 開始轉(zhuǎn)為 0 基 return torch.from_numpy(window).float(), torch.tensor(label, dtypetorch.long)這段代碼里有幾個關(guān)鍵參數(shù)和邏輯值得單獨說明。第一是use_pca參數(shù)。高光譜原始波段數(shù)可能有 200 個直接做 3D 卷積時第一層參數(shù)量就會很大而且相鄰波段高度相關(guān)。PCA 降維到 30 維后顯存占用和訓(xùn)練速度都會有明顯改善。但注意PCA 必須在全數(shù)據(jù)集上擬合不能在__getitem__里逐樣本做否則會出現(xiàn)信息泄漏。第二是邊界處理。我用的是只保留窗口完整像素的策略也就是從pad位置開始遍歷去掉四周邊緣。這樣每個樣本都是完整窗口不會引入 padding 噪聲。代價是邊緣像素樣本丟失但對于大部分高光譜數(shù)據(jù)集有效樣本數(shù)量仍然足夠。第三是標簽處理。很多公開數(shù)據(jù)集的標簽背景為 0而實際類別從 1 開始。在__getitem__里統(tǒng)一減 1讓類別變成 0 基這樣可以直接用于nn.CrossEntropyLoss。第四是__getitem__返回的是torch.Tensordtype 為float32和long。這是 PyTorch 數(shù)據(jù)加載的最基本要求如果返回 numpy 數(shù)組后續(xù)拼接 batch 時會因為類型不一致報錯。2.3 構(gòu)建 Dataloader 時的關(guān)鍵參數(shù)設(shè)置有了 DatasetDataloader 的設(shè)置直接決定訓(xùn)練速度和穩(wěn)定性。下面是我常用的配置from torch.utils.data import DataLoader, random_split dataset HyperspectralDataset( data_pathPaviaU_data.h5, label_pathPaviaU_labels.npy, window_size11, use_pcaTrue, n_components30 ) # 按比例劃分訓(xùn)練集和測試集保持類別分布大致均勻 train_len int(len(dataset) * 0.8) test_len len(dataset) - train_len train_dataset, test_dataset random_split(dataset, [train_len, test_len]) train_loader DataLoader( train_dataset, batch_size128, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) test_loader DataLoader( test_dataset, batch_size256, shuffleFalse, num_workers4, pin_memoryTrue, drop_lastFalse )這里重點講三個容易被忽視的參數(shù)shuffle、drop_last、pin_memory。shuffle在訓(xùn)練集必須為True否則每個 epoch 內(nèi)樣本順序完全一致模型會逐漸記住數(shù)據(jù)順序?qū)е买炞C集精度虛高。測試和驗證集則必須為False保證評估時每個樣本都被看到且順序固定。drop_lastTrue是很多人在小數(shù)據(jù)集上翻車的隱性因素。如果訓(xùn)練樣本總數(shù)不是 batch_size 的整數(shù)倍最后一個 batch 不足 128批歸一化層會在該 batch 上統(tǒng)計錯誤的均值和方差并且如果網(wǎng)絡(luò)中有固定形狀的輸入限制小 batch 也可能報錯。我一般會打開drop_lastTrue但要注意最后一個 batch 被舍棄后實際參與訓(xùn)練的樣本數(shù)比總數(shù)少一點這在樣本量幾萬時完全無感。pin_memoryTrue適合 GPU 訓(xùn)練它讓數(shù)據(jù)加載器使用頁鎖定內(nèi)存?zhèn)鬏數(shù)?GPU 的速度會有明顯提升。代價是占用更多主機內(nèi)存如果你的機器內(nèi)存只有 8 GB 且數(shù)據(jù)較大可以先關(guān)掉。num_workers也不是越多越好。實際經(jīng)驗是當數(shù)據(jù)預(yù)處理只有裁剪和轉(zhuǎn)置時num_workers4通常就能打滿數(shù)據(jù)加載帶寬如果設(shè)為 8 以上Windows 上容易出現(xiàn)多進程撕裂錯誤Linux 上則可能因為共享內(nèi)存不足報錯。我一般先在 2 和 4 之間試再往上加。3. 歸一化與數(shù)據(jù)增強的四種常見做法讓模型不被波段量級帶偏3.1 為什么不能只做簡單的 Min-Max 歸一化高光譜數(shù)據(jù)的原始數(shù)值范圍極其混亂。有些波段的反射率在 0 到 1 之間有些經(jīng)過輻射定標后可能到幾千尤其是某些噪聲波段。如果直接把這些數(shù)據(jù)喂給卷積網(wǎng)絡(luò)梯度會被大數(shù)值波段的特征主導(dǎo)小數(shù)值波段的信息基本被淹沒。常見的 Min-Max 歸一化是把每個波段壓到 [0,1]公式是(x - min) / (max - min)。但高光譜數(shù)據(jù)中個別波段可能存在極端的離群值比如傳感器壞點產(chǎn)生的尖峰這會讓 Min-Max 后的絕大多數(shù)數(shù)值被壓縮到 0 附近動態(tài)范圍反而被破壞。更穩(wěn)妥的做法是按波段做 Z-score 標準化也就是對每個波段分別減去均值、除以標準差。這樣處理后每個波段的數(shù)值中心為 0尺度統(tǒng)一而且不受離群值影響。我在 Dataset 里用的StandardScaler就是這么做的。注意標準化參數(shù)必須從訓(xùn)練集的每個波段上統(tǒng)計不能在整個數(shù)據(jù)集上統(tǒng)計后直接用到驗證集否則會有數(shù)據(jù)泄漏。3.2 2D 卷積下適合的鄰域窗口參數(shù)如果你的分類模型是 2D CNN也就是把每個波段當作一個通道那么輸入形狀是(bands, window_size, window_size)。窗口大小的選擇直接影響分類精度和訓(xùn)練耗時。窗口太小比如 3×3空間上下文太少模型只能看到中心像素附近的一小圈分類精度通常很低窗口太大比如 25×25雖然上下文豐富但每個樣本包含的冗余空間信息太多不僅計算量大還可能讓模型過擬合背景紋理。我跑過一組對比實驗窗口從 5 到 25間隔 2最終在模擬項目X數(shù)據(jù)集上 11×11 附近達到精度峰值繼續(xù)增大精度不再提升但訓(xùn)練時間幾乎線性增長。所以默認 11×11 是一個平衡點。如果用的是 3D 卷積輸入形狀變成(1, window_size, window_size, bands)或(bands, window_size, window_size, 1)這時窗口大小對計算量的影響更大我通常會降為 9×9 并配合 PCA 降維到 20 維左右。3.3 數(shù)據(jù)增強旋轉(zhuǎn)和翻轉(zhuǎn)要留意標簽一致性高光譜樣本的標簽是中心像素的類別鄰域窗口中的其他像素類別不一定相同。如果做隨機旋轉(zhuǎn)或翻轉(zhuǎn)中心像素不會變但空間特征的方向會改變這在物理意義上是合理的——地物類別不應(yīng)當因為傳感器視角旋轉(zhuǎn)而改變。所以高光譜數(shù)據(jù)增強比普通圖像更安全90 度旋轉(zhuǎn)、水平翻轉(zhuǎn)、垂直翻轉(zhuǎn)都不會破壞標簽語義。def augment(window): # window: (bands, h, w) if torch.rand(1) 0.5: window torch.flip(window, dims[2]) # 水平翻轉(zhuǎn) if torch.rand(1) 0.5: window torch.flip(window, dims[1]) # 垂直翻轉(zhuǎn) k torch.randint(0, 4, (1,)).item() if k 0: window torch.rot90(window, k, dims[1, 2]) return window這段增強代碼放在__getitem__中只在訓(xùn)練集上啟用驗證集不做。實現(xiàn)上有兩個注意點一是翻轉(zhuǎn)維度必須是空間維不能是波段維二是rot90的dims必須明確指定為[1,2]否則默認對最后兩維操作在高光譜數(shù)據(jù)上恰好是空間維但寫成顯式更安全。增強對精度的提升幅度因數(shù)據(jù)集而異。有些數(shù)據(jù)集本身類別內(nèi)樣本足夠多增強反而可能引入過度變形有些類別只有幾百個樣本增強能明顯緩解過擬合。我一般只在訓(xùn)練樣本少于每類 500 個時才開啟增強。3.4 波段選擇與 PCA 降維的取舍PCA 并不是所有場景都適用。如果你只是想快速驗證網(wǎng)絡(luò)結(jié)構(gòu)PCA 降維到 30 維幾乎不影響精度但訓(xùn)練速度能快 5 倍以上。如果追求最高分類精度PCA 可能會剔除一些對類別區(qū)分有微弱貢獻的波段此時可以試試不降維直接使用全波段。我實際遇到過一個案例某圖像處理 Demo 使用的數(shù)據(jù)只有 103 波段但其中 20 多個波段噪聲很大。直接全波段訓(xùn)練時驗證集精度在 94% 左右徘徊把噪聲波段手動剔除后精度提升到了 96.5%。這說明 PCA 是自動降噪手段但手動剔除已知噪聲波段在某些情況下效果更佳。如果你對數(shù)據(jù)集的波段質(zhì)量有先驗知識建議先用envi或matplotlib逐波段可視化把明顯是水汽吸收帶或傳感器故障帶的波段刪掉再做 PCA。4. 避坑指南高光譜 Dataloader 的六個常見陷阱4.1 張量形狀變成(1, bands, h, w)導(dǎo)致模型輸入不匹配現(xiàn)象訓(xùn)練時第一次前向就報錯提示期望輸入是 4 維但得到 3 維或第一層卷積輸出尺寸不對。原因數(shù)據(jù)讀取后維度是(h, w, bands)直接torch.from_numpy(data).unsqueeze(0)會在最前面加一維變成(1, h, w, bands)模型的Conv2d期望(batch, channels, height, width)此時通道維在最后完全對不上。解決嚴格按np.transpose(data, (2, 0, 1))把波段維挪到最前再加 batch 維。也就是在__getitem__中輸出(bands, h, w)Dataloader 自動堆疊成(batch, bands, h, w)。4.2 標簽中有背景類但沒過濾導(dǎo)致類別數(shù)比實際多現(xiàn)象模型輸出類別數(shù)是 10但真實類別只有 9 個精度始終很低訓(xùn)練損失下降緩慢。原因高光譜數(shù)據(jù)集的標簽矩陣里背景像素往往標為 0在做樣本索引時沒有過濾掉導(dǎo)致 background 被當成第 0 類參與訓(xùn)練而實際類別從 1 開始模型被迫學(xué)習(xí)一個不存在的背景類。解決在統(tǒng)計valid_indices時加上if self.labels[i, j] 0并且__getitem__中l(wèi)abel self.labels[i, j] - 1。如果某些數(shù)據(jù)集的類別從 0 開始比如某公開模擬數(shù)據(jù)集的類別編號是 0 到 8其中 0 是背景那就去掉減 1 操作只保留大于 -1 的索引即可。4.3num_workers設(shè)置過大導(dǎo)致程序卡死或內(nèi)存爆炸現(xiàn)象Windows 上設(shè)置num_workers8迭代第一個 batch 時程序無響應(yīng)或直接拋出BrokenPipeError。原因Windows 下每個 worker 都會復(fù)制一份數(shù)據(jù)集對象如果 Dataset 初始化里做了 PCA 擬合和數(shù)據(jù)標準化這些計算會被每個 worker 重復(fù)執(zhí)行大量內(nèi)存被占用并且進程間通信效率低容易掛起。解決Windows 下老實把num_workers2或 3。如果確實需要更多 worker把數(shù)據(jù)預(yù)處理放到 Dataset 初始化之前生成一個預(yù)處理后的.npy文件然后在 Dataset 里只做簡單的讀取和裁剪。4.4 不同 epoch 之間樣本劃分不一致導(dǎo)致驗證集失真現(xiàn)象訓(xùn)練損失正常下降但每輪驗證集精度波動很大甚至出現(xiàn)驗證精度比訓(xùn)練精度還高的情況。原因使用random_split劃分訓(xùn)練集和驗證集后如果在同一個 Dataset 實例上又創(chuàng)建了多個 DataLoader 但沒有固定隨機種子每次遍歷時 shuffle 產(chǎn)生的順序不同本質(zhì)上樣本劃分變了。解決在創(chuàng)建random_split之前固定全局隨機種子import torch, random, numpy as np random.seed(42) np.random.seed(42) torch.manual_seed(42)同時給DataLoader的generator參數(shù)傳入一個固定的torch.Generator。這樣每次運行實驗訓(xùn)練集和驗證集的劃分完全一致結(jié)果可復(fù)現(xiàn)。4.5 PCA 在訓(xùn)練集和驗證集上分別擬合導(dǎo)致數(shù)據(jù)泄漏現(xiàn)象測試集精度異常高遠高于訓(xùn)練集或者驗證集精度比訓(xùn)練集還高 5 個百分點以上。原因如果在構(gòu)建驗證集 Dataset 時又調(diào)用了一次PCA.fit驗證數(shù)據(jù)也被降維到同一主成分空間相當于模型間接見過驗證集統(tǒng)計信息。解決PCA 只允許在訓(xùn)練集上擬合一次然后把擬合好的 PCA 對象保存下來驗證集和測試集直接調(diào)用transform。具體實現(xiàn)是先整體加載數(shù)據(jù)手動切分出訓(xùn)練像素和測試像素再對訓(xùn)練像素擬合 PCA然后把變換參數(shù)保存到.pkl驗證時加載并transform。4.6 數(shù)據(jù)增強導(dǎo)致樣本邊緣像素類別被“翻轉(zhuǎn)”到錯誤空間位置現(xiàn)象開啟旋轉(zhuǎn)增強后分類精度反而下降 3 個百分點以上。原因鄰域窗口旋轉(zhuǎn)后中心像素雖然沒變但窗口內(nèi)其他像素的相對位置發(fā)生了變化。對于某些紋理依賴極強的類別比如道路的直線邊緣旋轉(zhuǎn)后特征方向錯亂模型學(xué)不到穩(wěn)定的方向不變特征反而被增強樣本干擾。解決如果你發(fā)現(xiàn)增強后精度下降先關(guān)閉旋轉(zhuǎn)只保留翻轉(zhuǎn)和少量平移。高光譜鄰域窗口的任務(wù)本質(zhì)是對中心像素分類空間方向信息往往是有益的強行做旋轉(zhuǎn)有時適得其反。我的經(jīng)驗是只有在每類樣本數(shù)少于 200 時才考慮旋轉(zhuǎn)增強。5. 進階技巧把高光譜樣本緩存為 .npy 文件用查表法加速加載5.1 為什么每次訓(xùn)練前切窗口是浪費算力每一輪 epoch 都要通過 Dataset 的__getitem__重新裁剪鄰域窗口。如果樣本量是 5 萬窗口大小 11×11一次 epoch 就要執(zhí)行 5 萬次數(shù)組切片、轉(zhuǎn)置、標準化操作。雖然 numpy 切片很快但加上 Python 層的索引計算和類型轉(zhuǎn)換整體會拖慢數(shù)據(jù)加載速度。尤其是在num_workers0的調(diào)試環(huán)境下GPU 可能會長時間等待 CPU 數(shù)據(jù)準備。一個很實用的優(yōu)化方法在訓(xùn)練之前把所有的樣本和標簽一次性切好存成.npy文件。因為高光譜數(shù)據(jù)的樣本是固定尺寸的切完之后每個樣本是一個(bands, window_size, window_size)的數(shù)組標簽是一個長向量。這樣 Dataset 的__getitem__就簡化成了self.samples[index]和self.labels[index]幾乎零計算。5.2 緩存腳本的實現(xiàn)與使用import numpy as np from sklearn.decomposition import PCA from sklearn.preprocessing import StandardScaler import h5py def preprocess_and_cache(data_path, label_path, cache_path, window_size11, n_components30): # 讀取原始數(shù)據(jù) with h5py.File(data_path, r) as f: data f[data][:] labels np.load(label_path) # 標準化 h, w, b data.shape flatten data.reshape(-1, b) flatten StandardScaler().fit_transform(flatten) data flatten.reshape(h, w, b) # PCA 降維 pca PCA(n_componentsn_components) flatten data.reshape(-1, b) flatten_pca pca.fit_transform(flatten) data flatten_pca.reshape(h, w, n_components) # 提取所有有效樣本窗口 pad window_size // 2 samples [] sample_labels [] for i in range(pad, h - pad): for j in range(pad, w - pad): if labels[i, j] 0: window data[i - pad:i pad 1, j - pad:j pad 1] window np.transpose(window, (2, 0, 1)) # (bands, win, win) samples.append(window) sample_labels.append(labels[i, j] - 1) samples np.stack(samples).astype(np.float32) sample_labels np.array(sample_labels, dtypenp.int64) np.save(cache_path _samples.npy, samples) np.save(cache_path _labels.npy, sample_labels) print(f緩存完成: {samples.shape}, 標簽數(shù): {sample_labels.shape})使用緩存后的 Dataset 就非常輕量了class CachedHyperspectralDataset(Dataset): def __init__(self, samples_path, labels_path): self.samples np.load(samples_path) self.labels np.load(labels_path) def __len__(self): return len(self.labels) def __getitem__(self, idx): return (torch.from_numpy(self.samples[idx]), torch.tensor(self.labels[idx], dtypetorch.long))這個緩存做法有兩個額外收益。第一調(diào)試階段可以脫離原始數(shù)據(jù)路徑直接加載緩存避免每次都執(zhí)行耗時的 PCA。第二當你需要對比不同窗口尺寸時可以生成多份緩存文件實驗切換只是換路徑不用重跑預(yù)處理。5.3 查表法替代實時增強把增強后的樣本預(yù)生成如果數(shù)據(jù)增強是必須的但你又不想在__getitem__里實時計算可以在緩存階段直接把翻轉(zhuǎn)、旋轉(zhuǎn)后的樣本全部擴充進樣本庫。比如每類樣本數(shù)較少時把原樣本加上水平翻轉(zhuǎn)、垂直翻轉(zhuǎn)、旋轉(zhuǎn) 90 度生成 4 倍樣本量。這樣 Dataset 內(nèi)樣本數(shù)變成了原來的 4 倍訓(xùn)練時每次抽樣都是已經(jīng)增強后的樣本不需要任何隨機操作。def augment_samples(samples, labels): aug_samples [] aug_labels [] for s, lbl in zip(samples, labels): # s shape: (bands, h, w) aug_samples.append(s) aug_labels.append(lbl) aug_samples.append(np.flip(s, axis2)) aug_labels.append(lbl) aug_samples.append(np.flip(s, axis1)) aug_labels.append(lbl) aug_samples.append(np.rot90(s, k1, axes(1, 2))) aug_labels.append(lbl) return np.stack(aug_samples), np.array(aug_labels)這種做法適合樣本總數(shù)不超過 3 萬的中小型數(shù)據(jù)集。因為增強后的樣本在內(nèi)存中占用量成倍增加如果原始樣本已經(jīng)占 3 GB4 倍后就 12 GB很多機器扛不住。所以我只在必要時使用預(yù)生成增強否則還是用實時增強更靈活。5.4 驗證加載正確性的一個簡單檢查方法寫完 Dataloader 后不要急著開始訓(xùn)練先用一個 batch 做 sanity check。檢查三件事batch 的形狀是否符合預(yù)期、標簽是否正確對應(yīng)中心像素、數(shù)據(jù)分布是否合理。train_iter iter(train_loader) batch, labels next(train_iter) print(batch.shape, labels.shape) print(標簽范圍:, labels.min().item(), labels.max().item()) print(batch 均值:, batch.mean().item(), 標準差:, batch.std().item())我通常會再跑一個小循環(huán)統(tǒng)計一批樣本里標簽的類別直方圖確認沒有出現(xiàn)類別空缺。如果某個類別只在驗證集出現(xiàn)而訓(xùn)練集沒有那模型永遠學(xué)不會這一類最終測試分類時這一類會全部被分錯。從那以后我每次換數(shù)據(jù)集或調(diào)整窗口尺寸都會強制走一遍緩存預(yù)處理和 batch 檢查確定形狀和標簽分布都正常了才啟動訓(xùn)練。這樣能省掉很多訓(xùn)練到一半才發(fā)現(xiàn)數(shù)據(jù)加載問題的后悔藥時間。希望這篇關(guān)于高光譜 Dataloader 的拆解筆記能幫你少踩幾個坑。本文還有配套的精品資源點擊獲取