简介围绕PyTorch DataLoader高效加载高光谱数据这一核心问题这份7z压缩包提供了可直接运行的完整工程。高光谱图像包含成百上千个光谱通道数据量远大于普通RGB图像加载策略直接影响训练效率与内存占用。包内共有7个文件大小约5.69MB包含3个Python脚本、2个.mat高光谱示例数据及2个pyc缓存文件脚本按功能拆分覆盖数据集读取、预处理、批处理、训练入口等模块数据文件采用Indian Pines经典高光谱场景便于读者直接运行验证。目前已有612人学习下载是入门高光谱与PyTorch结合的实用参考。资源重点展示了如何自定义Dataset类实现__len__与__getitem__如何归一化光谱值、设置batch_size与num_workers、启用shuffle和pin_memory并针对多维数组提供自定义collate_fn的写法。这些细节能有效解决数据量巨大、内存占用高、加载缓慢等常见问题帮助开发者构建稳健的数据流水线。无论用于课程实验、论文复现还是实际项目改造这份代码都能提供清晰的起点避免重复踩坑。1. 高光谱加载为什么值得单独写别让DataLoader等你读文件高光谱数据用 PyTorch 加载第一反应通常是torch.from_numpy一把梭。实际一跑就发现训练卡在数据读取上GPU 占用率上不去显存还被一次 batch 塞爆。高光谱和 RGB 图最大的区别在于波段维度一张图往往有上百个波段数据立方体的体积是普通图像的几十倍直接套用 ImageNet 那套加载流程连数据类型和通道顺序都撑不住。这篇笔记完整拆解用 PyTorch 的 DataLoader 加载高光谱数据的全部步骤包括读取格式、通道顺序、Patch 采样、归一化、踩坑点和提速手段适合正在用 Python 做高光谱分类或者分割的开发者直接复现。2. 读取与张量约定把高光谱立方体理顺成 PyTorch 认识的输入高光谱数据最常见的落盘格式不是 PNG而是 HDF5 和 hdr/dat 配套的遥感格式。前者适合存大数组后者是传统的二进制裸数据加一个文本头。不管哪种最终都要落到同一个目标得到一个(H, W, B)的 numpy 数组其中 H 是行数、W 是列数、B 是波段数。之后要做的所有事情都围绕这个三维立方体展开。2.1 HDF5 与 hdr/dat 读取两种最常见的存储方式HDF5 在 Python 里直接用 h5py 读代码很短但很多人会卡在 key 的取值上。同一个 h5 文件里可能同时存了影像、标签、元数据读取前先打印 keys再决定读哪个数据集。另外f[key][()]是一次性读入内存文件很大的时候要小心内存压力。import h5py import numpy as np def load_h5(path: str, key: str data) - np.ndarray: with h5py.File(path, r) as f: print(available keys:, list(f.keys())) arr f[key][()] return arr # shape: (H, W, B)逻辑说明[()]是 h5py 把整个数据集拉进内存的写法等价于[:]但语义更明确。如果文件是分块存储的且单张影像超过几个 GB不建议用这个函数一次读全图而是保留f[key]的 Dataset 对象循环里按坐标切片读。参数说明key没有固定标准常见的有 data、cube、hsi最稳妥的办法是先打印 keys 再传参别硬编码。hdr/dat 格式更能暴露问题。hdr 文件本质是文本里面记录了行数、列数、波段数、数据类型和存储顺序 interleave。interleave 有三种BSQ波段按顺序整块存储、BIL行内波段连续、BIP像素内波段连续写法稍微错一个整个立方体就是花的。import numpy as np DTYPE_MAP { 1: np.uint8, 2: np.int16, 3: np.int32, 4: np.float32, 5: np.float64, 12: np.uint16, } def read_remote_cube(hdr_path: str, dat_path: str None) - np.ndarray: meta {} with open(hdr_path, r, errorsignore) as f: for line in f: if in line: k, v line.split(, 1) meta[k.strip().lower()] v.strip().strip(\) lines int(meta[lines]) samples int(meta[samples]) bands int(meta[bands]) dtype DTYPE_MAP[meta[data type]] interleave meta.get(interleave, bsq).lower() offset int(meta.get(header offset, 0)) if dat_path is None: dat_path hdr_path.rsplit(., 1)[0] .dat raw np.fromfile(dat_path, dtypedtype, offsetoffset) if interleave bsq: cube raw.reshape(bands, lines, samples).transpose(1, 2, 0) elif interleave bil: cube raw.reshape(lines, bands, samples).transpose(0, 2, 1) elif interleave bip: cube raw.reshape(lines, samples, bands) else: raise ValueError(funknown interleave: {interleave}) return cube # (H, W, B)逻辑说明BSQ 存储时文件中前lines * samples个数是第一个波段所以先 reshape 成(bands, lines, samples)再用 transpose 换成(lines, samples, bands)。BIL 则是每个行内先存所有波段的一行reshape 成(lines, bands, samples)后transpose 时把波段维挪到最后。BIP 最直接每个像素的所有波段连续排布reshape 成(lines, samples, bands)后不需要转置。参数说明offset对应 hdr 里的 header offset指的是二进制文件开头跳过的字节数np.fromfile直接支持这个参数不需要手动切片。我一般会先用一个已知数据集验证一次 interleave确认输出 shape 和数值范围再批量处理其他文件。2.2 Dataset 返回协议从 (H, W, B) 到 (B,) 的两条路线PyTorch 的 Dataset 只约定两件事__len__返回样本数量__getitem__返回一个样本和它的标签。高光谱数据在这里分岔成两种任务逐像素分类和 Patch 分类。逐像素分类的场景下每个样本是一个长度为 B 的光谱向量标签是类别 idPatch 分类的场景下每个样本是一个(B, p, p)的小立方体标签是中心像素或整块 mask。第一条路线的代码很轻from torch.utils.data import Dataset import torch class HsiPixelDataset(Dataset): def __init__(self, cube, labels): # cube: (H, W, B), labels: (H, W)-1 表示背景 self.cube torch.from_numpy(cube).float() self.labels torch.from_numpy(labels).long() valid labels.reshape(-1) 0 self.indices torch.nonzero(valid).squeeze(1).tolist() def __len__(self): return len(self.indices) def __getitem__(self, idx): y, x self.indices[idx] pixel self.cube[y, x] # (B,) label self.labels[y, x] return pixel, label逻辑说明self.indices在__init__里提前算好避免每次__getitem__都遍历全图找有效像素。背景像素 label 为 -1直接被过滤掉。__getitem__返回的pixel是(B,)DataLoader 在堆叠 batch 时会自动拼成(batch_size, B)正好是 Linear 层和 Conv1d 的输入格式。参数说明cube传入时保持(H, W, B)不要在 Dataset 里先把波段维提前因为后面做 Patch 采样和增强都要按空间坐标切片。如果直接拿(H, W, B)的数组转成 Tensor 后丢给模型等来的大概率是维度不匹配的报错。PyTorch 的 Conv2d 期望输入是(N, C, H, W)RGB 图是(H, W, 3)高光谱是(H, W, B)B 在这里等价于 C。所以返回样本前要把波段维挪到第 0 维这就是下面 Patch 采样要做的事。2.3 DataLoader 参数在三维影像上的真实作用加载高光谱数据时DataLoader 的常用参数不是摆设。batch_size受显存限制一个(B, 11, 11)的 patch 不占多少显存但如果错误返回整幅图一个 batch 就能把显存打满。num_workers决定有多少个进程在后台预处理数据高光谱读取慢建议从 4 起步但不是越大越好后面避坑章会专门说。pin_memory在 GPU 训练时值得开可以减少 host 到 device 的拷贝时间。from torch.utils.data import DataLoader loader DataLoader( dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue, )参数说明shuffleTrue对逐像素采样尤其重要因为同一张影像里相邻像素高度相关不打乱的话每个 batch 基本是同质的模型更新方向会来回抖。drop_lastTrue用来扔掉最后一批不完整的样本防止 BatchNorm 在 batch 尺寸过小时统计不准。num_workers4是常规起点如果机器有 16 核可以试 8但每次修改后都要看训练流程里 DataLoader 的耗时占比而不是光看 CPU 占用率。 注意shuffle参数只有在使用默认sampler时才生效如果自定义了 Sampler比如后面要讲的WeightedRandomSamplershuffle会被忽略。3. 像素级与 Patch 级采样两种高光谱任务的 DataLoader 写法高光谱数据的大头不是存储而是怎么从三维立方体里高效地切样本。逐像素分类把每个像素当作独立样本计算快但忽略了空间上下文Patch 分类把每个样本扩成一个窗口能利用邻域光谱相关性模型上限更高代价是样本重叠更多、加载更慢。3.1 像素级分类每个样本只有一个光谱向量2.2 里的HsiPixelDataset就是像素级方案。它的问题在于阈值为 -1 的像素可能只有几千个背景像素却有几十万个如果不做处理训练集会严重偏向背景。常见做法是在__init__里对每个类别做定额采样保证每个 epoch 里各类别数量接近。class BalancedPixelDataset(HsiPixelDataset): def __init__(self, cube, labels, samples_per_class200): super().__init__(cube, labels) self.samples_per_class samples_per_class self.indices [] classes torch.unique(self.labels) for c in classes: cls_mask self.labels.reshape(-1) c cls_idx torch.nonzero(cls_mask).squeeze(1).tolist() if len(cls_idx) samples_per_class: cls_idx random.sample(cls_idx, samples_per_class) self.indices.extend(cls_idx)逻辑说明继承上个 Dataset把self.indices替换成每个类别抽样后的坐标列表。这样训练时每类样本量一致避免背景类主导损失。参数说明samples_per_class可以按最小类别的样本数来定比如最小类只有 300 个样本就设 300保证所有类别都参与。注意这里没有做放回采样类别严重的场景建议用 3.3 的WeightedRandomSampler。3.2 Patch 级分类带邻域信息的立方体与边界过滤Patch 采样看起来只是多切了一个窗口实际上最容易翻车的是边界。以中心像素(cy, cx)裁(cy-r:cyr1, cx-r:cxr1)如果中心点离边界小于 r切片会越界numpy 会静默截断导致返回的 patch 尺寸不一致。处理办法有两种一是像下面这样在构建索引时把边界像素直接过滤二是用 padding但 padding 值会引入不存在的光谱信息。from torch.utils.data import Dataset import torch class HsiPatchDataset(Dataset): def __init__(self, cube, labels, patch_size11, valid_indicesNone): self.cube torch.from_numpy(cube).float() self.labels torch.from_numpy(labels).long() self.p patch_size self.h, self.w, self.b cube.shape r patch_size // 2 if valid_indices is None: valid_indices [] for yi in range(r, self.h - r): for xi in range(r, self.w - r): if labels[yi, xi] 0: valid_indices.append((yi, xi)) self.indices valid_indices def __len__(self): return len(self.indices) def __getitem__(self, idx): cy, cx self.indices[idx] r self.p // 2 patch self.cube[cy - r: cy r 1, cx - r: cx r 1] patch patch.permute(2, 0, 1) # (B, p, p) label self.labels[cy, cx] return patch, label逻辑说明构建索引时直接用range(r, h-r)和range(r, w-r)从源头排除越界情况比在__getitem__里 try-except 干净。返回前用permute(2, 0, 1)把(p, p, B)换成(B, p, p)这是 Conv2d 的标准输入。参数说明patch_size必须是奇数否则中心像素没法定义。实际使用时我一般先打印一个样本的 shape 确认是(B, p, p)再往下走这一步卡住的次数比想象中多。3.3 类别不均衡处理WeightedRandomSampler 的正确打开方式Patch 采样后类别不均衡问题依然存在而且因为样本重叠单纯过滤边界还不够。PyTorch 自带的WeightedRandomSampler是为这个场景设计的它按样本权重做放回采样每个 epoch 可以指定采样总数。from torch.utils.data import WeightedRandomSampler labels_np labels.reshape(-1) valid_mask labels_np 0 valid_labels labels_np[valid_mask] class_counts np.bincount(valid_labels) sample_weights 1.0 / class_counts[valid_labels] weights torch.from_numpy(sample_weights).double() sampler WeightedRandomSampler( weights, num_sampleslen(valid_labels), replacementTrue, ) loader DataLoader(dataset, batch_size64, samplersampler)逻辑说明sample_weights越大该样本被抽中的概率越高。用1.0 / class_counts让少数类权重更大。num_samples控制每个 epoch 采多少样本通常设成有效样本总数这样每轮迭代次数和其他方案一致。参数说明replacementTrue代表放回采样少数类样本可以被重复抽到配合num_samples可以完全替代手动平衡逻辑。 注意一旦传入samplerDataLoader 里的shuffle参数就不再生效这两个参数是互斥的。4. 归一化与数据增强在加载管线里的位置决定训练上限高光谱数据的数值量纲很乱有的波段范围是 0 到 1 的反射率有的是 0 到 65535 的原始 DN 值。如果不做归一化模型训练前期会花很长时间在尺度抖动上。而归一化做在哪一步、统计量从哪里来是很多人踩坑的重灾区。4.1 波段标准化先算统计量再进 Dataset最省事也最容易出问题的做法是在拿到整幅影像后直接(cube - mean) / std。问题是这个 mean 和 std 如果来自整幅图等于让验证集和测试集的分布参与了训练集的统计属于数据泄漏。正确顺序是先切出训练集只在训练集上计算每个波段的均值和标准差然后把这个统计量传给 Dataset 做变换。class BandNorm: def fit(self, cube, maskNone): # cube: (H, W, B) if mask is None: mask np.ones((cube.shape[0], cube.shape[1]), dtypebool) pixels cube[mask] self.mean_ pixels.mean(axis0) self.std_ pixels.std(axis0) 1e-6 def transform(self, cube): return (cube - self.mean_) / self.std_逻辑说明fit只接收训练部分的像素mask用来排除背景区域避免背景像素的极端值污染统计量。std_加1e-6是防止某个波段方差为 0 导致除零高光谱里近红外波段有时确实会出现近似常数的情况。使用方式是在 Dataset 构造时传入mean_和std_在__getitem__里对裁剪出的 patch 做变换而不是把整个cubetransform 一遍再存。参数说明有的任务更偏好 min-max 归一化把数值压到[0, 1]那就在fit里记录min_和max_本质都一样关键在于统计量只来自训练集。也可以把归一化直接放进 Datasetclass NormPatchDataset(HsiPatchDataset): def __init__(self, cube, labels, mean, std, patch_size11): super().__init__(cube, labels, patch_size) self.mean torch.from_numpy(mean).float() self.std torch.from_numpy(std).float() def __getitem__(self, idx): patch, label super().__getitem__(idx) patch (patch - self.mean.view(-1, 1, 1)) / self.std.view(-1, 1, 1) return patch, label逻辑说明mean.view(-1, 1, 1)把(B,)的向量变成(B, 1, 1)才能和(B, p, p)的 patch 做广播相减。这样每个 patch 都在进模型前完成标准化而且统计量是训练集预计算的不会泄漏。 注意别在__getitem__里对每个样本重新计算 mean/std那样每个 patch 的尺度都不一样模型会把尺度差异误当成有用特征。4.2 光谱与空间增强torch 函数直接作用在返回样本上高光谱增强分两个方向空间增强翻转、旋转和光谱增强加噪、波段扰动。空间增强跟普通图像处理几乎一样只是要在 patch 和中心标签上保持一致变换光谱增强是高光谱特有的手段模拟光照和传感器变化。import random import torch def spectral_augment(patch, p0.5): if random.random() p: return patch scale torch.rand(patch.shape[0], 1, 1) * 0.1 0.95 return patch * scale def spatial_augment(patch, label): k random.randint(0, 3) patch torch.rot90(patch, k, dims(1, 2)) if random.random() 0.5: patch torch.flip(patch, dims[2]) return patch, label逻辑说明spectral_augment对每个波段乘一个 0.95 到 1.05 之间的随机系数模拟波段增益变化因为相近波段的增益漂移通常是整体的。spatial_augment的旋转是围绕空间维操作的波段维不动。如果 label 是中心像素的类别旋转和翻转后中心不变label 不用动但如果是逐像素分割任务mask 也必须执行同样的旋转翻转否则标签对不上。4.3 波段选择在 getitem 里做索引比改原始数据更灵活高光谱的波段之间相关性强比如相近的可见光波段信息高度重复很多任务并不需要全部波段。与其在预处理阶段裁剪数据不如在 Dataset 里维护一个波段索引列表这样换波段集合不用重新生成文件。class BandSelectPatchDataset(HsiPatchDataset): def __init__(self, cube, labels, band_idx, patch_size11): super().__init__(cube, labels, patch_size) self.band_idx band_idx def __getitem__(self, idx): patch, label super().__getitem__(idx) return patch[self.band_idx], label逻辑说明patch[self.band_idx]在波段维做索引可以是连续切片如list(range(30, 180, 2))也可以是根据 PCA 排序选出的特定波段。好处是训练脚本里可以通过修改band_idx快速做消融实验不用反复读写磁盘。参数说明波段选择要在归一化之前还是之后我一般先选波段再归一化能少算很多无关波段的统计量如果用了BandNorm那统计量也要用选出的波段来 fit否则尺度会不一致。5. 高光谱 DataLoader 避坑5 条血泪经验与修复方法这一章写的都是实际跑 PyTorch 高光谱训练时最容易遇到的翻车现场。每条按现象、原因、解决的顺序说省得你在报错信息里反复猜。5.1 训练集 loss 正常验证集全程不降现象训练集准确率 95% 以上验证集准确率一直在 60% 徘徊换了模型结构也没有改善。原因Patch 采样的方式有重叠。相邻样本的中心像素距离只有 1 到 2 个像素它们的 patch 大部分区域是重叠的如果把同一个影像的这些 patch 既放进训练集又放进验证集验证集里全是训练集见过或几乎没变过的内容模型记忆的是像素位置而不是光谱特征。解决按影像切分数据而不是按 patch 切分。比如一个数据集有 10 张高光谱影像拿 8 张参与训练集采样另外 2 张影像全部划给验证集。这样训练集和验证集完全不共享 patch。如果再狠一点可以按地理坐标区域切分但通常按图切已经够了。5.2 显存溢出把整幅图当 sample 放进了 batch现象batch_size16输入模型前向时 CUDA out of memory把 batch_size 降到 2 还是炸。原因Dataset 的__getitem__返回了整幅(H, W, B)影像DataLoader 把 16 个这样的立方体堆在一起直接塞满显存。高光谱影像单张 100 MB 起步一个 batch 就是 1.6 GB再加上模型参数和中间激活显存不够很正常。解决确认__getitem__返回的是(B,)像素向量或(B, p, p)patch。打印一行print(dataset[0][0].shape)如果看到的是(512, 512, 224)这种三维 shape那就是把整幅图返回了。把采样逻辑改成 3.2 的 Patch 版本patch_size从 9 或 11 开始试。5.3 num_workers 越大越慢文件句柄和类型转换在作怪现象num_workers0时一个 epoch 10 分钟改成num_workers8后变成 15 分钟CPU 占用率还不低。原因每个 worker 进程在每次__getitem__里都重新打开了一次 HDF5 文件句柄或者反复做 numpy 转 tensor、维度转置。文件打开和类型转换都是重操作多进程并行反而把磁盘 IO 和内存拷贝拖满了。解决把文件的读取放在worker_init_fn里让每个 worker 启动时只打开一次文件后续直接索引读取。另外如果影像尺寸能塞进内存第一个 epoch 就把转好的float32tensor 缓存下来后面不重复转换。def worker_init_fn(worker_id): global f f h5py.File(h5_path, r) loader DataLoader( dataset, batch_size64, num_workers4, worker_init_fnworker_init_fn, )逻辑说明worker_init_fn会在每个 worker 进程启动时执行一次全局变量f在该进程内被复用避免每次__getitem__都h5py.File(...)。如果你的数据本来就缓存成了 numpy这个坑基本不会遇到但 HDF5 大文件场景下很值得做。5.4 DataLoader 返回维度奇怪模型第一层就报错现象模型第一层是nn.Linear(in_featuresB, ...)结果前向时报错说输入最后一维不是 B。原因Dataset 返回的pixel是(B,)DataLoader 堆叠成(batch_size, B)这个没问题。但 Patch 模式里如果没做permute返回的是(p, p, B)DataLoader 堆叠后变成(batch_size, p, p, B)而 Conv2d 期望(batch_size, B, p, p)维度顺序错了。还有一种情况是标签 shape 不小心带进来了__getitem__返回了[patch, label]以外的多余字段。解决在写 Dataset 后、开训前的第一个动作就是打印两个 shapedataset[0][0].shape和loader.dataset[0][0].shape。Patch 模式确认是(B, p, p)像素模式确认是(B,)。维度不对就加permute这一步花两分钟能省半夜的排查时间。5.5 每轮 epoch 结果不一样随机源没有固定现象同一个模型、同一份数据两次训练的结果差异很大甚至跑第一个 epoch 时 loss 初始值都不同。原因DataLoader 的 worker 进程使用了 Python 的随机源shuffleTrue时没有固定 seed同时数据集里如果用了random.sample或 PyTorch 的torch.rand它们的随机状态不受torch.manual_seed控制。解决训练脚本开头统一设置种子并为每个 worker 独立设置随机状态。import torch import random import numpy as np def set_seed(seed42, worker_idNone): random.seed(seed worker_id if worker_id is not None else seed) np.random.seed(seed worker_id if worker_id is not None else seed) torch.manual_seed(seed) set_seed(42) loader DataLoader( dataset, batch_size64, shuffleTrue, num_workers4, worker_init_fnlambda wid: set_seed(42, wid), )逻辑说明worker_init_fn里调用set_seed保证每个 worker 的 shuffle 顺序固定。torch.manual_seed管的是 PyTorch 的随机源但 DataLoader 里的索引打乱走的是 Python 的随机模块所以两个都要设。参数说明seed选什么数字无所谓关键是固定下来后每次实验结果可比复现问题才有意义。6. 把加载时间压到五分之一缓存、流式读取与验证方法高光谱训练的 GPU 利用率低很多时候不是网络太慢而是 DataLoader 在拖后腿。针对不太大的数据集比如内存能装下全部影像最直接的办法是第一个 epoch 构建内存缓存之后所有 epoch 直接读缓存。class CachedPatchDataset(HsiPatchDataset): def __init__(self, cube, labels, patch_size11): super().__init__(cube, labels, patch_size) self.cache {} def __getitem__(self, idx): if idx in self.cache: return self.cache[idx] patch, label super().__getitem__(idx) self.cache[idx] (patch, label) return patch, label逻辑说明这个缓存在第一个 epoch 被逐步填满第二个 epoch 开始不再走 numpy 切片、permute 和类型转换DataLoader 变成纯内存拷贝。配合num_workers4和pin_memoryTrue加载耗时能明显下降。代价是内存占用会涨缓存(B, 11, 11)的 float32 样本一个样本只有 4 * 224 * 121 约 108 KB10 万个样本也就 10 GB 上下很多机器扛得住。参数说明如果内存不够可以用torch.save把预处理后的 patch 缓存成文件或者把cache改成 LRU 限制最大条数。对于超大影像流式读取是更稳的选择。h5py 支持按坐标切片读取Dataset 里只取需要的 patch 块而不是[()]把全图拉进内存。配合worker_init_fn打开文件句柄每个 worker 按需读取能跑完内存放不下的数据。验证加速效果时不要只看单步__getitem__速度要在训练脚本里统计每个 epoch 的 DataLoader 耗时占比。import time loader DataLoader(dataset, batch_size64, num_workers4) t0 time.time() for batch in loader: pass data_time time.time() - t0 print(fdata loading time: {data_time:.2f}s)逻辑说明这里只遍历 DataLoader 不执行模型测的是纯数据时间。如果这个时间占整个训练流程的一半以上优先优化 DataLoader如果只占 10%没必要在缓存上过度设计。从那以后我每次拿到新的高光谱数据第一件事都是先看数据类型和 interleave再写一个最小 Dataset 跑通一个 batch确认 shape 和数值范围最后才上完整训练管线。这套流程跑熟之后高光谱数据从原始文件到模型的调试时间基本都能控制在半小时以内希望帮到你。本文还有配套的精品资源点击获取