PyTorch DataLoader高光谱数据加载与patch切分调优

简介：面向使用深度学习处理高光谱图像的研究者与工程实践者，这份代码资源围绕 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 / 65535，numpy 在某些版本下会先升到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 cube

2.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)), mode=mode) return padded

第二是标签对齐。gt是(H, W)的整型图，0 通常表示未标注或背景。切 patch 时索引要基于原始坐标，填充只影响 cube 的读取偏移，不能把gt也 pad 了再去取值，否则中心位置会整体偏移r个像素，训练出来的模型精度看着还行，混淆矩阵却怎么都对不上。这个偏移错误非常隐蔽，因为 loss 会正常下降。

3. 写一个能复用的 Dataset：patch 切分、标签映射与光谱增强

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_size=11, ignore_label=0, label_map=None, transform=None): # cube: (C, H, W) float32，已归一化 self.cube = cube self.patch_size = patch_size self.transform = transform self.r = patch_size // 2 # 标签重映射：把原始类别 id 压成 0..K-1，ignore_label 单独处理 self.label_map = label_map or {} gt_mapped = np.full_like(gt, -1, dtype=np.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_index=255)用，这是 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, y

patch 级要处理填充。注意这里我用的是“先 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_flip=0.5, p_band_drop=0.3, snr_db=30): 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 x

p_band_drop控制置零概率，width上限取波段数的 1/20，避免一次抹掉太多信息。噪声那几行里的snr_db是分贝值，30 dB 大约对应 3% 的噪声水平，可以根据数据质量在 20 到 40 之间调。增强只在训练集开，验证和测试一律关闭，这个开关放在 Dataset 外面控制，不要写死在类里。

4. DataLoader 参数怎么定：num_workers、采样器与 collate_fn

Dataset 写对之后，DataLoader 的参数决定了你的 GPU 是满负荷跑还是有一半时间在等数据。高光谱的数据量比自然图像大一个量级，一个11 × 11 × 200的 patch 展开是 24200 个浮点数，接近 100 KB，这个尺度下 num_workers 和内存的关系必须算清楚。

4.1 num_workers 与 persistent_workers：吞吐和内存的平衡

num_workers=0意味着数据加载在主进程里做，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_size=64, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True, prefetch_factor=4, drop_last=True, )

逐项说。pin_memory=True把 batch 放进锁页内存，配合训练循环里的x.cuda(non_blocking=True)能做异步拷贝，这一步在数据量大时能省下可观的等待。persistent_workers=True让 worker 进程在 epoch 之间存活，避免每个 epoch 重新 fork 一次，样本多、epoch 多的时候差别明显。prefetch_factor=4表示每个 worker 预取 4 个 batch，默认是 2，内存紧张就调回 2。drop_last=True是为了避免最后一个不满的 batch 影响 BatchNorm 的统计，样本量本来就少的话可以关掉。

注意：persistent_workers=True必须配合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, minlength=train_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( weights=sample_w, num_samples=len(sample_w), replacement=True, )

np.clip(counts, 1.0, None)防止某个类频次为 0 时除零。num_samples设成样本总数，意味着一个 epoch 的样本量和普通 shuffle 一致；如果想让小类多出现几次，可以把这个值调大到 1.5 倍，但要注意 epoch 会变长。replacement=True是有放回抽样，必须开，否则权重只是摆设。

用 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, dim=0) # (B, C, p, p) ys = torch.tensor(ys, dtype=torch.long) coords = torch.tensor(coords, dtype=torch.int32) return xs, ys, coords

torch.stack要求每个样本形状完全一致，所以 patch 级 Dataset 里必须保证边界处理统一，不能一部分 pad 一部分裁剪，否则 collate 直接报 size mismatch。ys显式指定torch.long，因为CrossEntropyLoss要求标签是 int64，从 numpy 转过来如果是 int32 会报错。

参数选型的经验值可以归纳成一张表，实际项目里按数据规模微调：

这张表不是标准答案，只是起步值。真正该看的是nvidia-smi里的 GPU 利用率和top里的 CPU 占用：GPU 忽高忽低说明数据供不上，加 worker；CPU 打满而 GPU 空闲，说明__getitem__里有重计算，该往缓存方向优化。

5. 避坑排查：五个高光谱 DataLoader 翻车现场

下面五条都是实际项目里踩过的，按“现象 → 原因 → 解决”写，遇到对应症状可以直接对号入座。

5.1 训练 loss 一直不降，模型像在学随机噪声

现象：数据能跑通，loss 在前几个 epoch 缓慢下降然后卡住，验证精度和随机猜差不多。

原因：维度顺序错了。(H, W, C)没有转成(C, H, W)，DataLoader拼出来是(B, H, W, C)，模型第一层Conv2d(in_channels=C)收到的却是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_index=255，Dataset 返回时把-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, dtype=np.float32): if os.path.exists(cache_path): return np.memmap(cache_path, dtype=dtype, mode="r", shape=shape) cache = np.memmap(cache_path, dtype=dtype, mode="w+", shape=shape) for i in range(len(dataset)): x, _ = dataset[i] cache[i] = x.numpy() cache.flush() return np.memmap(cache_path, dtype=dtype, mode="r", shape=shape)

shape要提前算好，是(样本数, C, p, p)，dtype用float32，别用float64，体积直接翻倍。mode="w+"创建，mode="r"只读映射，后者在多 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 震荡剧烈，检查学习率和归一化。下面这张自检表可以贴在实验记录里，每次换数据集过一遍：

6.3 我后来固定下来的两个习惯

第一个习惯是每换一份数据，先写一个十行的脚本，只做读取、打印形状、打印值域、打印标签分布这四件事，跑通了再动 Dataset。这十行代码看着不起眼，但它把“数据本身的问题”和“管道的问题”隔离开了，后面无论报什么错，都能确定不是数据格式的锅。

第二个习惯是给每个 Dataset 配一个_sanity_check方法，随机抽三个样本，把输入和标签打印出来，训练前手动跑一次。有一次标签整体偏移了一个像素，就是这个方法发现的——loss 正常下降，验证精度也不难看，只有把 patch 和标签可视化出来才能看出中心对不上。从那以后我每次改完切 patch 的逻辑，都强制走一遍这个小检查，再开始跑训练。

数据管道这部分没有太多花哨的技巧，难的是把每个边界情况想清楚。希望这些经验能帮你少走几个弯路。

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