★ FEATURED ARTICLE 简介面向使用深度学习处理高光谱图像的研究者与工程实践者这份代码资源围绕 PyTorch 的 Dataset 与 DataLoader 展开聚焦高光谱波段多、数据量大所带来的读取与批处理效率问题。压缩包共 7 个文件、约 5.69MB其中 3 个 py 脚本分别对应数据集类定义、DataLoader 构建与训练流程串联2 个 pyc 为配套的模块缓存2 个 mat 文件提供 Indian Pines 等实测高光谱数据可直接跑通示例。代码中体现了归一化、裁剪缩放、batch_size 与 num_workers 调优、shuffle 采样、pin_memory 加速及自定义 collate_fn 合并多维样本等常用做法读者可据此搭建自己的数据管线并排查内存不足与加载瓶颈。已有 614 人学习适合需要现成高光谱加载范式的中级 PyTorch 用户参考。1. 高光谱数据加载为什么 ImageFolder 那套搬过来会翻车第一次拿到高光谱数据的人十有八九会先去找类似ImageFolder的目录结构一个类别一个文件夹里面放图片DataLoader一包就能训练。真到动手才发现根本对不上——高光谱不是一堆散图而是一个(H, W, C)的三维立方体C 动辄上百个波段标签往往是另一张同尺寸的二维图背景像素还标着 0。你没法按文件分目录也没法用PIL直接读。这篇笔记围绕 PyTorch DataLoader 加载高光谱数据这条主线把文件读取、Dataset 构造、patch 切分、参数调优和排错串起来讲。适合正在做地物分类、波段分析、小样本高光谱实验的从业者也适合刚接手一份.mat立方体、不知道从哪下手的同学。下面按“数据长什么样 → 怎么写 Dataset → DataLoader 怎么调 → 坑在哪”的顺序推进。2. 先看清立方体长什么样.mat、ENVI 与 TIFF 的三条读取路线在写任何一行 Dataset 代码之前先把数据读进来打印形状、类型和值域。这一步偷懒后面所有报错都会变成玄学。2.1 形状约定HWC 还是 CHW先统一再谈别的高光谱社区里流传着两套张量顺序。MATLAB 一路下来的数据习惯是(H, W, C)因为 MATLAB 的索引是先行后列再波段PyTorch 的卷积层要的是(N, C, H, W)中间差一次转置。我的习惯是读进来之后立刻在内存里统一成(C, H, W)Dataset 返回的单个样本是(C, p, p)让DataLoader默认 collate 直接拼成(B, C, p, p)中间不再做任何转置。统一顺序的好处不只是省一次permute。当你后面要做波段归一化、波段丢弃、按波长切片这些操作时第一个维度是波段会让索引写得非常自然反过来如果一直是(H, W, C)每次切波段都要写cube[:, :, band_idx]稍微一疏忽就把空间维切掉了。import numpy as np def to_chw(cube: np.ndarray) - np.ndarray: 把 (H, W, C) 统一转成 (C, H, W)已是 CHW 则原样返回。 if cube.ndim ! 3: raise ValueError(f期望三维数组实际 {cube.ndim} 维) # 高光谱的波段数通常远小于空间边长用这个启发式判断顺序 if cube.shape[2] cube.shape[0] and cube.shape[2] cube.shape[1]: cube np.transpose(cube, (2, 0, 1)) return np.ascontiguousarray(cube)np.ascontiguousarray这步别省。转置之后数组在内存里是不连续的PyTorch 从 numpy 转 tensor 时会报警告某些版本下还会悄悄复制一份既慢又费内存。判断顺序的启发式用的是“波段数小于空间边长”对绝大多数航空/星载数据成立如果拿到的是空间很小的 patch 数据集这个判断会失效那就老老实实手动指定。提示不要依赖自动判断初次接触一份数据时把cube.shape打印出来人工确认一次写死在配置里比任何启发式都可靠。2.2 .mat 读取键名、dtype 和值域三件事.mat是最常见的分发格式也最容易出问题。第一个坑是键名。scipy.io.loadmat返回的是一个字典除了你关心的变量还有__header__、__version__、__globals__这些元信息。变量名可能是data、HSI、cube、X没有任何标准。import scipy.io as sio import numpy as np mat sio.loadmat(scene.mat) # 先看有哪些业务变量过滤掉 __ 开头的元信息 keys [k for k in mat.keys() if not k.startswith(__)] print(可用变量:, keys) cube mat[data] # 键名按上一步打印结果替换 print(shape:, cube.shape, dtype:, cube.dtype) print(min/max:, cube.min(), cube.max())打印dtype和值域是为了决定归一化方式。常见情况是uint16值域 0 到几千也有uint8偶尔会遇到已经归一化到[0, 1]的float32。这三种的处理完全不同uint16要除以 65535 或按波段最大值归一化uint8除以 255已经是浮点的就不要再除否则全图变成 0.00x模型前几十个 epoch 学不到东西。还有一个隐蔽问题uint16转float32要显式写astype(np.float32)。如果你写成cube / 65535numpy 在某些版本下会先升到float64转成 tensor 时就是torch.float64送进float32的模型直接报 dtype mismatch。这个错在第一次跑通时最容易遇到。def normalize_cube(cube: np.ndarray) - np.ndarray: cube cube.astype(np.float32) # 先转 float32避免 float64 if cube.max() 1.5: # 粗略判断是否已归一化 cube cube / max(float(cube.max()), 1.0) return cube2.3 ENVI 与 TIFF元数据比像素更值得看ENVI 格式是一对文件.hdr头文件加一个无扩展名的数据文件。头文件里最关键的几项是samples、lines、bands、data type和interleave。interleave决定了内存布局bsq是按波段顺序存bil是按行bip是按像素。读进来之后数组形状是对的但如果你自己手写解析interleave弄错会导致图像看起来像被切碎的花屏。常见做法是用spectral库的open_image它把头文件解析和数组重排都封装好了import spectral img spectral.open_image(scene.hdr) cube img.load() # 形状 (lines, samples, bands) wavelength img.metadata.get(wavelength, None) print(cube.shape, cube.dtype) print(波段数:, img.nbands, 波长信息:, 有 if wavelength else 无)波长信息值得单独看一眼。有些下游任务比如植被指数、特定吸收峰分析需要按波长选波段如果.hdr里没有wavelength字段就得靠wavelength units和起始波长自己推算这一块在预处理阶段就要补齐别等到训练时才发现选不了波段。TIFF 格式的问题在维度顺序。多页 TIFF 用tifffile读出来通常是(bands, H, W)而单页的多通道 TIFF 可能是(H, W, C)。同一个.tif后缀两种布局都合法import tifffile with tifffile.TiffFile(scene.tif) as tf: cube tf.asarray() print(页数:, len(tf.pages), 形状:, cube.shape) # 若读成 (H, W, C)用 2.1 的 to_chw 统一 cube to_chw(cube).astype(np.float32)2.4 切 patch滑动窗口和标签对齐逐像素分类的样本是单个像素的光谱向量形状(C, 1, 1)信息量少模型容易过拟合。更常见的是 patch 级样本以某个像素为中心取p × p的空间邻域标签取中心像素的类别。p一般取 9、11、13奇数保证中心对齐。切 patch 有两个细节必须处理。第一是边界。图像四角的像素取不到完整邻域要么丢弃要么填充。我一般先对整个 cube 做一次np.pad模式用reflect这样所有像素都能取到完整 patch索引也简单def pad_cube(cube: np.ndarray, p: int, mode: str reflect): r p // 2 # 只 pad 空间维波段维不动 padded np.pad(cube, ((0, 0), (r, r), (r, r)), modemode) return padded第二是标签对齐。gt是(H, W)的整型图0 通常表示未标注或背景。切 patch 时索引要基于原始坐标填充只影响 cube 的读取偏移不能把gt也 pad 了再去取值否则中心位置会整体偏移r个像素训练出来的模型精度看着还行混淆矩阵却怎么都对不上。这个偏移错误非常隐蔽因为 loss 会正常下降。3. 写一个能复用的 Datasetpatch 切分、标签映射与光谱增强Dataset 写得好不好决定了后面调参是享受还是受罪。核心原则只有一条__init__只做一次性的准备工作__getitem__里只做轻量的切片和变换。3.1init里放什么、不放什么放进去的已经读进内存的 cube归一化后、gt图、预先算好的样本索引表、标签映射字典、增强开关。不放的每次取样都要重新算的东西比如对整幅图做统计、动态生成随机 patch 坐标。样本索引表我一般用两个 numpy 数组存(rows, cols)而不是 Python list of tuple。几万个样本时差别不大但到几十万级tuple 列表的内存和遍历开销就显出来了。构造时一次性把有效像素筛出来import numpy as np import torch from torch.utils.data import Dataset class HSIPatchDataset(Dataset): def __init__(self, cube, gt, patch_size11, ignore_label0, label_mapNone, transformNone): # cube: (C, H, W) float32已归一化 self.cube cube self.patch_size patch_size self.transform transform self.r patch_size // 2 # 标签重映射把原始类别 id 压成 0..K-1ignore_label 单独处理 self.label_map label_map or {} gt_mapped np.full_like(gt, -1, dtypenp.int64) for src, dst in self.label_map.items(): gt_mapped[gt src] dst self.gt gt_mapped # 有效样本 标签不为 -1 的像素 rows, cols np.where(self.gt 0) self.rows rows.astype(np.int32) self.cols cols.astype(np.int32) self.num_classes int(self.gt.max()) 1这段代码里有两个决定后续走向的选择。label_map用字典显式映射而不是靠“减一”这种隐含规则是因为不同数据集的类别编号毫无规律有的是 1 到 16有的跳号写死减一会把背景和第一类混在一起。ignore_label统一映射成-1在__getitem__里再转成255给CrossEntropyLoss(ignore_index255)用这是 PyTorch 生态里比较通行的约定。3.2 两种getitem像素级和 patch 级像素级最简单取(C, 1, 1)def __getitem_pixel(self, idx): r, c self.rows[idx], self.cols[idx] x self.cube[:, r, c] # (C,) x torch.from_numpy(x).unsqueeze(-1).unsqueeze(-1) # (C,1,1) y int(self.gt[r, c]) return x, ypatch 级要处理填充。注意这里我用的是“先 pad 整个 cube”的方案__init__里已经 pad 过所以__getitem__直接按偏移取def __getitem_patch(self, idx): r, c self.rows[idx], self.cols[idx] r0, c0 r, c # 原始坐标未 pad rp, cp r self.r, c self.r # pad 后坐标 p self.patch_size x self.cube[:, rp - self.r: rp self.r 1, cp - self.r: cp self.r 1] # (C, p, p) y int(self.gt[r0, c0]) x torch.from_numpy(np.ascontiguousarray(x)) if self.transform: x self.transform(x) return x, y两个变量r0/c0和rp/cp分开写就是为了避免 2.4 里说的标签偏移。这种“看起来多此一举”的写法在 review 别人代码时救过我好几次——很多人直接在 pad 后的坐标上取gt边界样本的标签全错但只有几十个像素精度掉零点几个点根本发现不了。3.3 高光谱专用增强翻转、波段丢弃与噪声注入高光谱的增强和自然图像不太一样。空间翻转、90 度旋转可以直接用因为地物在空间上没有方向偏好。但颜色抖动这一套不能用波段之间是有物理连续性的你把波段顺序打乱等于把光谱曲线撕碎模型学到的全是噪声。我常用三类增强。空间翻转用torch.flip零成本。波段丢弃spectral erasing随机把连续几个波段置零模拟传感器坏带或大气吸收能让模型对波段缺失更鲁棒。高斯噪声按信噪比注入而不是固定标准差因为不同数据集的反射率量级差别很大import torch class HSIAugment: def __init__(self, p_flip0.5, p_band_drop0.3, snr_db30): self.p_flip p_flip self.p_band_drop p_band_drop self.snr_db snr_db def __call__(self, x): # x: (C, p, p) if torch.rand(1).item() self.p_flip: x torch.flip(x, dims[2]) # 水平翻转 if torch.rand(1).item() self.p_flip: x torch.flip(x, dims[1]) # 垂直翻转 if torch.rand(1).item() self.p_band_drop: c x.shape[0] width torch.randint(1, max(2, c // 20), (1,)).item() start torch.randint(0, c - width, (1,)).item() x x.clone() x[start:start width] 0.0 # 连续波段置零 # 按信噪比注入噪声 signal_power x.pow(2).mean() noise_power signal_power / (10 ** (self.snr_db / 10)) x x torch.randn_like(x) * noise_power.sqrt() return xp_band_drop控制置零概率width上限取波段数的 1/20避免一次抹掉太多信息。噪声那几行里的snr_db是分贝值30 dB 大约对应 3% 的噪声水平可以根据数据质量在 20 到 40 之间调。增强只在训练集开验证和测试一律关闭这个开关放在 Dataset 外面控制不要写死在类里。4. DataLoader 参数怎么定num_workers、采样器与 collate_fnDataset 写对之后DataLoader 的参数决定了你的 GPU 是满负荷跑还是有一半时间在等数据。高光谱的数据量比自然图像大一个量级一个11 × 11 × 200的 patch 展开是 24200 个浮点数接近 100 KB这个尺度下 num_workers 和内存的关系必须算清楚。4.1 num_workers 与 persistent_workers吞吐和内存的平衡num_workers0意味着数据加载在主进程里做GPU 每个 batch 都要等 CPU。设成 4 到 8 通常能把 GPU 喂饱但不是越大越好。Linux 下 worker 用 fork 启动父进程的 cube 数组通过写时复制共享只读的话不会真复制可一旦 Dataset 里有任何在__getitem__里修改数组的操作写时复制立刻退化成真实复制4 个 worker 就是 4 份 cube几百 MB 的数据直接把内存打满。from torch.utils.data import DataLoader loader DataLoader( train_set, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue, persistent_workersTrue, prefetch_factor4, drop_lastTrue, )逐项说。pin_memoryTrue把 batch 放进锁页内存配合训练循环里的x.cuda(non_blockingTrue)能做异步拷贝这一步在数据量大时能省下可观的等待。persistent_workersTrue让 worker 进程在 epoch 之间存活避免每个 epoch 重新 fork 一次样本多、epoch 多的时候差别明显。prefetch_factor4表示每个 worker 预取 4 个 batch默认是 2内存紧张就调回 2。drop_lastTrue是为了避免最后一个不满的 batch 影响 BatchNorm 的统计样本量本来就少的话可以关掉。注意persistent_workersTrue必须配合num_workers 0使用否则 PyTorch 会直接报错。4.2 类别不平衡WeightedRandomSampler 的权重怎么算高光谱地物分类的类别分布极度不均大片农田加零星建筑最多的类可能占 40%最少的类不到 1%。直接 shuffle 采样的结果是模型倾向于预测大类整体精度看着不错小类召回惨不忍睹。WeightedRandomSampler给每个样本一个权重按权重有放回地抽。权重按类别频率的倒数算import numpy as np from torch.utils.data import WeightedRandomSampler # 每个训练样本的标签 labels np.array([train_set[i][1] for i in range(len(train_set))]) counts np.bincount(labels, minlengthtrain_set.num_classes).astype(np.float32) # 类别权重 1 / 频次再按样本映射 class_w 1.0 / np.clip(counts, 1.0, None) sample_w class_w[labels] sample_w torch.from_numpy(sample_w).double() sampler WeightedRandomSampler( weightssample_w, num_sampleslen(sample_w), replacementTrue, )np.clip(counts, 1.0, None)防止某个类频次为 0 时除零。num_samples设成样本总数意味着一个 epoch 的样本量和普通 shuffle 一致如果想让小类多出现几次可以把这个值调大到 1.5 倍但要注意 epoch 会变长。replacementTrue是有放回抽样必须开否则权重只是摆设。用 sampler 时 DataLoader 里的shuffle必须设成False两个同时开会报错这个错很直白但第一次遇到还是会愣一下。4.3 collate_fn 与 dtype 对齐默认的default_collate能处理(tensor, int)这种返回拼出来是(B, C, p, p)和(B,)。需要自定义collate_fn的场景主要是两种返回里带了不定长信息比如原始坐标或者做了 mixup 之类的 batch 级增强。def hsi_collate(batch): xs, ys, coords [], [], [] for x, y, rc in batch: xs.append(x) ys.append(y) coords.append(rc) xs torch.stack(xs, dim0) # (B, C, p, p) ys torch.tensor(ys, dtypetorch.long) coords torch.tensor(coords, dtypetorch.int32) return xs, ys, coordstorch.stack要求每个样本形状完全一致所以 patch 级 Dataset 里必须保证边界处理统一不能一部分 pad 一部分裁剪否则 collate 直接报 size mismatch。ys显式指定torch.long因为CrossEntropyLoss要求标签是 int64从 numpy 转过来如果是 int32 会报错。参数选型的经验值可以归纳成一张表实际项目里按数据规模微调参数小数据集 5 万样本大数据集 20 万样本说明batch_size32 ~ 64128 ~ 256受显存限制patch 越大越要降num_workers2 ~ 48 ~ 16观察 GPU 利用率低于 80% 就加prefetch_factor24内存充裕时提高persistent_workersFalseTrue减少 fork 开销pin_memoryTrueTrue有 GPU 就开这张表不是标准答案只是起步值。真正该看的是nvidia-smi里的 GPU 利用率和top里的 CPU 占用GPU 忽高忽低说明数据供不上加 workerCPU 打满而 GPU 空闲说明__getitem__里有重计算该往缓存方向优化。5. 避坑排查五个高光谱 DataLoader 翻车现场下面五条都是实际项目里踩过的按“现象 → 原因 → 解决”写遇到对应症状可以直接对号入座。5.1 训练 loss 一直不降模型像在学随机噪声现象数据能跑通loss 在前几个 epoch 缓慢下降然后卡住验证精度和随机猜差不多。原因维度顺序错了。(H, W, C)没有转成(C, H, W)DataLoader拼出来是(B, H, W, C)模型第一层Conv2d(in_channelsC)收到的却是H要么直接报通道数不匹配要么在某些写法下被广播成了错误形状模型实际在学空间维上的模式。解决在 Dataset 返回前打印一次x.shape确认是(C, p, p)。更稳妥的做法是在__init__里做一次断言assert cube.shape[0] expected_bands把错误挡在训练开始之前。5.2 第一个 batch 就卡死或者报 shared memory 相关错误现象num_workers设成 0 一切正常设成 4 之后进程挂起或者报unable to open shared memory object、Bus error。原因两个可能。一是容器环境的/dev/shm默认只有 64 MB多 worker 加pin_memory很快耗尽二是 Dataset 里存了 lambda、打开的文件句柄这类不可 pickle 的对象worker fork 时传不过去。解决容器里把/dev/shm调大或者把num_workers降到 2 先验证检查 Dataset 的__init__参数把 lambda 换成模块级函数把文件句柄在__getitem__里现开现关不要挂在self上。5.3 验证精度高得离谱换一批数据就崩现象验证集精度 98% 以上测试集或者新场景一跑掉到 60%。原因训练集和验证集是按像素随机划分的相邻像素的 patch 高度重叠几乎等于把训练样本原样喂给了验证集。高光谱的空间相关性极强这种泄漏比自然图像严重得多。解决按空间块划分比如按行号切成若干条带整条带进训练或验证或者按地理区域划保证两边的 patch 没有重叠像素。划分完打印一下两边的坐标范围确认没有交叠。5.4 混淆矩阵对不上某个类永远预测不出来现象训练正常精度也还行但混淆矩阵里某一类的召回是 0或者背景类占据了大部分预测。原因标签映射没做干净。原始gt里 0 既表示背景又表示未标注直接送进CrossEntropyLoss会被当成一个正常类别参与训练模型很快就学会“全预测 0”这个捷径。解决显式建立label_map把未标注像素映射成-1并在 Dataset 里过滤掉背景类如果需要单独成类就给它一个明确的 id不要和 0 混用。CrossEntropyLoss里设ignore_index255Dataset 返回时把-1转成 255。5.5 训练到一半内存持续上涨最后 OOM现象前几个 epoch 正常越往后内存占用越高最终被系统杀掉。原因__getitem__里做了原地修改比如x / 255或x[band] 0触发了写时复制每个 worker 都持有一份完整的 cube 副本或者在 transform 里不断往列表里 append 样本。解决Dataset 里所有数组按只读使用需要修改的地方先clone()transform 保持无状态不要持有跨样本的缓存。排查时用id(cube)在不同 worker 里打印确认是不是同一块内存。6. 进阶验证memmap 缓存与小样本过拟合自检数据管道调通之后还剩两件事值得做把重复计算缓存下来以及在正式训练前用一个小 batch 验证整条链路没有静默错误。6.1 用 np.memmap 把 patch 缓存到磁盘patch 级 Dataset 每次取样都要做切片和增强。切片本身不贵但如果预处理里包含辐射校正、波段重采样这类重计算每个 epoch 重算一遍就是纯浪费。np.memmap可以把预处理结果落到磁盘用的时候按需映射进内存既不占 RAM又比每次重算快。import numpy as np import os def build_cache(dataset, cache_path, shape, dtypenp.float32): if os.path.exists(cache_path): return np.memmap(cache_path, dtypedtype, moder, shapeshape) cache np.memmap(cache_path, dtypedtype, modew, shapeshape) for i in range(len(dataset)): x, _ dataset[i] cache[i] x.numpy() cache.flush() return np.memmap(cache_path, dtypedtype, moder, shapeshape)shape要提前算好是(样本数, C, p, p)dtype用float32别用float64体积直接翻倍。modew创建moder只读映射后者在多 worker 下更安全不会出现两个进程同时写同一块。缓存文件通常比原始 cube 大一个数量级磁盘空间要留够一个 10 万样本、11 × 11 × 200的缓存大约 9.7 GB。缓存做好之后Dataset 的__getitem__改成从 memmap 读增强照常做。注意 memmap 返回的是 numpy 数组转 tensor 时仍然要np.ascontiguousarray因为按索引切出来的一片在磁盘上是连续的但经过增强之后不一定。6.2 一个 batch 的过拟合自检正式训练之前我习惯拿一个 batch 反复训练几十步看 loss 能不能降到接近 0。这一步能抓出大部分静默错误标签偏移、归一化错误、维度顺序错、损失函数和标签 dtype 不匹配都会在这几十步里暴露出来。loss 降不下去说明模型在输入里找不到任何和标签相关的信息先别急着调超参。具体做法是固定一个 batch关掉增强用较大的学习率跑 100 步观察 loss 曲线。正常情况下应该单调下降到 1e-3 量级。如果降到某个值就平台化多半是标签有问题如果 loss 震荡剧烈检查学习率和归一化。下面这张自检表可以贴在实验记录里每次换数据集过一遍检查项期望结果不通过时优先排查单个 batch 能否过拟合loss 降到 1e-3 以下标签映射、维度顺序输入张量值域大致落在 [0, 1]归一化除数、dtype标签 dtypetorch.int64Dataset 返回时的类型转换一个 epoch 耗时与样本量线性相关num_workers、缓存命中GPU 利用率稳定在 80% 以上prefetch、worker 数量各类别样本数最小值不为 0划分策略、label_map6.3 我后来固定下来的两个习惯第一个习惯是每换一份数据先写一个十行的脚本只做读取、打印形状、打印值域、打印标签分布这四件事跑通了再动 Dataset。这十行代码看着不起眼但它把“数据本身的问题”和“管道的问题”隔离开了后面无论报什么错都能确定不是数据格式的锅。第二个习惯是给每个 Dataset 配一个_sanity_check方法随机抽三个样本把输入和标签打印出来训练前手动跑一次。有一次标签整体偏移了一个像素就是这个方法发现的——loss 正常下降验证精度也不难看只有把 patch 和标签可视化出来才能看出中心对不上。从那以后我每次改完切 patch 的逻辑都强制走一遍这个小检查再开始跑训练。数据管道这部分没有太多花哨的技巧难的是把每个边界情况想清楚。希望这些经验能帮你少走几个弯路。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
