PyTorch高光谱数据加载全指南：Dataset、DataLoader与踩坑优化

简介：面向 PyTorch 与高光谱图像处理学习者的代码示例包，演示如何借助 DataLoader 高效加载多通道高光谱数据，解决 HSI 通道多、数据量大的批处理与预处理难题。压缩包共 7 个文件，包含 3 个 Python 脚本、2 个 .mat 数据文件及 2 个编译缓存文件，整体仅 5.69MB。其中 .mat 文件为 Indian Pines 高光谱数据集，可直接复现实验；Python 脚本完整实现了自定义 Dataset 类，具体涵盖数据读取、归一化到 0—1、随机打乱、batch_size 调参、num_workers 多线程加载，以及 pin_memory 加速与自定义 collate_fn 合并多维数组等进阶用法，并配有训练入口，无需额外配置即可运行观察完整流程。已有 612 人学习。该资源不仅呈现 DataLoader 在 HSI 场景下的标准用法，还通过分模块的代码结构展示数据加载与模型训练的衔接，尤其适合刚接触遥感或高光谱深度学习的研究者，既可作为入门参考，也能根据自身数据修改 Dataset 部分，快速迁移到新的项目。此外，代码中还对数据预处理与加载性能进行了调优说明，帮助读者避免常见的踩坑问题。

1. 高光谱数据加载：DataLoader 是训练链路里最容易翻车的环节

在 PyTorch 里做高光谱影像分类，第一道坎往往不是模型设计，而是数据加载。高光谱数据和普通 RGB 图有本质差异：单张图有上百个波段，空间尺寸却可能只有 145×145，整景影像直接读进内存毫无压力；可一旦做 Patch 级像素分类，每个样本都要截取一块空间窗口并完整复制所有波段，DataLoader 就成了训练循环里最慢也最容易出错的黑匣子。下面按数据读取、Dataset 实现、DataLoader 参数、踩坑记录、二次改造的顺序展开，适合刚接触高光谱分类、想把 PyTorch 训练代码从「能跑」推进到「跑得稳」的 Python 开发者。

2. 技术选型与工程结构：从 ENVI 文件到可迭代的数据集

2.1 高光谱数据的三种存储形态与读取方式

高光谱数据在公开资源里最常见的形态有三种，它们决定了预处理脚本要依赖哪个库，也决定了 Dataset 的__init__第一行怎么写。

第一种是 ENVI 标准格式，一个.hdr文本头文件加一个.dat二进制数据体。.hdr里记录了 samples、lines、bands、data type、interleave 五个关键元信息，分别对应宽、高、波段数、存储类型和波段交错方式。用 spectral 库读取最省事：

import spectral.io.envi as envi img = envi.open("scene.hdr", "scene.dat").load() # 返回 numpy ndarray，统一转成 (lines, samples, bands) = (145, 145, 200)

这里有个容易忽略的点：interleave 字段。ENVI 支持 BIP、BIL、BSQ 三种物理存储顺序，有些数据商导出时默认 BSQ（波段优先），spectral 会在 load 时根据 header 自动把形状统一成(lines, samples, bands)，读出来不用自己操心轴序。但load()默认整景读入内存，建议先用open()拿到 header 看 shape，再决定是直接 load 还是走后面的 memmap 路线。另外 ENVI 数据 dtype 常见是 uint16，值域 0 到 65535，喂给网络前要转 float32，我习惯在预处理阶段统一转，而不是在__getitem__里每次转，后者会拖慢 worker 进程。

第二种是 GeoTIFF 多波段文件，rasterio 是主力。它read()返回的轴序是(bands, rows, cols)，和 PyTorch 的 CHW 正好一致，但和 spectral 返回的 HWC 不同，两个库混用转轴时容易翻车：

import rasterio import numpy as np with rasterio.open("scene.tif") as src: img = src.read() # (200, 145, 145) img = np.moveaxis(img, 0, -1) # 统一成 (145, 145, 200)

第三种是.mat或.npy的预处理数组。很多公开数据集已经帮你做过去水汽吸收波段、样本切分的工作，scipy.io.loadmat或np.load直接得到(H, W, C)的 float 数组，这是实际写训练代码时最省事的一条路。

三种形态的取舍可以看这张表：

给一个简单的选择标准：原始数据是学术公开数据集，用 spectral；数据带地理配准信息、要和 GIS 工具链打交道，用 rasterio；只要跑模型，一律先转 npy。

2.2 工程目录结构与依赖清单

这个加载工程完整落地后的目录结构大致是这样：

hsi_pytorch_loader/ ├── data/ │ ├── raw/ # 原始 .hdr/.dat 或 .tif，只读不改 │ ├── processed/ # 预处理后的 img.npy、label.npy、band_stats.npy │ └── split/ # 训练/验证/测试的样本索引 ├── datasets/ │ ├── __init__.py │ ├── hsi_dataset.py # 自定义 Dataset 核心实现 │ └── transforms.py # 归一化、波段选择、在线增强 ├── configs/ │ └── train.yaml # 路径、patch_size、batch_size 等 ├── train.py # 训练入口，DataLoader 装配与循环 ├── preprocess.py # raw -> processed 的转换脚本 └── requirements.txt

依赖清单方面，核心只有 torch、numpy、spectral（或 rasterio 二选一），pandas、scikit-learn 用在样本划分和类别分布统计上，不是训练必需。requirements.txt 里建议固定 torch 与 numpy 的大版本，因为 DataLoader 的多进程行为和 numpy 版本有关，后文会回到这个话题。

我拆这类工程时坚持一个约定：原始格式读取只发生在 preprocess.py，训练代码只跟 processed 目录下的 npy 打交道。preprocess.py 负责去水汽吸收波段、归一化统计、划分样本索引三件事，产出img.npy（H×W×C，float32）和label.npy（H×W，int64）。这样换数据集时只改 preprocess.py，Dataset 一行都不用动。配置集中在 train.yaml：

data: img_path: data/processed/img.npy label_path: data/processed/label.npy band_stats: data/processed/band_stats.npy patch_size: 13 loader: batch_size: 256 num_workers: 4 pin_memory: true

训练时用yaml.safe_load读进来传给 Dataset 和 DataLoader。我做实验的习惯是每个实验存一份 yaml 副本，和训练日志一起归档，这样回头能精确复现当时的配置，不用靠记忆猜参数。

3. 核心代码拆解：Dataset 与 DataLoader 完整装配过程

3.1 自定义 Dataset：初始化与样本索引

Dataset 要回答的核心问题是「一个样本是什么」。高光谱像素分类任务里，一个样本 = 以某个有效像素为中心的空间 Patch（比如 13×13 窗口）× 全部波段，标签是中心像素的地物类别。这里有一个关键选择：返回整图还是返回 Patch。整图路线适合分割任务，一次__getitem__返回(200, 145, 145)，batch_size 只能压到个位数；Patch 路线适合分类任务，样本数等于有效像素数，batch 可以随便设。绝大多数高光谱分类论文走 Patch 路线，核心代码是下面这段：

class HSIPatchDataset(Dataset): def __init__(self, img, label, patch_size=13, valid_classes=None): # img: (H, W, C) float32；label: (H, W) int64，0 表示背景 self.img = img self.label = label self.patch_size = patch_size self.half = patch_size // 2 self.rows, self.cols, self.bands = img.shape mask = label > 0 if valid_classes is not None: mask = mask & np.isin(label, valid_classes) self.samples = np.argwhere(mask).astype(np.int32) self.class_names = np.unique(label[label > 0]) def __len__(self): return len(self.samples)

这段代码做了三件事：把背景像素（label=0）从候选样本里剔除；支持只保留指定类别做实验；用argwhere生成样本坐标表。__len__返回的是有效样本数而不是影像尺寸，和 CV 里按文件列表建 Dataset 的思路一致。

valid_classes参数值得多说一句。高光谱数据集的类别分布极不均衡，植被类可能占三成像素，某个稀有地物只有几十个像素。做「只训练某几类、看可分性」的对比实验时，这个参数比在预处理阶段重新裁剪数据灵活得多，改一行配置就能切换实验组合。patch_size 的选择也有讲究，13×13 是个常用起点，窗口太小感受野不足，太大则样本高度重叠、训练变慢，一般从 11 到 27 之间调。

3.2getitem实现：中心像素采样与 Patch 提取

__getitem__是加载性能的分水岭，也是容易写错的地方。朴素写法是直接按坐标切片：

def __getitem__(self, idx): r, c = self.samples[idx] patch = self.img[r - self.half: r + self.half + 1, c - self.half: c + self.half + 1, :] # 边界像素窗口不完整，零填充到 patch_size ph, pw = patch.shape[0], patch.shape[1] if ph != self.patch_size or pw != self.patch_size: padded = np.zeros((self.patch_size, self.patch_size, self.bands), dtype=self.img.dtype) pr = min(ph, self.patch_size) pw_ = min(pw, self.patch_size) padded[:pr, :pw_] = patch[:pr, :pw_] patch = padded patch = np.ascontiguousarray(patch.transpose(2, 0, 1)) # HWC -> CHW x = torch.from_numpy(patch).float() y = torch.tensor(self.label[r, c], dtype=torch.long) return x, y

单看逻辑不复杂，但高光谱的 C 维度是 200，一个 13×13×200 的 Patch 有 33800 个 float，一个 batch 256 个样本就是 860 多万个数值的搬运，数据加载时间很容易超过 GPU 计算时间，这是高光谱训练比普通 CV 慢的根源之一。想确认瓶颈时，在训练循环里对 DataLoader 的迭代计时，如果load耗时占比超过三成，就该调 num_workers 或走后面的缓存改造。

这段代码有两个细节值得注意。第一，边界零填充不是可选项而是必需项，否则边缘像素返回的 Patch 形状不齐，默认 collate 组装 batch 时直接报错，这个问题避坑章节还会再讲。第二，transpose 之后必须接ascontiguousarray，因为 numpy 的 transpose 返回的是视图、内存不连续，torch.from_numpy虽然能用，但后续 GPU 拷贝会触发隐式的内存整理，白白多一次复制。

如果整景数据能放内存，我一般会在__init__里预先做一次img.transpose(2,0,1)并转 float32，整理成 CHW 连续数组，这样__getitem__里只需要做空间维切片，省掉每次转轴和拷贝的开销，代价是内存里多一份副本，内存紧张时不要这样干。

3.3 DataLoader 装配：batch_size、num_workers 与常见参数

Dataset 写完后，装配 DataLoader 的代码很短，但每个参数在高光谱场景下都有讲究：

train_loader = DataLoader( train_dataset, batch_size=256, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, persistent_workers=True, prefetch_factor=2, )

参数选择上有几个经验值。batch_size 受 Patch 尺寸影响：13×13×200 用 256 很稳；Patch 涨到 27×27 或者模型换成 3D-CNN，batch 要降到 64 或 32，否则显存和训练时间都吃不消。num_workers 从 4 起步，看训练日志里数据加载耗时占比再决定加不加，一般到 8 就到头了，超过物理核数反而变慢，这个只能实测不能拍脑袋。

pin_memory=True对 GPU 训练几乎总是划算，它让 host 内存页锁定，Host-To-Device 拷贝更快。persistent_workers=True让 worker 跨 epoch 常驻，避免每轮重新 fork 进程、重新加载 numpy 数据的启动开销，但它会放大 CUDA fork 相关的问题，避坑章节会展开。prefetch_factor默认 2，瓶颈在磁盘 IO 时可以调到 4 或 8。

高光谱类别不均衡是个绕不开的问题，光靠 shuffle 不够，建议用 WeightedRandomSampler 替代 DataLoader 自带的 shuffle，让稀有类别的采样概率提上来：

class_weight = 1.0 / np.bincount(train_dataset.labels) sample_weight = class_weight[train_dataset.labels] sampler = WeightedRandomSampler(sample_weight, num_samples=len(sample_weight), replacement=True) loader = DataLoader(train_dataset, batch_size=256, sampler=sampler, num_workers=4, pin_memory=True)

提示：用了 sampler 之后，DataLoader 的 shuffle 参数必须保持默认 False，两者互斥，这是 PyTorch 里一个容易忽略的约束。

验证集的 DataLoader 则保持简单，关掉 shuffle 保证评估顺序稳定：

val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False, num_workers=4, pin_memory=True)

4. 踩坑与排查：高光谱加载的高频问题与修复记录

4.1 现象：训练还没开始，内存直接爆掉

某高校实验室的同学第一次跑 3D-CNN，用一个 512×512 像素、204 波段的城区场景，预处理时np.save没任何问题，一进 DataLoader 就 OOM。后来发现是__init__里做了一次img.transpose(2,0,1)生成全图 CHW 副本，同时保留了原始 HWC 数组，两个 float64 数组各占 2GB 多，加上 4 个 worker 进程 fork 时的内存开销，直接翻了好几倍。

原因很简单：高光谱单景数据看着不大，但 float64 的 HWC 和 CHW 双份副本会把内存占用翻倍甚至翻三倍。解决方法是统一 float32，__init__里只保留一份数据副本。如果场景真的超过单机内存，改用 numpy 的 memmap 按需映射：

self.img = np.load("img.npy", mmap_mode="r") # 磁盘映射，不占物理内存 patch = self.img[r1:r2, c1:c2, :].astype(np.float32) # 只读需要的窗口

注意：mmap 是随机磁盘读，频繁访问时比内存慢一个数量级，只建议在内存确实放不下的场景用，而且 num_workers 要适当降下来，避免多进程同时打磁盘。

4.2 现象：num_workers>0 时进程直接崩，报 CUDA 初始化错误

报错信息形如Cannot re-initialize CUDA in forked subprocess。原因是 DataLoader 的多进程 worker 由 fork 创建，而 CUDA context 无法在 fork 出的子进程里重新初始化。常见处理分两步：第一，模型和数据搬到 GPU 的代码只放在主进程，Dataset 内只做 CPU 侧 numpy 操作；第二，检查 Dataset 或自定义 collate 里有没有提前生成 CUDA 张量，有就去掉。

还有一个隐蔽场景：训练脚本 import 了某个可视化或日志库，这个库在 import 阶段就初始化了 CUDA，导致 worker fork 时已经带上 CUDA context。排查办法是把 import 逐个注释掉，定位到具体库之后改成懒加载，只有主进程需要时才 import。

4.3 现象：collate 阶段报 stack expects each tensor to be equal size

这个错误几乎都出在边界 Patch 上。影像边缘像素的窗口天然不满 patch_size，如果__getitem__没做零填充，返回的 tensor 就是 11×13×200 或 9×9×200，同一个 batch 里形状不齐，默认 collate 直接炸。修复就是 3.2 节里的零填充分支。

补充一个经验：这里不要用反射填充或边缘复制代替零填充，虽然视觉上更自然，但会引入边缘之外并不存在的像素值，干扰中心像素的分类。高光谱分类比 RGB 分类对数值敏感得多，一个波段值错位，模型可能把同一个像素判成另一个地物。

4.4 现象：验证集精度虚高，换个数据集就崩

用全图统计量做归一化是最隐蔽的数据泄漏。常见写法是在 Dataset 里对整景影像算 min/max 或 mean/std，再对训练、验证的 Patch 一起归一化。问题是验证 Patch 已经「见过」训练 Patch 的统计信息，精度虚高，换一个分布完全不同的场景就崩。正确做法是只从训练像素里统计归一化参数，存成文件，验证和测试时复用同一组参数：

# 预处理阶段，只基于训练集掩码统计 band_mean = img[mask_train].reshape(-1, bands).mean(axis=0) band_std = img[mask_train].reshape(-1, bands).std(axis=0) np.save("data/processed/band_stats.npy", np.stack([band_mean, band_std]))

Dataset 里加载 band_stats.npy，统一做(img - mean) / std即可。顺手给每个波段存一份统计量，后面做归一化方式的对比实验会省很多时间。

4.5 现象：同一份数据，两次训练精度对不上

再次运行训练时精度波动超过 1 个百分点，先怀疑 shuffle 随机性而不是模型。高光谱数据划分如果没有固定种子，训练和验证像素每次都不一样，模型看到的数据分布变了，精度自然对不上。解决方法是把数据集划分逻辑固定下来：设置torch.manual_seed和np.random.seed，把划分出的样本索引存到 data/split 目录，每次训练加载同一个索引文件，保证可复现。

更进一步的做法是让 Dataset 里__getitem__的随机性也受控，尤其是做了随机裁剪、随机波段扰动等在线增强时，否则即使索引一致，增强路径不同也会带来精度抖动。论文复现和对比实验时，这一步很关键。

5. 二次开发：从单机训练到多模态与分布式加载改造

5.1 改造一：波段子集采样，支撑「少波段」对比实验

高光谱论文里经常要做「全波段 vs 选定波段」的可分性对比。给 Dataset 加一个band_indices参数即可，完全不用改数据文件：

def __init__(self, img, label, band_indices=None, patch_size=13): self.band_indices = band_indices def __getitem__(self, idx): ... if self.band_indices is not None: patch = patch[:, :, self.band_indices]

把波段列表放进 configs/train.yaml，每次实验改配置而不是改代码。我习惯把波段选择放在 Dataset 内部而不是预处理脚本里，因为实验阶段组合换得特别勤，全量重跑预处理不现实。配合 3.1 的 valid_classes，就能组合出「指定类别 + 指定波段」的任意实验矩阵。

5.2 改造二：缓存机制与分布式训练兼容

单机加载跑通后，下一步值得做两个改造。一是加缓存层：把高频访问的 Patch 用 LRU 缓存放在 Dataset 内部。注意 num_workers>0 时每个 worker 有独立缓存，命中率会下降，所以缓存更适合 worker=1 或在线推理场景；如果内存充足，更简单粗暴的方式是在预处理阶段把全部 Patch 落盘成顺序 npy，训练时按索引顺序读，用空间换时间。

二是分布式训练兼容。PyTorch DDP 下每个进程持有各自的 Dataset 实例，必须保证每个 epoch 的 shuffle 在所有 rank 上一致，否则 loss 曲线异常抖动。常见做法是数据划分脚本生成全量样本索引文件，DDP 各进程用 DistributedSampler 对索引做切分，Dataset 本身保持无状态（不持有随机种子、不持有划分状态），这样数据和模型可以独立扩展。

高光谱加载这个模块，我前后重构过三版。第一版直接读原始 ENVI 格式，能跑但慢；第二版统一转 npy、把边界填充写对，公开数据集全部跑通；第三版加波段子集和缓存之后，才敢说它能交给别人维护和二次开发。从那以后，我每拿到一个新数据集都强制先走一遍「预处理输出 npy → Dataset 只做索引与切片 → 统计量只基于训练集」这条链路，整套流程最快十分钟走完，但省掉的调试时间是以天来计的。希望帮到你。

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