简介面向使用PyTorch处理高光谱图像HSI的开发者与研究人员这份资源系统梳理了利用DataLoader高效加载与预处理高光谱数据的完整流程从自定义Dataset类、数据读取与归一化到设置batch_size、num_workers多线程加载、shuffle随机采样、pin_memory缓存加速以及自定义collate_fn合并批量样本每个环节均给出清晰示例可有效解决高光谱数据量大、通道数多导致的内存与训练效率问题。压缩包共7个文件包含3个Python源码脚本、2个pyc预编译文件与2个MAT数据文件整体大小仅5.69MB结构紧凑、开箱即用。资源以IndianPines真实高光谱数据集为示例数据集定义、加载器实现与训练脚本一应俱全方便读者对照学习并快速迁移到自己的项目中。目前已有609人学习浏览适合刚接触高光谱深度学习或希望优化数据加载流程的PyTorch使用者下载参考。1. 用 PyTorch 的 DataLoader 加载高光谱数据为什么第一步就走不稳用 PyTorch 的 DataLoader 加载高光谱数据很多人第一反应是“这有什么难的”真动手才发现处处是坑。普通 RGB 图像只有 3 个波段而高光谱图像少则几十个波段如 GF-5 的 AHSI 有 330 个波段多则两百多个如 AVIRIS 的 224 个波段单帧数据量比普通图像大两到三个数量级文件格式也从 jpg 变成了 ENVI、MAT、TIFFtorchvision 那套现成的加载方案基本失灵。这篇文章按“数据读入 → 自定义 Dataset → DataLoader 参数配置 → 踩坑排查 → 效率优化”的路径把一条能直接落地的加载链路讲清楚适合做高光谱分类、地物识别、蚀变信息提取的从业者和研究生读完能直接照抄代码改到自己项目里。2. 数据读入与预处理从 ENVI/MAT/TIFF 到 numpy 的四条通路2.1 ENVI 格式spectral 库读入与反射率换算高光谱从业者手里最常见的格式就是 ENVI 标准格式一个.hdr头文件加一个.dat或.img数据文件。GF-5、高分五号这类星载数据的 L1 级产品经常以这种形式发布。读 ENVI 我一般直接用spectral库它比 GDAL 轻量而且读出来就是内存友好的 numpy 数组。import numpy as np from spectral.io import envi # 读 ENVI 格式hdr 和 dat 放在同一目录 img envi.open(gf5_scene.hdr, gf5_scene.dat) data img.load().astype(np.float32) print(data.shape) # (H, W, C)例如 (1024, 1024, 330)img.load()会把所有波段读入内存返回的数组维度是(H, W, C)这和 PyTorch 期望的(C, H, W)不一样后面转换时要记着。astype(np.float32)这一步非常关键ENVI 原始数据常见的是uint16或int16如果不转成 float32后续做归一化和反向传播都会出问题。高光谱如何转反射率是绕不开的问题。ENVI 头文件里通常会写scale factor之类的系数常见做法是读入后逐波段乘系数再除以 10000 或 100000 把 DN 值变成 0~1 范围的反射率。具体系数以产品头文件或说明文档为准不要凭经验写死不同级别产品的定标系数不一样。scale_factor 10000.0 # 示例值以头文件实际值为准 data data / scale_factor data np.clip(data, 0.0, 1.0)clip 到[0, 1]是为了把异常像元压掉但要注意如果后续要做归一化clip 之后再统计 min/max 会更稳定。反射率换算完成后建议立刻把结果保存成.npy避免每次跑实验都重新读一遍原始文件。2.2 MAT 格式scipy.io 与 h5py 双方案ICVL 这类公开高光谱数据集经常以.mat格式发布文件里一般包含data和labels两个变量。读这种文件我用scipy.io.loadmat但有个坑必须提前说MATLAB 保存的 v7.3 格式文件loadmat直接报错必须改用h5py。import scipy.io as sio # 普通 v7 及以下版本 mat sio.loadmat(icvl_data.mat, squeeze_meTrue) data mat[data] # 可能是 (N, C) 或 (H, W, C) labels mat[labels] # 对应的标签squeeze_meTrue会把单元素维度去掉比如某个变量维度是(1, N)读进来就变成(N,)省去后面手动 squeeze 的麻烦。如果读的时候报NotImplementedError: Please use HDF5 reader说明文件是 v7.3 格式import h5py with h5py.File(icvl_data.mat, r) as f: data f[data][:] labels f[labels][:]h5py 返回的数组维度顺序和 MATLAB 里正好相反比如 MATLAB 里是(C, N)h5py 读出来就是(N, C)拿到数据后第一件事就是确认 shape别急着往下传。2.3 TIFF 格式rasterio 读取与波段顺序很多预处理后的高光谱数据会合成为 GeoTIFF尤其是和 GIS 工具链配合的时候。读 TIFF 我首选rasterio它读出来的数组直接是(C, H, W)和 PyTorch 的输入布局一致省一次转置。import rasterio with rasterio.open(hyperspectral_scene.tif) as src: data src.read() # (C, H, W) profile src.profilesrc.read()不加参数时读出所有波段返回 float32 或 uint16 取决于 TIFF 内部存储类型。要注意有些软件导出的 TIFF 波段顺序是 BIP 或 BSQ 排列rasterio已经帮你处理成(C, H, W)但波段之间的顺序仍然保持文件里的顺序做训练前建议打印波段名或波长信息核对一次。对于超大场景src.read()一次性读入内存可能直接吃掉几 GB我一般配合src.read(window...)按窗口读或者先用block_shapes看一下 TIFF 内部分块情况再决定怎么切。2.4 边界与切块把大图拆成可训练的小样本整幅高光谱图动辄上千万像素直接丢进 CNN 不现实常规做法是按窗口切 patch。切块时只记录索引而不复制数据等__getitem__真正访问时才切片内存占用能省一个量级。def build_sample_index(data_shape, patch_size11, stride5): 返回所有 (i, j) 左上角坐标不复制数据 h, w data_shape[:2] indices [] for i in range(0, h - patch_size 1, stride): for j in range(0, w - patch_size 1, stride): indices.append((i, j)) return indicespatch_size是窗口边长stride是滑动步长。窗口越小、步长越小样本数越多训练越充分但耗时也越长。如果做蚀变信息提取这种大范围任务我一般先用 stride 稍微大于 patch_size 的方式采样减少相邻样本的冗余同时用类别分布做一次过滤防止某些稀有类别在切块时被漏掉。3. 自定义 Dataset 类从一整幅图到一个个训练样本3.1 内存布局numpy 的 HWC 与 PyTorch 的 CHW 转换numpy 数组读出来是(H, W, C)但 PyTorch 的卷积网络默认输入是(C, H, W)所以__getitem__里必须做一次转置。这里有个容易忽略的细节transpose(2, 0, 1)得到的新数组是原始内存的非连续视图直接torch.from_numpy有时会报错或者带来额外拷贝开销稳妥做法是先np.ascontiguousarray。patch data[i:ipsize, j:jpsize, :] # (H, W, C) patch np.ascontiguousarray(patch) # 确保内存连续 patch patch.transpose(2, 0, 1) # (C, H, W) patch torch.from_numpy(patch).float() # 转 PyTorch Tensor这四行已经是我每次写高光谱 Dataset 的固定开头。torch.from_numpy默认共享内存后面在训练循环里做归一化或增强时如果对 Tensor 做 inplace 操作会反向影响原始 numpy 数组这个副作用在调试时最容易让人怀疑人生。3.2 归一化放哪里提前算好再进 Dataset高光谱波段多如果每个 epoch 都重新算一次 min/max 归一化纯属浪费算力。我的习惯是读入数据后立刻统计每个波段的 min、max保存成 npy 文件Dataset 里只用保存好的参数做线性变换。这样训练、验证、测试用的归一化参数完全一致不会因为数据集划分不同导致数据分布漂移。# 预处理阶段执行一次 band_min data.min(axis(0, 1)) band_max data.max(axis(0, 1)) np.save(band_min.npy, band_min) np.save(band_max.npy, band_max)归一化参数只保存原始数据统计结果不做其他全局标准化。高光谱数据每个波段的动态范围差异很大比如近红外波段反射率普遍偏高全局 min/max 会严重影响低反射率波段的信息。3.3 一个能直接抄的 Dataset 实现下面这个 Dataset 类覆盖了前面所有要点索引采样、HWC 转 CHW、归一化、标签读取。标签图通常是单通道整数矩阵每个像素的类别编号直接索引即可。from torch.utils.data import Dataset import numpy as np import torch class HSIDataset(Dataset): def __init__(self, data, label, sample_indices, patch_size11, band_minNone, band_maxNone, transformNone): self.data data.astype(np.float32) self.label label self.indices sample_indices self.patch_size patch_size self.band_min band_min self.band_max band_max self.transform transform def __len__(self): return len(self.indices) def __getitem__(self, idx): i, j self.indices[idx] patch self.data[i:iself.patch_size, j:jself.patch_size, :] if self.band_min is not None: patch (patch - self.band_min) / (self.band_max - self.band_min 1e-6) patch np.ascontiguousarray(patch).transpose(2, 0, 1) patch torch.from_numpy(patch).float() label torch.tensor(self.label[i, j], dtypetorch.long) return patch, label这个类里band_min和band_max是提前算好的每波段向量维度(C,)广播到(H, W, C)时自动逐波段操作。 1e-6防止某些波段 max 和 min 相等造成除零。transform参数预留着给数据增强用——高光谱增强和 RGB 不一样必须逐波段一致变换。3.4 为什么用索引列表而不是生成全部 patch很多新手会把所有 patch 在初始化阶段一次性切好存成一个大数组这不是不行但高光谱数据量太大时很容易翻车。一幅 1024×1024×330 的 float32 数据按 11×11 窗口切光存储就是全图的 11×11 倍直接超出内存。索引列表方案每次只切一个小窗口内存占用恒定代价是训练时每个 batch 都做一次切片但现代 CPU 处理这点切片绰绰有余。4. DataLoader 参数实战batch、worker、pin_memory 与采样策略4.1 关键参数对照表参数默认值作用高光谱场景建议batch_size1每个 batch 的样本数分类任务 64~128语义分割看显存num_workers0数据加载进程数物理核数的一半Windows 从 2 起步pin_memoryFalse锁页内存加速 CPU→GPU 拷贝GPU 训练时设为 Trueprefetch_factor2每个 worker 预取的 batch 数内存紧张时调成 1persistent_workersFalse是否保活 worker 进程训练轮次多时设 Truedrop_lastFalse是否丢弃最后不足整 batch 的数据常规分类任务建议 True4.2 num_workers 到底设多少高光谱单样本读取本身不重but 一旦数据在磁盘上按 ENVI/TIFF 存储随机读取的 IO 开销就会成为瓶颈num_workers就是用来掩盖这个开销的。常见的设法是os.cpu_count() // 2但不建议在 Windows 上直接拉满Windows 的 worker 进程模型开销大设到真实核数一半就够。from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size128, shuffleTrue, num_workers8, pin_memoryTrue, persistent_workersTrue, prefetch_factor2, drop_lastTrue )persistent_workersTrue配合多轮 epoch 能避免每轮都重新创建 worker 进程的开销代价是占用的内存不释放。如果训练脚本在第一个 epoch 结束时报BrokenPipeError大概率是num_workers设太高或者没加if __name__ __main__保护先把 worker 降到 2 验证一次。4.3 类别不平衡WeightedRandomSampler 怎么接高光谱地物分类的标签分布天然不平衡——裸地、植被这类大类占大多数稀有地物类别可能只有几百个像素。如果直接shuffleTrue模型会被大类带偏。常规做法是用WeightedRandomSampler按类别频率反比采样。from torch.utils.data import WeightedRandomSampler flat_label label.flatten() class_counts np.bincount(flat_label) weights 1.0 / class_counts[flat_label] sampler WeightedRandomSampler(weights, num_samples80000, replacementTrue) train_loader DataLoader(train_dataset, batch_size128, samplersampler)weights给每个样本一个采样概率稀有类别权重高训练时被抽中的次数就多。num_samples80000表示每个 epoch 采 8 万个样本采样带放回。注意设置了sampler后DataLoader 里的shuffle参数必须为 False两者冲突时会直接报错。4.4 自定义 collate_fn什么时候真的要写默认的collate_fn要求 batch 内所有样本的 tensor shape 完全一致。普通分类任务每个 patch 都是(C, 11, 11)用默认行为就够了。但做分割或者多尺度任务时patch 大小不统一就需要自己写一个填充版本。def collate_pad(batch): patches [item[0] for item in batch] labels [item[1] for item in batch] max_h max(p.shape[1] for p in patches) max_w max(p.shape[2] for p in patches) C patches[0].shape[0] padded torch.zeros(len(patches), C, max_h, max_w) for k, p in enumerate(patches): padded[k, :, :p.shape[1], :p.shape[2]] p return padded, torch.stack(labels) loader DataLoader(train_dataset, batch_size32, collate_fncollate_pad)拼 batch 时用零填充对齐到当前 batch 内的最大尺寸模型在训练时最好配合 mask 忽略填充区域否则填充的零值会变成真实像素参与梯度计算。我一般只在不得已时才用自定义 collate能保持形状一致就保持形状一致省心。5. 高光谱 DataLoader 避坑记录五条真实踩坑5.1 Windows 下 num_workers 一开大就报错现象训练到第一个 epoch 快结束时控制台抛出BrokenPipeError或RuntimeError: DataLoader worker (pid xxx) is killed by signal。原因Windows 下数据加载 worker 是 multiprocessing 实现的数据集初始化或索引生成不在if __name__ __main__保护内worker 进程重新导入主模块时递归执行了数据集构建逻辑导致进程崩溃。解决把整个训练流程包进if __name__ __main__:保护块同时把num_workers先降到 0跑通后再逐步加到 2、4、8。另外确认不是杀毒软件拦截了临时文件读写这个在 Windows 上真遇到过。5.2 patch 尺寸设太大显存直接爆掉现象batch_size 64 时 GPU 显示 OOM程序崩溃把 batch_size 降到 8 才能跑训练速度慢到怀疑人生。原因高光谱 patch 的通道数是普通 RGB 的几十倍。假设 330 个波段、patch 大小 32×32一个样本就是330*32*32*4约 1.35 MBbatch_size 64 时单 batch 就要 86 MB 显存往后还有模型参数和梯度OOM 是必然的。解决测量单样本的内存占用并估算全 batch 需求公式是C * H * W * 4 * batch_size。通道数多的数据集优先用 11×11 小 patch必须大 patch 时就把batch_size按比例缩小或者考虑混合精度训练。5.3 标签和图像错位训练集验证集空间重叠现象训练集准确率接近 99%验证集只有 60% 出头无论怎么调参都补不上差距。原因按滑窗切 patch 时相邻窗口之间只有 stride 的距离训练集和验证集如果来自同一幅图的相邻区域空间上高度相关模型在训练集上学到的空间纹理特征在验证集上同样适用但这种性能是虚高的。真正部署时换一片区域效果立刻掉下来。解决切分数据时按场景或按空间区域划分比如把整幅图分成左半训练右半验证或者用不同的航带数据分别做训练和验证。评估时也要看类别级别的 recall而不是只看整体准确率。5.4 uint16 数据被当成 uint8 截断现象图像整体偏暗某些波段过曝异常像元奇多反射率曲线呈锯齿状。原因很多库默认读 TIFF 时返回 uint8 或按 uint8 处理高光谱原始 DN 值经常超过 65535截断后所有超范围像元变成 255信息全部丢失。解决读入后立刻打印data.dtype确认。uint16 数据在归一化前保持原始 dtype不要提前转 uint8。如果已经截断了没有后悔药只能重新读原始文件。5.5 每个 epoch 越来越慢最终卡死现象第一个 epoch 30 秒第二个 50 秒第三个直接卡住不动。原因persistent_workersTrue加上prefetch_factor设过大worker 进程持续占用的内存越堆越多触发了操作系统的内存回收机制或者磁盘缓存被耗尽导致 IO 阻塞。解决先看任务管理器里的内存占用曲线。内存涨得很快就把prefetch_factor改成 1如果数据量本身巨大可以把 patch 预先缓存成 npy 文件减少每次随机读的开销。这套组合拳基本能解决百分之九十五的“越跑越慢”。6. 把加载链路再压一截缓存中间结果与光谱可视化验证6.1 把切好的 patch 落盘成 npy去掉重复读取每次训练都从原始 ENVI/TIFF 文件里切 patch对磁盘 IO 是巨大的浪费。我现在的做法是预处理阶段把训练集和验证集分别切好并存成 npy 文件显式地把磁盘随机读变成了顺序读训练时 IO 开销几乎可以忽略。train_patches [] train_labels [] for i, j in train_indices: patch data[i:i11, j:j11, :].astype(np.float32) train_patches.append(patch) train_labels.append(label[i, j]) np.save(train_patches.npy, np.stack(train_patches)) np.save(train_labels.npy, np.array(train_labels))这里注意磁盘占用11×11×330 的 float32 patch10 万个样本大约需要 15 GB。如果磁盘紧张可以压缩成 uint16 或只缓存归一化前的原始值训练时再现场转 float32。6.2 用光谱曲线确认归一化没有做错数据集加载做完第一件事不是开训而是抽一条光谱曲线画出来看。高光谱数据的 DNA 就是每个像素的光谱曲线归一化做没做对画一次图立刻见分晓。import matplotlib.pyplot as plt wavelengths np.arange(data.shape[2]) # 实际波段对应的波长 sample_pixel data[100, 100, :] plt.plot(wavelengths, sample_pixel) plt.xlabel(Band) plt.ylabel(Reflectance) plt.savefig(check_curve.png, dpi120)正常反射率曲线应该是平滑的相邻波段之间的反射率变化符合地物光谱特征。如果曲线是锯齿状或出现大量异常尖峰说明原始数据波段排列有问题或预处理除零了。训练过程中如果 loss 一直不降我也会先看几条随机样本的光谱曲线而不是急着改模型结构。从那以后我每次新建高光谱项目都强制走一遍“缓存 patch → 抽画光谱曲线 → 再开训练”的流程归一化参数只算一遍模型跑偏先看数据再调参。这个习惯帮我定位过好几次数据错位的问题希望帮到你。本文还有配套的精品资源点击获取