简介面向PyTorch开发者的高光谱数据加载配套代码包针对高光谱图像光谱通道多、数据量大的特点提供基于数据集类与加载器的完整实现解决批处理、多线程读取、随机打乱与数据增强等常见难题。包内包含印度松等公开数据集的读取示例适合正在学习高光谱分类或需要优化数据管线的中初级开发者。压缩包共有七个文件包括三个Python源码、两个编译类文件及两个矩阵格式数据文件整体大小约五点六九兆字节。源码覆盖数据集类构建、加载器配置与训练入口数据文件可供实验测试编译文件便于快速验证项目结构区分了数据存放与工具模块可对照学习。已有六百余人学习下载。通过实际源码可掌握数据集接口实现、归一化与预处理方法、加载器线程数与内存优化技巧以及自定义批量合并函数处理多维数组的方法能够快速迁移到自己的高光谱项目中从而提升模型训练效率。1. 高光谱数据配 PyTorch DataLoader从手撕 Numpy 到一键迭代第一次拿高光谱数据练手的时候我干的蠢事是把整块 HDF5 读进内存然后手写循环一行行喂给模型。跑了半天发现数据是 200 个波段、上万条像元内存直接爆掉训练步数还没走完就卡死。后来我把高光谱数据的加载逻辑重构成 PyTorch DataLoader 的标准形态问题立刻消失显存占用稳定、shuffle 和数据增强随手即用、多进程预取把加载时间压到几乎为零。这篇笔记就是把这段重构过程完整拆开讲清楚怎么用 PyTorch 的 Dataset DataLoader 正确加载高光谱数据适合正在做高光谱图像分类、地物识别或者波段筛选的开发者也适合刚把 Python 和 PyTorch 装好、手头有一份高光谱数据却不知道怎么喂给模型的新手参考。2. 数据准备与整体架构先把 .h5 和 .tif 变成模型认识的数字2.1 高光谱数据到底长什么样高光谱数据跟普通 RGB 图像最大的区别在于波段数。相机拍的是三个通道高光谱传感器一次能拿到几十到几百个窄波段比如 400nm 到 1000nm 区间按 10nm 间隔采样就得到 60 个波段。每一条像元对应一条连续光谱曲线这也是高光谱能够区分精细地物的原因。常见的高光谱数据存储格式有这么几种HDF5.h5适合存大型遥感影像、ENVI 标准格式.dat .hdr常见于航空高光谱数据、GeoTIFF 则更有普适性。还有一种更朴素的做法数据直接以原始二进制.raw保存配合一个 shape 元组就能读出来。不管格式怎么变数据本身的维度通常是 H × W × B也就是高度、宽度、波段数和 PyTorch 默认的 B × C × H × W 相差一个维度顺序。这个顺序问题我在第 4 章的避坑里会展开讲它是新手翻车的高频点之一。另外还有标签数据分类任务里每一条像元或者每一个 patch 对应一个类别编号0 到 N-1需要和特征数据严格对齐。2.2 预处理管道逃不掉的标准化高光谱数据的数值范围随波段差异很大某些波段反射率集中在 0.1 附近另一些可能在 0.8 上下波动。如果原样喂给模型梯度计算容易被数值大的波段主导。常见做法是先做逐波段的标准化也就是对每个波段单独计算均值和标准差让数值分布接近零均值单位方差。也有人喜欢用 MinMax 归一化压到 [0, 1] 区间这取决于模型是否有 BatchNorm 层。我在项目里的习惯是先查一遍每个波段的统计量再决定策略。假如某个波段有大量空值或饱和值MinMax 会被极端值拉偏搞出一片全黑或全亮的图像这种情况下逐波段标准化更稳。预处理逻辑我会写成一个函数放在 Dataset 内部而不是单独开线程做理由是它只做数值变换不涉及数据读取放到getitem里跑的代价很便宜。2.3 目录结构一个高光谱加载工程的标准姿势拆解这类项目的时候我喜欢先看目录结构。下面的布局是某个高光谱分类模拟项目里我用过的组织形式结构简单但适合二次扩展hyperspectral_loader/ ├── data/ # 原始数据与标签存放目录 │ ├── scene.h5 # 高光谱影像数据 │ ├── label.npy # 标签图H×W 整数数组 │ └── metadata.json # 波段信息、波长范围 ├── loader/ │ ├── __init__.py │ ├── dataset.py # Dataset 子类核心加载逻辑 │ ├── transforms.py # 波段标准化、数据增强相关函数 │ └── split.py # 训练集、验证集、测试集划分 ├── train.py # 训练入口 └── config.yaml # 参数配置文件dataset.py 负责把原始数据按索引取出来transforms.py 负责数值变换和增强split.py 负责样本划分。划分逻辑单独放一个文件是因为高光谱数据的样本划分经常要避免同一区域的数据串到两个集合里写成独立模块方便后续加空间分块逻辑。2.4 样本组织方式Patch 级别还是像元级别高光谱数据落到 DataLoader 里每一条样本的形态取决于任务类型。像元级别分类每一条样本是一个 B 维的光谱向量标签是对应位置的类别编号Patch 级别分类每一条样本是以某个像元为中心的邻域窗口比如取 9×9 的窗口样本形状是 9 × 9 × B模型能看到空间上下文精度一般更好但数据量会膨胀。我在模拟项目 X 里同时实现了两种模式切换开关就写在 dataset.py 里。Patch 模式需要额外处理边界窗口滑到影像边缘时常见做法是零填充或者镜像填充镜像填充对高光谱数据的连续性更友好我默认用后者。3. 核心模块拆解一个 Batch 从索引到张量的完整链路3.1 Dataset 子类的骨架PyTorch 的 DataLoader 本身不负责读取数据它只负责调度。真正干活的是继承自 torch.utils.data.Dataset 的子类必须实现len和getitem两个方法。我的 dataset.py 长这样import torch import numpy as np import h5py from torch.utils.data import Dataset class HyperspectralDataset(Dataset): 高光谱数据加载器支持像元级与Patch级两种取样本模式 def __init__(self, h5_path, label_path, patch_size1, modepixel, transformNone): # 只在这里打开HDF5文件句柄读取数据主体避免反复IO with h5py.File(h5_path, r) as f: self.data f[scene][:] # shape: (H, W, B) self.labels np.load(label_path) # shape: (H, W) self.H, self.W, self.B self.data.shape self.patch_size patch_size self.mode mode self.transform transform # 生成所有有效样本的坐标索引 self.sample_indices self._build_index() def _build_index(self): 收集所有标签非空的像元坐标 ys, xs np.where(self.labels 0) indices list(zip(ys.tolist(), xs.tolist())) return indices def __len__(self): return len(self.sample_indices) def __getitem__(self, idx): y, x self.sample_indices[idx] if self.mode pixel: patch self.data[y, x, :] else: half self.patch_size // 2 top max(0, y - half) bottom min(self.H, y half 1) left max(0, x - half) right min(self.W, x half 1) patch self.data[top:bottom, left:right, :] # 边缘补齐到patch_size × patch_size patch self._pad(patch) label self.labels[y, x] patch torch.from_numpy(patch).float() if self.transform: patch self.transform(patch) return patch, label关键逻辑在getitem里。每次 DataLoader 通过索引调用这个函数函数根据事先算好的坐标索引定位到影像上某个像元或邻域窗口取出特征和标签经过标准化处理后返回一个 torch.Tensor 和一个整数标签。DataLoader 只负责把这些返回值堆叠成 batch所以返回值必须保证维度一致。数据主体在init里一次性载入内存防止训练时反复读磁盘。HDF5 文件句柄用完立刻关闭数据以 numpy 数组保存在内存里后续操作全部基于内存数组。3.2 边界补齐函数Patch 模式下最头疼的就是边界。影像边缘的窗口滑出范围时会截断截断后样本形状不一致DataLoader 在 collate 时报错。我的处理方式是镜像填充保留光谱连续性def _pad(self, patch): 把Patch补齐到统一大小用镜像填充保持光谱趋势连续 Hp, Wp patch.shape[0], patch.shape[1] target self.patch_size pad_h_before (target - Hp) // 2 pad_h_after target - Hp - pad_h_before pad_w_before (target - Wp) // 2 pad_w_after target - Wp - pad_w_before if pad_h_after 0 and pad_w_after 0 and \ pad_h_before 0 and pad_w_before 0: return patch return np.pad(patch, ((pad_h_before, pad_h_after), (pad_w_before, pad_w_after), (0, 0)), modesymmetric)np.pad 的 modesymmetric 做镜面扩展扩展后的 Patch 内部不会出现生硬的数值跳变。三个维度的填充增量不一样前两个维度是空间坐标最后一个波段维度不填充。边界补齐的代价是单条样本多一点内存换来统一形状DataLoader 的 collate 流程才能顺利地把样本堆成 batch。3.3 标准化写进 transform标准化不放在 Dataset 之外的另一个原因是train/val/test 三份数据要共用同一套均值方差。如果每份数据算一次自己的均值和方差等于给验证集也做了和训练集不同的归一化这属于数据泄漏。我在 transforms.py 里写了一个可调的标准化器class SpectralStandardize(): 逐波段标准化均值方差从训练集统计得到 def __init__(self, mean, std): self.mean np.array(mean, dtypenp.float32) self.std np.array(std, dtypenp.float32) def __call__(self, patch): return (patch - self.mean) / (self.std 1e-6)这里的 mean 和 std 长度必须等于波段数计算时只在训练集上统计然后实例化三个 transform 分别给训练、验证、测试用。统计均值和标准差的代码我放在 split.py 里先按训练集索引取数据再对每个波段汇总def compute_mean_std(dataset, indices): 采集训练集上所有样本的逐波段统计量 all_patches [] for i in indices: patch, _ dataset[i] all_patches.append(patch.numpy()) concat np.stack(all_patches, axis0) # (N, H, W, B) mean concat.mean(axis(0, 1, 2)) std concat.std(axis(0, 1, 2)) return mean, std这种先把所有训练样本统计一遍再进入训练的做法适合数据量可控的高光谱单场景。数据量特别大时改成抽样统计取训练集的子集算均值方差即可差别不大。3.4 DataLoader 组装与 batch 生成Dataset 写好后DataLoader 的组装就比较机械了from torch.utils.data import DataLoader from loader.dataset import HyperspectralDataset from loader.transforms import SpectralStandardize train_dataset HyperspectralDataset( h5_pathdata/scene.h5, label_pathdata/label.npy, patch_size9, modepatch, transformtrain_transform ) train_loader DataLoader( train_dataset, batch_size64, shuffleTrue, num_workers4, drop_lastTrue, pin_memoryTrue )batch_size 表示一个批次包含多少条样本Patch 模式下每单位内存消耗跟 patch_size 平方成正比显存不够就调小 batch_size。shuffleTrue 会在每个 epoch 开始前打乱样本顺序保证类别分布随机。num_workers4 开 4 个子进程做数据预取瓶颈在 IO 时效果明显但 Windows 下 num_workers 大于 0 需要把训练代码放进 ifname main 保护块否则多进程递归报错。pin_memoryTrue 配合 GPU 训练能减少 H2D 拷贝时间CPU 训练时开不开无所谓。drop_lastTrue 丢弃最后不足一个 batch 的样本避免归一化统计出错。4. 避坑指南高光谱加载高频翻车现场与排查手段4.1 通道顺序错乱训练 Loss 不降反升现象模型可以跑但 loss 从头到尾几乎不下降或训练曲线异常震荡。检查数据形状发现输入是 (B, H, W, C)模型实际接收的是 (C, H, W)。原因高光谱原始数据默认是 H × W × BPyTorch 卷积层期望的输入是 B × C × H × W。把两者直接喂给模型相当于把波段维度当成了 batch 维度空间维度被当成通道整个数据的语义全变了。解决在 transform 里或 Dataset 返回前把张量从 HWC 转成 CHW一步到位# 在 _pad 之后标准化之前或之后都行 patch torch.from_numpy(patch).float() patch patch.permute(2, 0, 1) # (H, W, B) - (B, H, W)4.2 Windows 下 num_workers 报错程序反复重启现象设置 num_workers4 后运行就报 RuntimeError内容涉及 spawn、freeze_support进程反复被拉起。原因Windows 上多进程数据加载使用 spawn 方式启动子进程会重新导入主模块导致递归执行训练代码。解决训练脚本入口加保护块。这不是高光谱特有的问题但高光谱数据单次加载耗时较长很多人下意识把 num_workers 调大反而被这个坑绊住。if __name__ __main__: train()4.3 所有样本标签全为 0模型学了个寂寞现象训练起来很快准确率始终在某个常数附近混淆矩阵发现模型只预测少数几个类别。原因标签文件读取时把背景类通常用 0 表示当成了真实类别参与训练或者标签和特征数据的空间位置没对齐导致特征和标签错位。解决加载标签后先做一次类别分布统计unique, counts np.unique(labels, return_countsTrue) print(dict(zip(unique, counts)))如果发现某个类别占了 95% 以上需要检查标签语义以及样本划分是否均衡。高光谱影像里背景像元数量普遍很多我习惯在 _build_index 里直接把标签等于背景类比如 -1 或 255的坐标过滤掉只保留有效样本。4.4 HDF5 文件句柄泄漏内存只涨不降现象长时间训练过程中内存占用持续上升最终系统卡顿或 OOM 崩溃。原因在init中反复打开 HDF5 文件却不关闭或者每次getitem都执行 h5py.File 操作。解决init里一次性把数据载入 numpy 数组并确保文件句柄用 with 语句管理不止是省内存还能避免大量小文件打开带来的 IO 抖动。5. 验证与可视化Dataloader 跑通后先看一眼数据5.1 人工检查 Batch 输出训练之前先手动迭代一个 batch确认数据形状、数值范围、标签值都符合预期。这一步能拦截掉 90% 以上的低级错误写几行调试代码的代价远小于训练半天之后才发现数据有问题。sample next(iter(train_loader)) features, labels sample print(features.shape) # 期望 (batch_size, B, H, W) 或 (batch_size, B) print(features.min(), features.max()) print(labels.unique()) # 期望类别编号都在有效范围内特征最小值最大值如果出现 NaN 或极端值优先检查原数据是否有空值。标准化后的数据 min/max 通常不会离均值太远但出现几百几千的值也不一定错个别波段可能有强反射峰关键是看分布是否和预处理逻辑匹配。5.2 伪彩色合成可视化高光谱数据没法像 RGB 图一样直接显示通常从几百个波段里抽 3 个代表性波段合成伪彩色图。我常用的方案选近红外、红、绿三个波段对应 RGB 三个通道。可视化能直观看出 patch 窗口取得到底对不对有没有整片黑掉标签边界和影像地物轮廓是否吻合。这段代码可以单独放在可视化脚本里不进训练主流程import matplotlib.pyplot as plt def visualize_rgb(data, rgb_indices(30, 20, 10)): 从高光谱立方体中抽取三个波段合成伪彩色图 bands [data[:, :, i] for i in rgb_indices] rgb np.stack(bands, axis-1).astype(np.float32) # 逐个通道分位数拉伸避免个别极端值压掉整体亮度 result np.zeros_like(rgb) for c in range(3): vmin, vmax np.percentile(rgb[:, :, c], (2, 98)) result[:, :, c] np.clip((rgb[:, :, c] - vmin) / (vmax - vmin 1e-6), 0, 1) return resultpercentile 拉伸比单纯除以最大值靠谱高光谱某些波段存在镜面反射或噪声像元数值可能高出正常范围几个量级不裁剪会把图像整体压暗。5.3 单个样本的可视化检查Batch 输出验证之后再抽查几个单条样本。我从 train_dataset 里按索引取样本打印标签对应的类别名再把 patch 的逐波段均值光谱画出来。观察同一类别不同 patch 的光谱曲线走势是否接近不同类别曲线是否区分明显。如果同一类别的光谱曲线分布散乱说明标签对齐有问题要回头检查标签图是否做了几何校正。这一步在高光谱数据加载项目里特别值得做因为它能直接暴露数据源层面的错位问题。6. 二次开发与改造把基础 Loader 扩展成顺手的数据管道拿到这套 Dataset/DataLoader 之后最常做的改造集中在数据增强、采样策略、缓存优化三个方向。高光谱数据做随机裁剪或旋转时波段维必须和空间维保持同步numpy 的高维数组操作天然满足这个条件所以增强函数和普通图像的写法区别不大只需保证 transform 接收的是 (H, W, B) 形状的张量再按需返回并调整顺序。第一个建议是加缓存层。高光谱 patch 加载虽然快但每个 epoch 重复读取同一批 patch 仍然浪费 IO。我在 Dataset 里加了一个简单的 dict 缓存按索引把最近访问过的样本存下来内存够用就直接缓存全部训练集样本省掉重复切片和填充运算。显存紧张时改用 lru_cache 控制缓存上限实测训练速度能快 10% 到 20%尤其是 patch_size 较大时收益更明显。第二个建议是类别均衡采样。高光谱地物分类的标签分布天然不平衡某个类别可能只占全部像元的 0.5%。直接在init里对每个类别的样本数量做 weight 统计然后传给 WeightedRandomSamplerDataLoader 每次采样时按权重调整各类别被抽中的概率。实现时先统计每个类别样本数再算权重只改 DataLoader 的 sampler 传参Dataset 代码不用动。第三个建议是波段子集选择。有些任务不需要全部波段比如植被指数只需要红边和近红外附近几个波段。我在 Dataset 初始化里加了一个 band_indices 参数传入需要保留的波段索引数组getitem内部按索引切片即可省内存也省训练时间。具体做的时候可以先跑一个随机森林或者互信息筛选把特征重要性排前 N 的波段下标存进配置再传给 Dataset。这套代码后来被我用到某高校实验室的高光谱地物分类项目中从原始 HDF5 到模型训练的时间从原来的一个下午缩短到一顿午饭的工夫。回头再看最有价值的其实是第 3 章的坐标索引和数据对齐那一步——不管增强、采样还是缓存只要索引和可视化保持正确改起来就顺畅。从那以后我每次上手新数据都会强制走一遍先可视化、再训练的标准流程。希望帮到你。本文还有配套的精品资源点击获取