简介面向使用PyTorch处理高光谱图像HSI的开发者提供一套基于DataLoader的高效数据加载实现。资源围绕高光谱数据通道多、内存占用大的特点完整展示了Dataset子类封装、数据读取支持.mat等常见格式、归一化预处理、多线程批处理加载以及自定义collate_fn、shuffle随机采样等关键操作可显著提升数据加载效率并保障训练稳定性。压缩包为7z格式共7个文件包含3个Python脚本实现dataloader、工具函数与训练流程、2个编译后的pyc文件及2个MAT格式高光谱样本数据如IndianPines整体大小5.69MB。工程目录结构清晰utils与DataSet模块划分明确便于直接复用或二次开发。该资源已有612人学习下载。通过阅读代码读者可快速掌握高光谱数据从读取到批处理的全流程实现并能根据自身硬件灵活调节batch_size、num_workers与pin_memory等参数将示例迁移至自备数据集。1. 用 PyTorch Dataloader 加载高光谱为什么常规写法必然翻车用 Python 和 PyTorch 加载高光谱数据时最直觉的写法是把整幅光谱立方体torch.from_numpy(...)再扔给 DataLoader这也是我最初踩过的坑。高光谱图像和普通 RGB 图完全是两类数据一个波段通道数少则十几、多则数百单幅影像动辄几百 MB。DataLoader 的高效与否直接决定训练能不能跑起来、跑多快。这篇笔记写给用 PyTorch 做高光谱分类或分割的从业者内容包括数据格式处理、Dataset 子类设计、DataLoader 参数调优和一份真实的避坑清单照着复现就能绕开「爆显存、段错误、光谱维错位」这三个最典型的翻车点。2. 先处理数据格式HDF5/ENVI 的光谱维、空间维与读取库选择2.1 高光谱与普通图像的数据组织差异为什么不能直接读图高光谱数据的核心是「一摞图叠在一起」空间上是行、列光谱上是波段。常见存储格式有两种拿到手先别急着写 Dataset先分清格式再说。一种是 HDF5.h5/.h5py内部按(样本数, 波段数, 高度, 宽度)组织维度顺序大概率是(N, C, H, W)这正好符合 PyTorch 的NCHW输入习惯。但 HDF5 文件里数据不一定全在一个 dataset 里有些项目会把主数据和标签拆成两个 key还有可能把波段维放在最后一位打开文件后必须打印shape确认。另一种是 ENVI 格式.hdr.dat/.bil在遥感领域非常常见。ENVI 有三种像元存储布局BSQBand Sequential波段顺序、BILBand Interleaved by Line行交叉、BIPBand Interleaved by Pixel像元交叉。BSQ 的shape通常接近(bands, rows, cols)BIP 的shape更接近(rows, cols, bands)。如果直接按(bands, rows, cols)去读 BIP 数据等于把每个像元的光谱向量当成了横切面后面模型学到的特征就是乱的。我一般先用 GDAL 或spectral库读 ENVI 文件读出后用np.moveaxis把波段维挪到第 0 维再交给 Dataset。读取库的选择上HDF5 用h5pyENVI 用osgeo.gdal或spectral。别贪图省事用scipy.io.loadmat去读h5py之外的格式很多高光谱公开集虽然后缀是.mat但内部是 HDF5 结构loadmat直接报错用h5py反而能读。2.2 Dataset 子类把读取、波段选择、缓存放一起构建 Dataset 时最忌讳在__init__里一次性把全部数据读进内存。一幅 600×600×224 的 float32 立方体约占 322 MB600×600×224×4 字节 / 1024² ≈ 322 MB几十个样本就把内存吃光了这还没算后续 DataLoader 多进程复制开销。我一般把读取动作放在__getitem__里并加一个「读后缓存」开关。import h5py import numpy as np import torch from torch.utils.data import Dataset class HSIDataset(Dataset): HDF5 高光谱数据集读取后保持 (bands, rows, cols) 返回 def __init__(self, h5_path, keyhsi, band_sliceNone, use_cacheTrue): self.h5_path h5_path self.key key self.use_cache use_cache # band_slice 只支持 slice 对象例如 slice(10, 200, 2) # 选连续区间用 slice不要用 list后面避坑章节会展开 self.band_slice band_slice with h5py.File(h5_path, r) as f: hsi f[key] self.num_samples hsi.shape[0] self.band_total hsi.shape[1] self.h hsi.shape[2] self.w hsi.shape[3] # 用 dict 做样本级缓存内存吃紧时把它关掉 self.cache {} def __len__(self): return self.num_samples def __getitem__(self, idx): # 命中缓存直接返回避免重复读 HDF5 if self.use_cache and idx in self.cache: cube self.cache[idx] else: # 这里才真正打开文件。多进程 DataLoader worker 会 fork # 在 __init__ 里保留文件句柄跨进程共享会导致 HDF5 崩溃 with h5py.File(self.h5_path, r) as f: cube f[self.key][idx] # 波段维切片优先用 slice if self.band_slice is not None: cube cube[self.band_slice] # 转成连续内存的 float32再进 Tensor cube np.ascontiguousarray(cube, dtypenp.float32) if self.use_cache: self.cache[idx] cube return torch.as_tensor(cube)参数含义keyHDF5 内 dataset 的路径。有的文件里叫data、cube先用f.keys()确认。band_slice只接受slice例如选 10 到 200 波段且每隔 2 步采样。不要传列表列表索引会触发 fancy indexing产生额外拷贝后面第四章验证过。use_cache样本数不多就开样本多或内存吃紧就关。缓存住的是处理好的 float32不是原始 uint16能省一半内存。逻辑上有两点容易被忽略。其一__init__里只读取 shape不读整块数据这样 Dataset 初始化极快其二__getitem__里每次用with打开文件虽然看起来有点浪费但实际上 HDF5 支持并发只读多进程 worker 每次打开自己独立的句柄比共享父进程句柄稳定得多。如果文件特别大且单句柄打开耗时长可以换成在__init__中只保留路径、在__getitem__里懒加载一个全局句柄但那个方案在多进程下要额外加锁我自己的项目里不推荐。2.3 先跑一遍单进程时序验证Dataset 写完别直接训。先用batch_size1、shuffleFalse跑一遍单进程加载确认返回张量的shape和dtype顺带测一次完整遍历耗时。这一步能把「模型里报维度错误」和「DataLoader 加载过慢」两类问题在训练前干掉。from torch.utils.data import DataLoader import time ds HSIDataset(hsi.h5, band_sliceslice(0, 200, 2)) loader DataLoader(ds, batch_size1, shuffleFalse) sample next(iter(loader)) print(sample shape:, sample.shape) print(sample dtype:, sample.dtype) start time.time() for _ in loader: pass print(f单进程遍历耗时: {time.time() - start:.3f} s)预期输出里sample.shape应该是torch.Size([1, 100, rows, cols])如果出现 5 维或波段维在最后说明数据原始布局有问题回到上一节用np.moveaxis纠正而不是硬着头皮往下写模型。这一步跑完数据加载部分心里就有底了之后调num_workers才有对照基准。3. DataLoader 参数调参num_workers、pin_memory 与预取策略3.1 DataLoader 参数表中高光谱环境该动哪些DataLoader 大多数参数在普通图像上可以照抄默认值但高光谱数据体量大、通道数多参数组合会直接影响显存和吞吐。先给一份速查表后面逐条展开。参数默认值高光谱场景下的推荐取向batch_size1先按模型输入上限算显存高光谱常设 8~32别拍脑袋shuffleFalse训练开 True验证开 False做滑窗切块时可用自定义 Samplernum_workers0从 0 起步验证再按 CPU 核数的一半起步往上加pin_memoryFalseGPU 训练推荐 True但会多占一部分主机内存drop_lastFalse遇到最后一批 shape 不齐、BN 层报错时设 Trueprefetch_factor2数据读取慢就提高到 4 或更高但要盯着内存占用persistent_workersFalse高频重启 worker 的 epoch 循环推荐 Truenum_workers是高光谱加载最容易“玄学”的参数。它不是越大越好因为每个 worker 会各复制一份 Python 环境还要各自打开 HDF5 文件句柄worker 数量超过物理核心时线程切换成本甚至比读取成本还高。我的经验是先从 0 开始确认数据管道正确再跳到os.cpu_count() // 2观察 GPU 利用率是否上升如果上升不明显优先查磁盘 IO而不是继续加 worker。pin_memoryTrue是对 GPU 训练友好的选项PyTorch 会在主机内存中锁页再把数据异步拷到显存省去一次分页拷贝。但锁页内存是不能被换出的锁太多会导致系统内存不足甚至触发 OOM。高光谱单样本本来就大要对batch_size × 单样本字节数 × prefetch_factor有个预估别把pin_memory当成无脑加速开关。3.2 用“先 0 后多”的顺序排除加载黑匣子多进程 DataLoader 出了诡异问题很难直接定位因为报错往往不在数据管道本身而是在训练循环的中后期。我习惯用“先 0 后多”的验证顺序像剥洋葱一样定位瓶颈。先跑num_workers0让主进程单线程加载数据这样任何异常都会直接暴露完整的堆栈错误信息也会指向__getitem__里的具体行。确认无误后再跑一组对比实验分别用num_workers2, 4, 8加载同一个 epoch记录耗时和峰值内存。注意这时候要保证数据已经预热过否则第一次 HDF5 读取会把磁盘缓存时间算进去干扰判断。import os import time from torch.utils.data import DataLoader def bench_loader(num_workers): ds HSIDataset(hsi.h5, band_sliceslice(20, 200, 2)) loader DataLoader( ds, batch_size16, shuffleTrue, num_workersnum_workers, prefetch_factor2, pin_memoryTrue, persistent_workersnum_workers 0, ) start time.time() for _ in loader: pass print(fnum_workers{num_workers} 耗时 {time.time() - start:.3f} s) for n in [0, 2, 4, 8]: bench_loader(n)逻辑说明每次循环重新创建 DataLoader是因为要干净地比较不同 worker 数下的启动成本。persistent_workersTrue会让 worker 在 epoch 间常驻而不是每个 epoch 结束就销毁重启这一步在 epoch 数多、单 epoch 耗时短时收益明显但 worker 多了以后常驻内存也高。对比结果一般呈 U 形worker 太少IO 和预处理算不过来worker 太多内存带宽先见顶。找到拐点后固定下来再回头调prefetch_factor。3.3 可复用的 build_loader 辅助函数把上面验证后的参数固化成函数方便不同项目复用。需要注意prefetch_factor在 worker 为 0 时不能设置否则会直接报参数错误所以要先做判断。def build_loader(h5_path, batch_size16, num_workers4, shuffleTrue, is_trainTrue, band_sliceNone, use_cacheTrue): ds HSIDataset( h5_path, band_sliceband_slice, use_cacheuse_cache, ) loader_kwargs dict( batch_sizebatch_size, shuffleshuffle, num_workersnum_workers, pin_memoryTrue, persistent_workersnum_workers 0, ) # num_workers0 时不能设置 prefetch_factorPyTorch 会直接拒绝 if num_workers 0: loader_kwargs[prefetch_factor] 2 loader DataLoader(ds, **loader_kwargs) return loader这个函数把pin_memory固定写成 True因为高光谱上 GPU 训练的收益普遍大于内存占用代价如果跑纯 CPU 推理手动改成 False 即可。is_train参数在这里只留了一个开关位后续可以接Sampler控制样本均衡目前没实际使用。重点是num_workers 0时prefetch_factor才生效这个坑在 PyTorch 低版本上尤其隐蔽启动时不会报错但数据读取一直单线程。4. 高光谱 DataLoader 避坑清单从段错误、显存溢出到光谱维度错位4.1 h5py 句柄在多进程 worker 下段错误现象num_workers从 0 改成 4 后训练刚开始或中途直接报Segmentation fault有时连带报HDF5 library has been compiled without thread safety。原因在__init__里打开 HDF5 文件并保存句柄然后交给多进程 DataLoader。Linux 的fork会把父进程的 HDF5 文件描述符和内部状态复制给子进程多个子进程同时操作同一个句柄HDF5 底层库线程不安全直接崩溃。解决不要在__init__里持有 HDF5 句柄改为在__getitem__内部用with h5py.File(...)打开读取。这样每个 epoch 里每个样本都会独立打开文件看起来重复开销大但 HDF5 打开文件本身是轻量操作比崩溃后排查成本低得多。另外h5py打开文件时显式传入moder不要用默认写权限。4.2 整幅 float32 立方体直接灌进显存现象模型和数据管道都跑通训练第一个 batch 就报CUDA out of memory把batch_size降到 4 依然报错。原因高光谱输入是 4 维张量(batch, bands, H, W)中间层特征图又会放大通道数。我遇到过 224 波段的立方体即使空间分辨率只有 100×100一个 batch 的输入就有 224×100×100×4 字节 ≈ 9 MB加上网络中间特征显存占用远超预期。解决优先压缩波段band_sliceslice(10, 120, 2)把 224 波段砍到 55精度损失可控。其次是空间降采样在__getitem__里对空间维做最近邻或双线性缩小而不是在进入模型后才发现尺寸不对。最后才是调batch_size。如果数据集本身波段数固定可以把预处理统一到__getitem__里训练循环不碰原始维度的计算。4.3 光谱维错位BIP/BSQ 布局与 moveaxis现象loss 下降正常但验证集准确率一直偏低或者训练无崩溃但 loss 波动很大。原因高光谱数据原始布局是(height, width, bands)但是读取出来直接当作(bands, height, width)用了。尤其 BIP 格式每个像元的光谱向量是连续存储的布局更像(height, width, bands)。如果 Dataset 里不做转置模型的通道维就莫名其妙变成了空间维训练再久也学不到光谱特征。解决打印一次cube.shape是第一步。若确认波段维在最后统一用np.moveaxis(cube, -1, 0)把最后一个维度挪到第一个位置。注意做np.ascontiguousarray因为moveaxis得到的可能是非连续视图PyTorch 转 Tensor 时虽然能转但后续view/permute容易出幺蛾子先转连续内存最稳妥。这一步要放在 band_slice 之前还是之后先moveaxis再切片因为切片索引依赖波段维的位置。4.4 band 列表索引不省时切片与 fancy indexing现象用cube[[10, 20, 30, ..., 200]]选取 20 个波段加载耗时反而比全量 224 波段还慢。原因列表索引属于 numpy 的 fancy indexing会触发逐索引拷贝产生一段不连续的内存访问尤其在高光谱大矩阵上开销被放大。slice(10, 201, 10)是视图语义numpy 直接按跨度切少了中间层拷贝。解决能用slice就不用列表。只有当波段是任意非等距挑选比如只选 450nm、550nm、650nm 特征波段才需要np.take并且实验前做一次耗时对比。经验数据是 200 波段里挑 10 个等距波段用slice比列表索引快 3 到 5 倍在 DataLoader worker 里被放大成训练吞吐差异。另外如果 selection 是固定不变的把它放进band_slice参数里Dataset 内部只执行一次不要在每次__getitem__里重复构造索引数组。5. 进阶技巧把整幅高光谱切成任意 Patch 的高效数据流5.1 tile 式 Dataset 最短实现高光谱场景里训练样本经常不是一张张小图而是一整幅大立方体。直接当成一个样本喂模型空间信息保留太多显存也扛不住。常见做法是把整幅图切成固定大小的 patch再喂给 DataLoader。我把 tile 逻辑放进 Dataset让 DataLoader 只负责调度。import h5py import numpy as np import torch from torch.utils.data import Dataset class TileHSIDataset(Dataset): 从单幅高光谱立方体中随机切 patch返回 (1, bands, ps, ps) def __init__(self, h5_path, patch_size64, num_tiles1000, seed42): self.h5_path h5_path self.ps patch_size with h5py.File(h5_path, r) as f: hsi f[hsi] self.bands hsi.shape[0] self.h, self.w hsi.shape[1], hsi.shape[2] # 预先采样 num_tiles 个左上角坐标避免随机种子在多进程下不一致 rng np.random.default_rng(seed) self.ys rng.integers(0, self.h - self.ps, sizenum_tiles) self.xs rng.integers(0, self.w - self.ps, sizenum_tiles) def __len__(self): return len(self.xs) def __getitem__(self, idx): x, y self.xs[idx], self.ys[idx] with h5py.File(self.h5_path, r) as f: hsi f[hsi] # 只读需要的窗口HDF5 底层按 chunk 读取避免全量载入 patch hsi[:, y:yself.ps, x:xself.ps] patch np.ascontiguousarray(patch, dtypenp.float32) return torch.as_tensor(patch).unsqueeze(0) # 补一个样本维关键参数patch_size是窗口尺寸不要求能被整图长宽整除最后一行/列采样不到的坐标在随机采样时自然避开num_tiles决定一个 epoch 里的 patch 数量按训练步数需求反推seed必须在__init__里固定否则多进程下每个 worker 各持一份随机数epoch 内 patch 可能出现重复或覆盖不均。这里没有直接用random.randint(0, H - ps)而是提前把坐标数组算好。原因是多进程 DataLoader 里每个 worker 是独立进程共享不了 Python 的random状态不同 worker 可能采到几乎同一批 patch导致一个 batch 内重复样本过多。提前用 numpydefault_rng固定坐标再用索引取坐标从根源上保证了可复现性。5.2 加载后强制验证的不变量每次训练跑起来之前我不管 Model 怎么改都强制检查两条不变量。第一返回张量shape[1]必须等于波段数第二torch.unique前几个 patch 的值域必须在预处理后的预期范围内不能出现 NaN 或全零样本。loader build_loader(large.h5, batch_size16, num_workers4) sample_batch next(iter(loader)) assert sample_batch.shape torch.Size([16, 224, 64, 64]), sample_batch.shape assert torch.isfinite(sample_batch).all(), 数据里出现 NaN/Inf first_patch sample_batch[0, 0] print(min/max:, first_patch.min().item(), first_patch.max().item())这两行 assert 看起来多余但高光谱数据预处理链条长从 HDF5 读取到归一化再到增广任何一步出错都不会立刻让训练崩溃只会让 loss 卡在异常区间。从那以后我每次新拿到一块高光谱数据都强制走一遍单进程加载、shape 断言、值域检查再上num_workers多进程训练这套流程帮我挡掉了至少三次数据错位引发的返工希望帮到你。本文还有配套的精品资源点击获取