简介：这套资料面向正在使用PyTorch处理高光谱影像的深度学习开发者，系统讲解如何通过DataLoader高效完成数据读取、预处理、批处理与多线程加载。内容覆盖Dataset子类编写、光谱归一化、batch_size与num_workers调优、shuffle增强、pin_memory加速及自定义collate_fn等关键环节，适合需要处理HSI大规模数据、想提升训练效率的入门到中级学习者。压缩包共7个文件，包含3个Python脚本、2个编译后的pyc文件及2个MAT格式数据集，脚本涵盖工具函数、数据加载器与训练入口，MAT文件为IndianPines等实验数据，整体大小约5.69MB。资源已有612人学习，代码结构清晰，可直接参考或二次修改，帮助读者少走弯路并快速搭建属于自己的高光谱数据加载流程。
1. 聊一聊高光谱数据加载的那些事
最近在跟进一个矿物填图任务，数据是某高光谱传感器推扫出来的 200 波段影像，单景 20000×4000 像素。部署阶段队友用torch把整块 20G 的 GeoTIFF 一次性read()进来，程序直接给干到休眠。原来问题不只出在内存规划上，PyTorch 的 DataLoader 加载高光谱数据跟加载 RGB 图完全不是一回事，波段维度的顺序、样本组织方式、normalize 策略都得重新设计。高光谱数据通常被存成 (Bands, Height, Width) 或 (Height, Width, Bands)，如果拿普通图像分类那套读图逻辑去套，缓冲区爆炸几乎是可以预见的。本文想聊的是我在做“高光谱使用 PyTorch 的 dataloader 加载高光谱数据”任务时攒下的一套方案，覆盖数据组织、Dataset 子类写法、DataLoader 参数微调，以及几个容易让人翻车的细节，写给正在跟高光谱推扫数据搏斗的 Python 开发者和遥感算法工程师。
2. 高光谱数据在 PyTorch 中的存储与组织方式
2.1 高光谱数据格式与通道顺序的坑
高光谱影像的原始格式通常分为三类：BIP（波段按像素交替存储）、BIL（波段按行交替存储）、BSQ（波段按面存储）。这三种格式最直接影响的是你在 Python 端用np.memmap读数据时 shape 的解读。BIP 保存的文件读出来 shape 是 (Height, Width, Bands)，BSQ 是 (Bands, Height, Width)，BIL 介于两者之间，是 (Height, Bands, Width)。大多数 PyTorch 图像任务默认期望 CHW 顺序，所以读进来之后第一件事通常是np.transpose(data, (2, 0, 1))。
工程中我建议不要直接在内存里保留整景高光谱数据。以 200 波段、16bit 存储、4000×1000 像素为例，单景算一下就是 200×4000×1000×2 字节，大约 1.6G。双景就顶不住。常见做法是先把原始数据按块切好，存储成npy或h5py，再做训练集组织。我一般会写下面这段逻辑做预处理和缓存的落地：
import numpy as np import tifffile import os def preprocess_and_store(data_root, output_root, band_groups=4): os.makedirs(output_root, exist_ok=True) for tif_name in os.listdir(data_root): if not tif_name.endswith('.tif'): continue img = tifffile.imread(os.path.join(data_root, tif_name)) # 假设 tif 是 H,W,C 排列，先转成 C,H,W if img.ndim == 3: img = np.transpose(img, (2, 0, 1)) # 按波段分块，方便后面做取块操作 band_step = img.shape[0] // band_groups for g in range(band_groups): sub = img[g*band_step:(g+1)*band_step] np.save(os.path.join(output_root, f"{tif_name[:-4]}_g{g}.npy"), sub)tifffile.imread读出来的通道轴默认在最后，直接存 npy 的时候务必转成波段轴在前。参数band_groups=4是分块数，如果你的显存很紧张，可以再细分一点。这个预处理的产出是多个小 npy 文件，每一个都能被独立映射到 Dataset，避免训练代码直接接触几 G 级的大文件。
2.2 三种常见样本组织模式
高光谱任务里，样本的组织方式决定了 DataLoader 里__getitem__的写法。我接触过三种场景：整图分类、像素分类、patch 分类。
第一种整图分类最简单，每张图是一个样本，__getitem__返回 (C,H,W) 张量和图的标签。这种模式要求每幅图尺寸一致，不一致就得 resize 或 padding。第二种像素分类，例如矿物填图，样本是 (C,) 的光谱向量，__getitem__返回单个像素的光谱和标签。这个方式如果从 npy 全量文件里随机采样，IO 开销极大，因为样本之间不连续。第三种 patch 分类，围绕中心像素切出一个固定大小的窗口，比如 9×9×C，适合 CNN 结构。
在实际工程中，我主推第三种。一个经验做法是先预先生成一个“索引表”，把每张样本的（文件路径、中心坐标、标签）记下来。__getitem__按索引去读取对应的窗口切片，天然支持随机化和多进程预取。下面这段是索引表生成的示意：
import pandas as pd import numpy as np def build_index_table(data_path, window=9, stride=5): img = np.load(data_path, mmap_mode='r') bands, h, w = img.shape records = [] for y in range(window // 2, h - window // 2, stride): for x in range(window // 2, w - window // 2, stride): records.append((data_path, y, x)) return pd.DataFrame(records, columns=['path', 'y', 'x'])stride 参数控制相邻窗口的重叠度，值越大，索引表越短，训练一个 epoch 步数越少。如果你做的是地物分类，往往混合两种方法更实用：用整图做粗分类，再用 patch 级模型对边缘区域复核。
3. 核心改造：为高光谱场景重写 Dataset 与 DataLoader
3.1 Dataset 子类实现：三处必改的地方
PyTorch 原生torchvision.datasets.ImageFolder只支持单波段三通道，直接用在高光谱数据上约等于自杀。必须自己实现torch.utils.data.Dataset子类。核心要改三处：一是文件读取逻辑，二是返回样本的 shape 与 dtype，三是归一化方式。
下面是我在高光谱任务里比较常用的一段模板：
import torch from torch.utils.data import Dataset import numpy as np import os class SpectralPatchDataset(Dataset): def __init__(self, meta_df, window_size=9, transform=None): self.meta_df = meta_df self.window = window_size self.half = window_size // 2 self.transform = transform def __len__(self): return len(self.meta_df) def __getitem__(self, idx): path, y, x, label = self.meta_df.iloc[idx].values img = np.load(path, mmap_mode='r') # 文件句柄不落地 patch = img[:, y-self.half:y+self.half+1, x-self.half:x+self.half+1] patch = np.asarray(patch, dtype=np.float32) # 逐波段归一化，之后每个波段的数值范围独立 mean = patch.mean(axis=(1, 2), keepdims=True) std = patch.std(axis=(1, 2), keepdims=True) patch = (patch - mean) / (std + 1e-5) if self.transform is not None: patch = self.transform(patch) return torch.from_numpy(patch), torch.tensor(label, dtype=torch.long)逻辑说明：mmap_mode='r'的做法不会把整个 npy 文件读进内存，而是按页映射，__getitem__里每次只有 patch 大小的数据进入物理内存。逐波段归一化比整图归一化在高光谱场景里更稳，因为不同波段间信号强度差异太大，比如某些吸收谷波段均值接近 0，而反射峰波段均值几百，用全局归一化会淹没弱波段贡献。
三点要注意的：第一，window必须小心设置，当中心点距离边缘不足 half 时，切片会产生越界，索引表的生成逻辑要保证 y 和 x 在安全范围；第二，如果输入 npy 是 uint16，np.asarray转换成 float32 时会出现整数溢出吗？答案是不会，astype是安全的，但 numpy 的规则是直接截断还是先转换取决于 dtype 转换描述，建议显式写np.asarray(patch).astype(np.float32)，不要依赖隐式转换；第三，transform 的位置应该在归一化之后，不能先做增强再归一化，否则均值重置会抹掉增强带来的增益。
3.2 DataLoader 参数配置与批量读取策略
Dataset 写完了，DataLoader 的参数配置同样影响加载性能与稳定性。处理高光谱数据时，能看到明显收益的有四个参数：num_workers、pin_memory、drop_last和collate_fn。
from torch.utils.data import DataLoader train_loader = DataLoader( train_dataset, batch_size=64, shuffle=True, num_workers=8, pin_memory=True, drop_last=True, collate_fn=None, # 保持默认，因为每个样本 shape 固定 )num_workers设置成 0 时训练速度会断崖式下跌，因为切片裁剪是 CPU 密集操作。但也不要无脑开满，在多任务训练或混合训练场景下，worker 进程数过多会挤占数据增强所用的 CPU 资源。我一般按物理核数的一半设，例如 16 核机器设 8。
drop_last=True在高光谱训练里很重要。patch 数据集最后一小批样本往往窗口重叠度高，会让反向传播的 loss 抖动，直接丢弃更理智。pin_memory=True适合 GPU 训练，这里要求数据是 CPU 内存上的 tensor，__getitem__里我们已经转成 torch.Tensor，满足条件。如果你在__getitem__里返回 numpy 数组，则 DataLoader 会自动做转换，pin_memory 也仍有效。
关于collate_fn，默认行为是沿 batch 维度堆叠，高光谱样本 shape 固定，不需要定制。但如果未来做多模态或者带着坐标信息返回，那你就要自定义它。这块细节我放到改法章节展开。
4. 从数据源到模型输入：预处理流水线设计
4.1 波段选择、量化与归一化顺序
高光谱数据波段之间存在很强的相关性和冗余性。如果不做任何波段筛选，直接把 200 个波段全部喂给模型，显存占用是小事，模型容易被强相关的波段带偏，训练收敛慢还容易过拟合。常见做法是先做 PCA 或者波段选择（比如选信息量最大的 30～50 个波段），再造训练样本。
至于归一化顺序，我坚持先做空间裁剪再做归一化，理由如下：如果先整图归一化再裁剪，耗时严重，因为归一化必须访问所有像素；先裁剪后归一化则每次计算范围是局部窗口，虽然均值波动略有差异，但对于 patch 级模型来说，这种局部统计更接近推理时看到的数据分布。两段比较的直观对比：
# 顺序A：先归一化再裁剪 —— 慢，且均值受全图影响 img = np.load(path, mmap_mode='r') img = (img - img.mean(axis=(1,2), keepdims=True)) / (img.std(axis=(1,2), keepdims=True) + 1e-5) patch = img[:, y-half:y+half+1, x-half:x+half+1] # 顺序B：先裁剪再归一化 —— 快，且分布贴近局部 img = np.load(path, mmap_mode='r') patch = img[:, y-half:y+half+1, x-half:x+half+1] patch = (patch - patch.mean(axis=(1,2), keepdims=True)) / (patch.std(axis=(1,2), keepdims=True) + 1e-5)顺序 B 有两个潜在弱点：若窗口区域内像素全是同一地物，局部均值会偏移，导致同类样本光谱形态被拉平。但对于高光谱分类任务，这个缺点反而约束了模型过分依赖绝对反射率。如果你做的是定量反演任务，比如叶绿素含量反演，那就必须更新为全图归一化，因为它需要保留绝对的辐射信息。
4.2 数据增强的正确姿势
高光谱数据不同于普通图像，像RandomResizedCrop这类基于视觉语义的增强并不适合，因为光谱向量是逐波段的连续曲线，空间裁剪只是在空间维度上做变换，波段维度不能随便丢。常用的增强方式是光谱噪声注入与波段掩码。
class SpectralAugment: def __call__(self, patch): # patch shape: (C, H, W) # 加性噪声，模拟传感器热噪声 noise = np.random.normal(0, 0.01, patch.shape).astype(np.float32) patch = patch + noise # 随机遮盖一部分波段，模拟坏线或缺失 if np.random.rand() < 0.3: mask = np.random.binomial(1, 0.9, size=(patch.shape[0], 1, 1)) patch = patch * mask.astype(np.float32) # 随机翻转空间方向 if np.random.rand() < 0.5: patch = patch[:, :, ::-1] return patch这段增强为什么有效？高光谱数据每个波段本质上是对地物辐射特性的测量，噪声注入模拟了传感器暗电流和环境扰动，波段遮盖则强迫模型不依赖少数强特征波段，提升交叉场景的泛化性。注意几点：空间翻转仅限水平方向和垂直方向，不要做旋转 90 度，因为推扫式传感器的扫描方向会导致地物几何朝向不一致。
4.3 难例挖掘与样本选择策略
高光谱分类任务里，样本数量少且类别极不均衡是常态。矿物填图中，某种矿物可能只有几百个可用像素，而背景占据了绝大多数区域。直接用全图训练，模型很容易进入“全预测为背景”的假收敛状态。
遇到这种情况，我的做法是在索引表阶段做类别重采样，用WeightedRandomSampler替换RandomSampler。这个 sampler 接收一个权重序列，权重高的类别会被更频繁抽到。权重计算我用了平衡误差率的倒数，简单有效。还有一个大家容易忽略的技巧：先训练一个粗糙模型，把置信度低的样本挑出来再做一遍较密的窗口索引，相当于聚焦难例区域。
5. 避坑指南：高光谱 DataLoader 常见问题与排查
5.1 数据增量存档之后，训练集和验证集泄漏
现象：验证集准确率极高，但换一块影像测试立刻崩掉。 原因：索引表是按时序生成的，相邻窗口中心点高度重叠，如果按顺序切分训练集和验证集，同一地物的空间邻接像素同时出现在两边。 解决：按影像文件而不是按像素做划分。例如把索引表按文件分组，前 80% 文件做训练，后 20% 做验证。或者用sklearn.model_selection.GroupKFold，group 参数传文件名。
5.2 DataLoader 长时间卡死在第一个 epoch
现象：训练开始时 GPU 利用率上不去，过几分钟才进入稳定运行。 原因：__getitem__每次np.load(..., mmap_mode='r')都重新打开并映射文件，Windows 下 mmap 句柄不会立即释放，反复打开容易触发系统文件缓存污染。Linux 下一般没这么明显，但 Windows 机器上会表现为“假死”。 解决：用 LRU 缓存包一层文件句柄，或者把 npy 转成h5py的单文件格式，保持一个文件句柄。我通常会把热数据缓存起来：
from functools import lru_cache @lru_cache(maxsize=4) def load_cached(path): return np.load(path, mmap_mode='r')5.3 边界像素被静默裁掉
现象：训练集 map 精度很高，但推理阶段图幅四周边界区域大面积错分。 原因：索引表生成时窗口跨边界则跳过，导致边界像素在整个训练周期中从未参与训练，模型对边界区域天然盲区。 解决：在生成索引表时，对不足窗口的区域采用 padding，允许边界像素的 patch 通过零填充补足窗口。通常在遥感任务里边界低频信息不影响大局。
5.4 归一化后数据方差归零导致训练崩溃
现象：某个波段在某些窗口内所有像素值完全相同，patch 的 std 为 0，除以 0 得到 nan。 原因：高光谱影像经过大气校正后，某些波段在地物均匀区域几乎无波动，特别是水体或阴影区域。 解决：在std上加 epsilon 只能解决 nan 数字，更稳的是对连续 3～5 个波段做统计池化，或直接剔除这些方差为零的波段。加 eps 的做法：
std = patch.std(axis=(1, 2), keepdims=True) + 1e-6实测这种场景下 1e-5 的 eps 仍可能输出大数，我最终选 1e-3，模型稳定性明显提升。
5.5 半精度训练时损失突然变成 nan
现象：换成torch.float16训练几轮后 loss 出现 nan，退回 float32 就没有问题。 原因：半精度有效位有限，高光谱影像均值较大，逐波段归一化后某些波段数值仍大量分布在 0.01 以下，半精度下直接变成 0，梯度传到归一化层出现除以极小值的风险。 解决：对归一化后的数据做全局缩放，乘个系数 10 或者约束到 [-1, 1] 范围，同时给归一化输出 clamp 做下限保护：
patch = torch.clamp(patch, min=-1.0, max=1.0)6. 扩展改造：让 DataLoader 通用化的三处实战改动
这块我分享三个实际改造点。第一个是支持多文件目录，高光谱项目通常跨多景影像，每个文件波段数不一致（例如有的 128 波段，有的 200 波段）。强行要求所有文件统一波段数不太现实。我的方式是增加一个band_selector参数，在预读时自动把波段对齐到某一个固定集合。具体做法是对所有文件先跑一遍波段元信息提取，每个文件记下 (filename, bands)，然后把__getitem__里np.load的结果按band_selector索引重排。
第二个改造是自定义collate_fn实现动态 padding。当 patch 尺寸不一致或者要附带坐标与邻域光谱均值信息时，默认的堆叠方法会炸。通常我这样写：
def collate_patches(batch): patches, labels = zip(*batch) # 这里假设尺寸不同，需要双线性插值统一到 9x9 resized = torch.stack([ torch.nn.functional.interpolate(p.unsqueeze(0), size=(9, 9), mode='bilinear').squeeze(0) for p in patches ]) return resized, torch.stack(labels)第三个改造是持久化缓存。高光谱预处理代价高，每个 epoch 重复跑归一化和裁剪太可惜。我加了一个无损缓存步骤，将预处理好的 patch 按哈希文件名存到~/.cache/spectral，第二次读取直接跳过计算。这个方法让训练迭代速度提升约 40%。
再补充两个我自己的习惯：检查 dataloader 是否把标签同步 shuffle 正确，WeightedRandomSampler和BatchSampler混用时容易出现标签顺序错位；另外num_workers增大后内存峰值会接近num_workers × batch_size × patch_size，预估显存时要按 worker 数乘进去，不然容易出现 OOM 翻车。
有一次我排查了很久才发现是drop_last参数的锅——它只在len(dataset)不能整除以 batch size 时生效，但不少人误以为它会丢掉验证集里的尾部样本。从那以后我每次更换数据集组织方式，都会先打印一个 batch 的 shape 和标签分布，确认不是数据加载的玄学问题，再进模型训练循环。希望这篇文章能帮到正在跟高光谱 DataLoader 搏斗的你。
本文还有配套的精品资源，点击获取
