高光谱数据进PyTorch：dataloader设计与调优实战指南

简介：这份资源面向使用PyTorch做深度学习、需要处理高光谱图像的研究者与工程实践者，聚焦于用DataLoader高效加载与预处理高光谱数据这一关键环节。包内共7个文件，以3个Python脚本为核心，辅以2个编译缓存文件与2个mat格式数据文件，压缩包约5.69MB，脚本涵盖数据集封装、加载流程与训练入口，mat文件对应高光谱数据样本，便于直接运行与二次修改。已有613人学习下载，说明其在同类实践中有一定参考价值。读者可从中获得一套可复用的Dataset与DataLoader实现思路，理解数据读取、归一化预处理、批大小设置、多线程加载、随机打乱与自定义collate_fn等环节如何配合，并借助示例数据快速搭建训练流程、排查内存与加载效率问题，适合作为高光谱分类任务的入门与调试参考。

1. 高光谱数据进 PyTorch：为什么 dataloader 是整条训练链最容易被低估的一环

高光谱数据不像普通 RGB 图像那样丢进ImageFolder就能跑。一个典型的推扫式高光谱立方体是H×W×B的三维结构，B 动辄上百甚至数百个波段，单文件轻松几百 MB 到几个 GB，格式还可能是.mat、.tif、.hdr/.img、.npy混着来。很多人第一次用 PyTorch 做高光谱分类，模型结构抄得挺顺，结果卡在数据加载上：内存爆、训练慢、标签对不上、波段顺序错乱。这篇笔记拆的就是「高光谱 + PyTorch dataloader」这套组合拳，从数据组织、Dataset 写法、采样策略到多进程踩坑，给一份能直接照着复现的落地路径。适合已经会写基础Dataset、但一碰高光谱就翻车的从业者，也适合想把实验室里那套 MATLAB 预处理流程迁到 Python 训练链的同学。

2. 高光谱 Dataset 的骨架：从立方体到张量的四步拆解

高光谱 dataloader 的核心矛盾只有一个：数据是三维立方体，而 PyTorch 的DataLoader期望的是「一条样本 = 一个张量 + 一个标签」。所以整件事的本质是把立方体切成样本，再让__getitem__稳定地吐出一致的形状。这一步没设计好，后面 batch size、归一化、增强全是空中楼阁。

2.1 先定样本粒度：像素级、patch 级还是整图级

高光谱分类任务里，样本粒度直接决定 Dataset 怎么写。常见三种：

patch 级是绝大多数高光谱分类论文和工程的做法：以某个像素为中心取p×p的空间窗口，保留全部 B 个波段，标签取中心像素的类别。这样单样本是(B, p, p)，p一般取 9、11、13 这类奇数，保证中心对齐。选 patch 级的理由是它同时利用了空间邻域和光谱维度，而像素级丢掉了空间信息，整图级又对显存太不友好。

我一般会把预处理和采样分开：预处理阶段把整个立方体读进内存（或 memmap），采样阶段只做切片和归一化。这样__getitem__里没有 IO，多进程才跑得起来。

2.2 用 NumPy memmap 扛住大立方体

高光谱立方体动辄上 GB，直接np.load全读进内存，开num_workers>0时每个 worker 再复制一份，内存直接翻几倍。稳妥做法是先落成.npy，再用np.load(..., mmap_mode='r')映射，让操作系统按页调度。

import numpy as np import os def cube_to_npy(mat_path, npy_path): """把预处理好的立方体存成 npy，供后续 memmap 读取。 假设 mat 里已经取出 data 和 gt 两个数组。 """ from scipy.io import loadmat m = loadmat(mat_path) data = m['data'].astype(np.float32) # (H, W, B) gt = m['gt'].astype(np.int64) # (H, W) # 高光谱常见坑：波段在最后一维还是第一维，务必确认 assert data.ndim == 3 np.save(npy_path.replace('.npy', '_data.npy'), data) np.save(npy_path.replace('.npy', '_gt.npy'), gt) def open_cube(npy_path): data = np.load(npy_path.replace('.npy', '_data.npy'), mmap_mode='r') gt = np.load(npy_path.replace('.npy', '_gt.npy'), mmap_mode='r') return data, gt

这里的关键点是mmap_mode='r'：它不会把整个文件读进 RAM，而是返回一个内存映射数组，切片时才真正读盘。astype(np.float32)要在存 npy 之前做掉，别留到__getitem__里，否则每个样本都做一次类型转换，白白拖慢。assert data.ndim == 3是我血泪经验加的一道保险——不同来源的高光谱数据波段维位置不统一，.mat里可能是(B, H, W)，不检查就会在后面切片时静默出错。

2.3 Dataset 与 patch 采样：索引表比循环更靠谱

patch 级采样的经典写法是预先构建「有效像素索引表」，把所有标签非零（或非背景）的像素坐标存成(N, 2)数组，__len__返回 N，__getitem__按索引取中心坐标再切 patch。这样避免了在__getitem__里遍历找有效像素。

import numpy as np import torch from torch.utils.data import Dataset class HyperPatchDataset(Dataset): def __init__(self, data, gt, patch_size=11, ignore_label=0): self.data = data # memmap (H, W, B) self.gt = gt # memmap (H, W) self.p = patch_size self.half = patch_size // 2 # 预构建有效像素索引，排除背景/无效标签 ys, xs = np.where(gt != ignore_label) self.coords = np.stack([ys, xs], axis=1).astype(np.int64) # 按波段做全局统计，供归一化使用（只算一次） self.mean = None self.std = None def __len__(self): return len(self.coords) def _pad_slice(self, y, x): """边界像素用反射填充，保证 patch 尺寸恒定。""" H, W, _ = self.data.shape y0, y1 = y - self.half, y + self.half + 1 x0, x1 = x - self.half, x + self.half + 1 patch = self.data[max(0, y0):min(H, y1), max(0, x0):min(W, x1), :] # 需要填充的量 pad_top = max(0, -y0) pad_bottom = max(0, y1 - H) pad_left = max(0, -x0) pad_right = max(0, x1 - W) if any([pad_top, pad_bottom, pad_left, pad_right]): patch = np.pad( patch, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), mode='reflect' ) return patch def __getitem__(self, idx): y, x = self.coords[idx] patch = self._pad_slice(y, x) # (p, p, B) patch = patch.astype(np.float32) # 转成 (B, p, p)，符合 PyTorch 通道优先约定 patch = np.transpose(patch, (2, 0, 1)) label = int(self.gt[y, x]) return torch.from_numpy(patch), torch.tensor(label, dtype=torch.long)

逻辑说明：coords在__init__里一次性算好，__getitem__只做切片和转置，这是多进程能提速的前提。_pad_slice处理边界像素——高光谱图像边缘的像素如果直接切会得到小于p×p的 patch，batch 拼不起来，用reflect填充比补零更保留光谱连续性。np.transpose(patch, (2, 0, 1))把(p, p, B)变成(B, p, p)，这一步顺序错了模型第一层卷积的in_channels就对不上，是新手最常见的翻车点之一。

参数上，patch_size取奇数，ignore_label按数据集背景值设（很多高光谱 GT 用 0 表示未标注）。mean/std留了字段但没在__getitem__里用，是因为归一化我倾向放到collate_fn或模型前的 transform 里统一做，避免每个样本重复算。

2.4 collate_fn 与归一化：把统计量算在 batch 上

逐样本归一化会让每个 patch 的均值方差不同，破坏光谱一致性。更稳的做法是在collate_fn里对整批做归一化，或者干脆用训练集全局统计量。下面给一个在collate_fn里做逐波段标准化的写法：

import torch def hyper_collate(batch): patches, labels = zip(*batch) x = torch.stack(patches, dim=0) # (N, B, p, p) y = torch.stack(labels, dim=0) # (N,) # 逐波段标准化：在 (N, p, p) 维度上统计每个波段 mean = x.mean(dim=(0, 2, 3), keepdim=True) std = x.std(dim=(0, 2, 3), keepdim=True) + 1e-6 x = (x - mean) / std return x, y

mean(dim=(0, 2, 3))保留的是波段维，得到(1, B, 1, 1)，正好广播到整批。+1e-6防除零。注意这个统计是 batch 内的，训练时 batch 越大越稳；如果要做严格的训练/验证一致，应该在训练集上先跑一遍算全局mean/std存下来，验证和推理时直接用固定值。两种做法都常见，前者省事，后者更严谨，看你对复现精度的要求。

3. DataLoader 参数调优：num_workers、batch 与采样器的组合边界

Dataset 写对了只是及格，DataLoader的参数配错照样让训练慢到怀疑人生。高光谱的特殊性在于单样本大（B×p×p浮点）、预处理重（切片 + 转置 + 归一化），所以 worker 数量和 batch size 的权衡比普通图像更敏感。

3.1 num_workers 不是越大越好

num_workers的本质是用多进程掩盖__getitem__的耗时。高光谱 patch 采样如果数据在 memmap 上，单次切片很快，worker 开 4～8 通常够；但如果每个样本还要做复杂增强（光谱扰动、随机波段丢弃），可以往上加。判断标准很简单：看 GPU 利用率，如果训练时 GPU 长期低于 70%，说明数据供给跟不上，加 worker；如果 CPU 已经打满、内存飙升，说明 worker 太多在抢资源。

from torch.utils.data import DataLoader train_loader = DataLoader( train_set, batch_size=64, shuffle=True, num_workers=8, pin_memory=True, persistent_workers=True, collate_fn=hyper_collate, drop_last=True, )

pin_memory=True让数据进 pinned 内存，配合 GPU 的异步拷贝能省一点时间，前提是你确实在用 CUDA。persistent_workers=True避免每个 epoch 重建 worker 进程，高光谱数据集初始化（建索引、算统计）不便宜，这个参数能省掉重复开销。drop_last=True在 batch 内做归一化时建议开，避免最后一个不满的 batch 统计量偏差过大。

3.2 分层采样解决类别极不平衡

高光谱场景分类的类别分布往往极度不均，某几类占了大半像素，稀有类只有几十个样本。直接shuffle=True会让稀有类在 batch 里几乎不出现，模型学不动。常见做法是用WeightedRandomSampler按类别频率反比给权重。

import numpy as np from torch.utils.data import WeightedRandomSampler def build_weighted_sampler(dataset): labels = np.array([int(dataset.gt[y, x]) for y, x in dataset.coords]) classes, counts = np.unique(labels, return_counts=True) # 类别权重 = 总数 / (类别数 * 该类样本数) class_weight = {c: len(labels) / (len(classes) * n) for c, n in zip(classes, counts)} sample_weight = np.array([class_weight[l] for l in labels], dtype=np.float64) return WeightedRandomSampler( weights=torch.from_numpy(sample_weight), num_samples=len(sample_weight), replacement=True, )

replacement=True是必须的，稀有类要被重复抽才能凑够num_samples。num_samples设成和数据集等长，保证每个 epoch 的步数和普通采样一致。这里有个容易忽略的点：用了WeightedRandomSampler就不要再设shuffle=True，两者冲突，DataLoader会报错或行为异常。

3.3 训练/验证划分要按空间块切，不能随机打散

高光谱数据空间相邻像素高度相关，随机按像素划分训练/验证集会造成「空间泄漏」——验证集里的像素和训练集里的像素在空间上挨着，模型等于见过答案，指标虚高。正确做法是按空间块划分，比如把图像切成若干不重叠的矩形区域，整块进训练或验证。

def spatial_split(coords, H, W, block=50, val_ratio=0.2, seed=42): """按 block×block 空间块划分，避免相邻像素泄漏。""" rng = np.random.default_rng(seed) block_ids = (coords[:, 0] // block) * (W // block + 1) + (coords[:, 1] // block) uniq = np.unique(block_ids) rng.shuffle(uniq) n_val = int(len(uniq) * val_ratio) val_blocks = set(uniq[:n_val]) val_mask = np.array([b in val_blocks for b in block_ids]) return coords[~val_mask], coords[val_mask]

block取 50 左右是个经验值，太小起不到隔离作用，太大又会让验证集空间分布偏离。这个划分逻辑要在构建 Dataset 之前做，把coords拆成训练和验证两份分别喂给两个 Dataset 实例。很多人图省事用random_split，指标好看但上线就崩，这是高光谱任务里最隐蔽的坑之一。

4. 避坑与排查：高光谱 dataloader 的五条血泪记录

4.1 现象：训练 loss 不降，检查发现标签全错位

原因：gt和data的空间尺寸不一致，或者读入时做了 resize/裁剪但只处理了其中一个。高光谱的 GT 常是低分辨率标注图，需要和立方体对齐。

解决：在 Dataset 初始化时断言data.shape[:2] == gt.shape[:2]，不一致就先做最近邻上采样对齐，别用双线性插值——标签插值会产生不存在的类别。

4.2 现象：开多 worker 后内存暴涨，甚至被系统 kill

原因：num_workers>0时每个 worker 是独立进程，如果 Dataset 在__init__里把整个立方体np.load进内存，每个 worker 各持一份，内存乘以 worker 数。

解决：用mmap_mode='r'做内存映射，或者把大数组放到共享内存。我一般直接在__init__里只存路径，__getitem__里按需 memmap 切片，虽然每次多一次映射开销，但内存稳。

4.3 现象：GPU 利用率忽高忽低，epoch 时间波动大

原因：__getitem__里有随机性很强的重操作（比如随机选波段做复杂变换），或者 worker 之间负载不均。

解决：把重预处理前移到离线阶段，__getitem__只做轻量切片；persistent_workers=True减少进程重建；必要时用prefetch_factor调大预取。

4.4 现象：验证集准确率异常高，换一批数据就崩

原因：随机按像素划分导致空间泄漏，训练和验证像素空间相邻。

解决：改用空间块划分（见 3.3），并在报告指标时说明划分方式。这个坑不解决，后面所有调参都是在拟合泄漏。

4.5 现象：波段顺序和预训练模型对不上，迁移学习效果差

原因：不同来源的高光谱数据波段排列、波长范围不同，直接套用别人预训练权重，通道语义错位。

解决：建立波段-波长映射表，按目标模型的波段顺序重排；没有对应波段的做插值或截断。这一步在 Dataset 之外单独做一次，存成新 npy，别在__getitem__里每次重排。

5. 进阶：用固定统计量 + 缓存把复现精度锁死

前面几章把 dataloader 跑通了，但要做到「换台机器、换个时间跑，指标还能对上」，还得解决两个变量：归一化统计量的随机性和数据读取的重复开销。我的做法是训练前先跑一遍全量统计，把逐波段mean/std存成.npz，训练和验证都读这个固定值；同时对高频访问的 patch 做一层轻量缓存。

import numpy as np import torch def compute_global_stats(dataset, batch_size=256): """遍历训练集，算逐波段全局 mean/std，存盘复用。""" loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4, collate_fn=lambda b: torch.stack([x for x, _ in b])) n = 0 s = None sq = None for x in loader: # x: (N, B, p, p) x = x.double() b = x.shape[1] if s is None: s = torch.zeros(b, dtype=torch.float64) sq = torch.zeros(b, dtype=torch.float64) # 在 (N, p, p) 上累加每个波段的和与平方和 s += x.sum(dim=(0, 2, 3)) sq += (x ** 2).sum(dim=(0, 2, 3)) n += x.shape[0] * x.shape[2] * x.shape[3] mean = s / n std = torch.sqrt(sq / n - mean ** 2) + 1e-6 np.savez('hyper_stats.npz', mean=mean.numpy(), std=std.numpy()) return mean, std def load_stats(path='hyper_stats.npz'): d = np.load(path) return torch.from_numpy(d['mean']).float(), torch.from_numpy(d['std']).float()

这段的关键是double()累加：高光谱波段多、像素多，用 float32 累加平方和容易丢精度，mean和std会偏。n统计的是「样本数 × 空间像素数」，因为每个波段在每个空间位置都有一个值。算完存.npz，训练时load_stats读固定值，在collate_fn里用(x - mean[None, :, None, None]) / std[None, :, None, None]广播即可，彻底消除 batch 内统计带来的波动。

缓存那层我一般只在调试阶段用，比如把验证集所有 patch 预取成一个大的torch.Tensor放内存，省掉反复切片。生产训练不建议全缓存，内存扛不住，用 memmap + 固定统计量已经够稳。

还有一个验证技巧：固定随机种子后，连续跑两次训练，对比前 10 个 batch 的x.mean()和x.std()，如果两次不一致，说明 dataloader 里还有没锁死的随机源（常见是__getitem__里用了全局np.random而不是传入的 generator）。这个检查我每次搭新 dataloader 都会走一遍，比事后对着波动的 loss 曲线猜原因省事得多。

从那以后我每次接一个新的高光谱数据集，都强制先跑「形状断言 → 空间块划分 → 全局统计落盘 → 双跑一致性检查」这四步，再动模型。希望帮到你。

本文还有配套的精品资源，点击获取