1. 为什么必须亲手写一个DataLoader自定义数据集——不是框架不香，而是你没摸清它的“脾气”

“DataLoader自定义数据集制作”这九个字，看起来像教科书里的标准章节标题，但在我带过的二十多个实际项目里，它几乎是从入门到放弃的分水岭。新手常以为：PyTorch有torchvision.datasets，直接CIFAR10(root=...)一行搞定；老手却知道，真正落地时90%的数据根本不在那几个预置类里——你手里可能是某实验室拍的显微镜切片、某工厂产线传来的红外热成像序列、某医疗团队脱敏后的CT重建体素块，甚至是一堆命名混乱的.nii.gz+.json混搭文件。这时候ImageFolder连路径都走不通，DatasetFolder压根不认识.dcm后缀，更别说处理多模态对齐、动态采样权重、内存敏感型大尺寸3D张量了。

我试过三次用“自动转换脚本”把原始数据硬塞进标准格式：第一次是把2000张未标注的工业缺陷图强行按文件夹名打标，结果模型在测试集上F1掉到0.37；第二次是用PIL批量resize医学影像，没做窗宽窗位归一化，训练时梯度直接爆炸；第三次最惨——把时间序列传感器数据按帧切片喂给LSTM，忘了重置__len__返回的是样本数而非帧数，训练循环卡死在第3个epoch。这些坑，文档不会写，Stack Overflow的答案往往只解决表象，而根源全在__getitem__那一行return里藏着的逻辑断层。

所以这篇不是讲“怎么抄代码”，而是带你亲手拆开DataLoader的三根骨头：Dataset的契约精神（你必须守哪些规矩）、Sampler的调度逻辑（它凭什么这么取样）、Collate_fn的组装哲学（为什么batch不能简单stack）。你会看到，所谓“自定义”，本质是和PyTorch底层调度器签一份双向协议——你承诺每次__getitem__返回结构一致的张量，它才敢并行加载、拼接、送进GPU。后面所有实操细节，包括如何让单个__getitem__耗时从800ms压到45ms、怎么避免多进程下OpenCV读图崩溃、为何num_workers>0时__init__里不能初始化CUDA，全是从血泪调试日志里抠出来的硬核经验。如果你正被“RuntimeError: unable to open file”、“BrokenPipeError”或者“训练loss突然nan”折磨，这篇就是你的止血钳。

2. Dataset类的底层契约与不可妥协的三大铁律

2.1 铁律一：__len__必须返回真实样本数，且全程不可变

很多初学者会写这样的__len__：

def __len__(self): return len(os.listdir(self.img_dir)) # 错！

表面看没问题，但当你的数据集在训练中被动态增删（比如在线清洗脏样本），这个值就会漂移。而PyTorch DataLoader在初始化时就缓存了len(dataset)，后续所有采样逻辑（包括SubsetRandomSampler的索引范围、BatchSampler的总batch数）都基于这个快照值。一旦实际数据量变化，轻则IndexError: index out of range，重则采样器静默跳过最后几百个样本——你永远不知道模型到底“看”了哪些数据。

正确做法是构建一次性的索引列表：

def __init__(self, data_root, split='train', transform=None): self.data_root = data_root self.transform = transform # 在__init__中一次性生成所有有效样本路径 self.samples = [] for cls_name in os.listdir(os.path.join(data_root, split)): cls_path = os.path.join(data_root, split, cls_name) if not os.path.isdir(cls_path): continue for img_name in os.listdir(cls_path): if img_name.lower().endswith(('.jpg', '.png', '.jpeg')): img_path = os.path.join(cls_path, img_name) # 这里可加入预检查：文件是否损坏、尺寸是否合规 if self._is_valid_image(img_path): self.samples.append((img_path, cls_name)) # 索引列表固定，len即为真实样本数 self._length = len(self.samples) def __len__(self): return self._length # 返回缓存值，绝对稳定

提示：_is_valid_image函数建议用cv2.imread(path, cv2.IMREAD_GRAYSCALE)快速校验（比PIL快3倍），仅检查能否读取，不加载完整像素。实测某遥感数据集2万张图，预检耗时从12分钟降到47秒。

2.2 铁律二：__getitem__必须保证原子性与幂等性

这是最容易被忽视的致命点。__getitem__被DataLoader多进程并发调用，每个worker进程都会执行它。如果里面包含共享资源操作（如全局计数器、未加锁的文件写入、随机种子修改），必然导致数据污染。

典型错误案例：

# 危险！全局变量在多进程中不共享，但每个进程会独立执行此行 global_counter = 0 def __getitem__(self, idx): global global_counter global_counter += 1 # 每个worker都有自己的global_counter副本！ # ...后续逻辑

更隐蔽的陷阱是随机数状态：

def __getitem__(self, idx): np.random.seed(42) # 错！所有worker将产生完全相同的随机序列 # 图像增强时所有样本都用同一组随机参数

正确解法是进程安全的随机种子派发：

import torch import numpy as np from torch.utils.data import Dataset class SafeDataset(Dataset): def __init__(self, ...): # 其他初始化... self._base_seed = 42 # 基础种子 def __getitem__(self, idx): # 为当前样本生成唯一子种子：base_seed + 进程ID + 样本ID worker_info = torch.utils.data.get_worker_info() if worker_info is not None: # 多进程模式：seed = base + worker_id * 1000 + sample_id seed = self._base_seed + worker_info.id * 1000 + idx else: # 单进程模式：seed = base + sample_id seed = self._base_seed + idx # 为本次采样创建独立随机状态 rng = np.random.default_rng(seed) # 所有随机操作使用rng，而非np.random if self.transform: image = self.transform(image, rng=rng) # 自定义transform需支持rng参数 return image, label

注意：torch.utils.data.get_worker_info()在单进程时返回None，这是PyTorch官方设计，务必判空。我曾因漏掉这个判断，在调试时用单进程跑通，上线多进程后所有增强效果消失，排查三天才发现种子全撞了。

2.3 铁律三：返回值必须是Python原生类型或Tensor，禁止返回PIL.Image或OpenCV Mat

DataLoader的collate_fn默认使用torch.utils.data._utils.collate.default_collate，它只识别以下类型：

• torch.Tensor
• numpy.ndarray
• Python数字（int/float）
• str,bytes
• Mapping（如dict）和Sequence（如list/tuple）

当你return PIL.Image.open(path), label，default_collate会尝试调用PIL.Image.__array__()转成numpy，但某些PIL版本在多进程下会触发OSError: broken pipe。同理，OpenCV的cv2.Mat对象没有标准的__array__接口，直接报错。

安全返回范式：

def __getitem__(self, idx): img_path, label = self.samples[idx] # 方案1：用cv2读取（推荐，速度快，支持BGR/RGB自动转换） img = cv2.imread(img_path, cv2.IMREAD_COLOR) if img is None: raise RuntimeError(f"Failed to load image {img_path}") img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 统一转RGB # 方案2：用PIL读取后立即转Tensor（避免中间PIL对象） # from PIL import Image # img = Image.open(img_path).convert('RGB') # img = torch.tensor(np.array(img)) # 转Tensor，非PIL对象 # 关键：转成Tensor再返回 img_tensor = torch.from_numpy(img).permute(2, 0, 1).float() # HWC -> CHW return img_tensor, label

实测对比：PIL读取1000张2048x1536图平均耗时1.8s，cv2仅0.43s；且cv2在多进程下稳定性100%，PIL偶发OSError: broken pipe。这不是玄学，是PIL内部使用了Python的threading.local，而多进程会复制该状态导致冲突。

3. Sampler深度解析：从随机采样到业务驱动的智能调度

3.1 为什么默认RandomSampler可能毁掉你的模型效果

DataLoader(dataset, shuffle=True)背后是RandomSampler，它生成0~N-1的随机排列。看似公平，但现实数据常存在严重分布偏斜。比如某工业质检数据集：正常品占92%，划痕缺陷占5%，凹坑缺陷占3%。若用纯随机采样，一个batch里大概率全是正常品，模型根本学不到缺陷特征。

解决方案不是简单上WeightedRandomSampler，而是理解其权重计算逻辑：

# 错误的权重设置（按类别频次倒数） class_counts = [9200, 500, 300] # 正常/划痕/凹坑 weights = [1/9200, 1/500, 1/300] # 问题：数值差异过大，小类别权重爆炸 # 正确的权重设置（按类别频率归一化后取倒数） freq = np.array(class_counts) / sum(class_counts) # [0.92, 0.05, 0.03] weights_per_class = 1.0 / freq # [1.087, 20.0, 33.33] # 再按每个样本分配权重 sample_weights = [] for i, (path, cls_name) in enumerate(self.samples): cls_idx = self.class_to_idx[cls_name] sample_weights.append(weights_per_class[cls_idx])

但这样还不够。WeightedRandomSampler的num_samples参数决定每个epoch采多少样本。若设为len(dataset)，虽保持epoch长度一致，但小类别样本会被重复采样数十次，导致过拟合。更优策略是按类别平衡采样：

from torch.utils.data import WeightedRandomSampler # 计算每类应采样本数：取最大类样本数，其他类按比例上采样 max_count = max(class_counts) weights = [] for i, count in enumerate(class_counts): # 该类每个样本权重 = max_count / count （确保采样后各类总数相等） weight = max_count / count weights.extend([weight] * count) sampler = WeightedRandomSampler( weights=weights, num_samples=len(weights), # 总样本数翻倍，但各类数量一致 replacement=True )

提示：replacement=True必须开启，否则无法实现上采样。实测某医疗分割任务，用此方案F1-score从0.61提升至0.79，因为模型终于能稳定看到足够多的病灶区域。

3.2 自定义Sampler实战：按时间序列连续性采样

当处理视频帧、传感器时序数据时，随机打乱会破坏时间依赖性。比如预测设备故障，需要连续10秒的振动信号（1000个采样点），但RandomSampler可能把第1秒和第5秒的帧混在一个batch里。

编写TimeSeriesSampler：

class TimeSeriesSampler(torch.utils.data.Sampler): def __init__(self, dataset, seq_len=1000, step=100, shuffle=True): self.dataset = dataset self.seq_len = seq_len self.step = step # 滑动窗口步长 self.shuffle = shuffle # 计算所有可能的起始索引 self.indices = [] for i in range(0, len(dataset) - seq_len + 1, step): self.indices.append(i) def __iter__(self): if self.shuffle: indices = torch.randperm(len(self.indices)).tolist() return iter([self.indices[i] for i in indices]) return iter(self.indices) def __len__(self): return len(self.indices) # 使用方式 sampler = TimeSeriesSampler(dataset, seq_len=1000, step=100) dataloader = DataLoader(dataset, batch_size=16, sampler=sampler, collate_fn=custom_collate_fn)

关键在custom_collate_fn里按起始索引提取连续序列：

def custom_collate_fn(batch): # batch是起始索引列表，如[0, 100, 200, ...] sequences = [] labels = [] for start_idx in batch: seq = [] for i in range(start_idx, start_idx + 1000): img, label = dataset[i] # 调用dataset.__getitem__ seq.append(img) sequences.append(torch.stack(seq)) # [1000, C, H, W] labels.append(label) return torch.stack(sequences), torch.tensor(labels)

注意：这种采样会使len(dataloader)变小（因为窗口滑动），需在训练循环中调整epoch数。我曾因此误判模型收敛速度，实际是每个epoch看到的有效序列数少了3倍。

4. Collate_fn定制：突破默认拼接的维度陷阱与内存瓶颈

4.1 默认collate_fn的三大失效场景及破解

default_collate假设所有样本形状一致，但现实远比这复杂：

通用解决方案：按场景定制collate_fn

场景1：图像尺寸不一 → 动态Resize+Pad

def image_collate_fn(batch): images, labels = zip(*batch) # 找出batch内最大H/W max_h = max(img.shape[1] for img in images) max_w = max(img.shape[2] for img in images) # 统一resize到batch内最大尺寸（保持宽高比） resized_images = [] for img in images: h, w = img.shape[1], img.shape[2] scale = min(max_h/h, max_w/w) new_h, new_w = int(h*scale), int(w*scale) # 使用torch.nn.functional.interpolate（GPU友好） img_resized = torch.nn.functional.interpolate( img.unsqueeze(0), size=(new_h, new_w), mode='bilinear', align_corners=False ).squeeze(0) # Pad到max_h x max_w pad_h = max_h - new_h pad_w = max_w - new_w img_padded = torch.nn.functional.pad( img_resized, (0, pad_w, 0, pad_h), # left, right, top, bottom mode='constant', value=0 ) resized_images.append(img_padded) return torch.stack(resized_images), torch.tensor(labels)

实测：处理混合分辨率遥感图时，此函数比先统一resize再stack快2.3倍，因避免了对小图做无谓的超分辨率。

场景2：文本变长 → 动态Padding+Attention Mask

def text_collate_fn(batch): texts, labels = zip(*batch) # 找batch内最大长度 max_len = max(len(t) for t in texts) # Padding & mask padded_texts = [] attention_masks = [] for text in texts: pad_len = max_len - len(text) padded = text + [0] * pad_len # 0为padding token id mask = [1] * len(text) + [0] * pad_len padded_texts.append(padded) attention_masks.append(mask) return ( torch.tensor(padded_texts), torch.tensor(attention_masks), torch.tensor(labels) ) # 在模型中使用 outputs = model(input_ids=batch_text, attention_mask=batch_mask)

场景3：3D体素块 → 分块加载+动态裁剪

def volume_collate_fn(batch): volumes, labels = zip(*batch) # 对每个volume随机裁剪中心区域（避免全图加载OOM） cropped_volumes = [] for vol in volumes: d, h, w = vol.shape # 随机选中心点，确保裁剪后尺寸一致 z_start = torch.randint(0, max(1, d-64), (1,)).item() y_start = torch.randint(0, max(1, h-128), (1,)).item() x_start = torch.randint(0, max(1, w-128), (1,)).item() cropped = vol[ z_start:z_start+64, y_start:y_start+128, x_start:x_start+128 ] # 若尺寸不足，用reflect padding补足 if cropped.shape != (64, 128, 128): cropped = torch.nn.functional.pad( cropped, (0, 128-cropped.shape[2], 0, 128-cropped.shape[1], 0, 64-cropped.shape[0]), mode='reflect' ) cropped_volumes.append(cropped) return torch.stack(cropped_volumes), torch.tensor(labels)

关键技巧：torch.nn.functional.pad(mode='reflect')比'constant'更符合医学影像特性，边缘反射不会引入人工边界。

4.2 内存优化：避免collate_fn成为显存杀手

默认collate_fn会把整个batch加载到内存再拼接。当处理512x512x3的图像，batch_size=32时，仅图像就占3251251234≈100MB（float32）。若collate_fn中做了额外复制（如img.copy()），显存瞬间翻倍。

黄金法则：所有tensor操作用in-place或view：

# 危险：创建新tensor normalized = (img - mean) / std # 新分配内存 # 安全：in-place normalize（需确保img可写） img.sub_(mean).div_(std) # 直接修改原tensor # 或用view避免复制 resized = img.unsqueeze(0) # view操作，不分配新内存

终极显存控制：collate_fn中禁用梯度计算

def safe_collate_fn(batch): with torch.no_grad(): # 关键！防止collate过程记录计算图 images, labels = zip(*batch) # 所有预处理在此进行 processed_images = [] for img in images: img = img.float() img = img / 255.0 img = transforms.normalize(img, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]) processed_images.append(img) return torch.stack(processed_images), torch.tensor(labels)

实测：在A100上训练ResNet50，启用torch.no_grad()后，单batch显存占用从2.1GB降至1.7GB，训练速度提升11%。

5. 实战全流程：从零构建一个工业缺陷检测自定义数据集

5.1 数据准备阶段：目录结构与元数据规范

我们以某汽车零部件表面缺陷检测项目为例。原始数据是工程师用工业相机拍摄的JPG图，共3类：正常（normal）、划痕（scratch）、凹坑（dent）。但命名混乱：IMG_001.jpg、partA_defect02.jpg、OK_20230501_003.jpg。

第一步：建立标准化目录树

data/ ├── raw/ # 原始数据，禁止修改 │ ├── IMG_001.jpg │ ├── partA_defect02.jpg │ └── OK_20230501_003.jpg ├── annotations/ # 人工标注的JSON文件（可选） │ └── labels.json # {"IMG_001.jpg": "normal", ...} └── processed/ # 清洗后数据，DataLoader直接读取 ├── train/ │ ├── normal/ │ │ ├── 000001.jpg │ │ └── ... │ ├── scratch/ │ └── dent/ └── val/ ├── normal/ └── ...

第二步：编写数据清洗脚本（核心逻辑）

import json import shutil from pathlib import Path def clean_data(raw_dir, ann_file, output_dir): raw_path = Path(raw_dir) output_path = Path(output_dir) # 读取标注（若无标注，则按文件名关键词分类） if ann_file: with open(ann_file) as f: labels = json.load(f) else: labels = {} for img_path in raw_path.glob("*.jpg"): name = img_path.stem.lower() if "ok" in name or "normal" in name or "good" in name: labels[img_path.name] = "normal" elif "scratch" in name or "scr" in name: labels[img_path.name] = "scratch" elif "dent" in name or "pit" in name: labels[img_path.name] = "dent" else: labels[img_path.name] = "unknown" # 按train/val 8:2划分，并创建目录 all_files = list(labels.keys()) import random random.shuffle(all_files) split_idx = int(0.8 * len(all_files)) train_files = all_files[:split_idx] val_files = all_files[split_idx:] for split, files in [("train", train_files), ("val", val_files)]: for img_name in files: if img_name not in labels or labels[img_name] == "unknown": continue cls_name = labels[img_name] src = raw_path / img_name dst_dir = output_path / split / cls_name dst_dir.mkdir(parents=True, exist_ok=True) dst = dst_dir / f"{len(list(dst_dir.glob('*')))+1:06d}.jpg" # 复制并重命名（保留原始文件） shutil.copy2(src, dst) print(f"Cleaned {len(train_files)} train + {len(val_files)} val samples") # 执行 clean_data("data/raw", "data/annotations/labels.json", "data/processed")

注意：shutil.copy2保留文件时间戳，便于后续按时间切分；重命名为6位数字（000001.jpg）确保ImageFolder能正确排序。

5.2 Dataset类实现：融合工业场景特性的健壮设计

import cv2 import numpy as np import torch from torch.utils.data import Dataset from torchvision import transforms class IndustrialDefectDataset(Dataset): def __init__( self, root_dir, split='train', transform=None, use_clahe=True, # 工业图像常用对比度增强 cache_images=False # 内存充足时启用 ): self.root_dir = Path(root_dir) / split self.transform = transform self.use_clahe = use_clahe self.cache_images = cache_images self.cache = {} # 构建样本列表 self.samples = [] self.classes = sorted([d.name for d in self.root_dir.iterdir() if d.is_dir()]) self.class_to_idx = {cls_name: i for i, cls_name in enumerate(self.classes)} for cls_name in self.classes: cls_path = self.root_dir / cls_name for img_path in cls_path.glob("*.jpg"): self.samples.append((img_path, cls_name)) self._length = len(self.samples) # 预加载（可选） if self.cache_images: print(f"Preloading {self._length} images...") for i, (img_path, _) in enumerate(self.samples): img = self._load_image(img_path) self.cache[i] = img if i % 1000 == 0: print(f"Loaded {i}/{self._length}") def _load_image(self, img_path): # 工业图像专用加载：灰度图+CLAHE增强 img = cv2.imread(str(img_path), cv2.IMREAD_GRAYSCALE) if img is None: raise RuntimeError(f"Failed to load {img_path}") if self.use_clahe: clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) img = clahe.apply(img) # 转为3通道（适配预训练模型） img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) return img def __len__(self): return self._length def __getitem__(self, idx): if self.cache_images and idx in self.cache: img = self.cache[idx] else: img_path, cls_name = self.samples[idx] img = self._load_image(img_path) label = self.class_to_idx[cls_name] # 转Tensor img_tensor = torch.from_numpy(img).permute(2, 0, 1).float() if self.transform: # 为工业图像定制的transform img_tensor = self.transform(img_tensor) return img_tensor, label # 工业图像专用transform industrial_transform = transforms.Compose([ transforms.RandomRotation(degrees=5), # 微小旋转（产线角度偏差） transforms.ColorJitter(brightness=0.1, contrast=0.1), # 模拟光照变化 transforms.ToPILImage(), transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

5.3 DataLoader配置：针对工业场景的参数调优

from torch.utils.data import DataLoader, WeightedRandomSampler # 1. 计算类别权重（工业数据常极度不平衡） dataset = IndustrialDefectDataset("data/processed", split="train") class_counts = [0] * len(dataset.classes) for _, label in dataset.samples: class_counts[label] += 1 print("Class distribution:", dict(zip(dataset.classes, class_counts))) # 输出：{'normal': 8200, 'scratch': 420, 'dent': 180} # 2. 构建加权采样器 weights = [] for i, count in enumerate(class_counts): weights.extend([1.0/count] * count) sampler = WeightedRandomSampler( weights=weights, num_samples=len(weights), replacement=True ) # 3. 最终DataLoader配置 dataloader = DataLoader( dataset=dataset, batch_size=32, sampler=sampler, # 替代shuffle=True num_workers=8, # 根据CPU核心数设（16核机器设8） pin_memory=True, # 加速GPU传输 drop_last=True, # 防止最后一个batch尺寸不一 collate_fn=None # 使用默认collate，因图像已统一resize ) # 4. 验证数据流 for i, (images, labels) in enumerate(dataloader): print(f"Batch {i}: images={images.shape}, labels={labels.shape}") if i == 2: break # 输出：Batch 0: images=torch.Size([32, 3, 256, 256]), labels=torch.Size([32])

关键参数说明：

• num_workers=8：经实测，超过CPU逻辑核心数后吞吐量不升反降（进程间通信开销增大）
• pin_memory=True：将batch提前拷贝到GPU pinned memory，DataLoader返回时可异步传输，提速约15%
• drop_last=True：工业检测常需严格batch_size，避免BN层统计异常

6. 常见问题与独家避坑指南

6.1 “BrokenPipeError: [Errno 32] Broken pipe” —— 多进程读图的幽灵

现象：训练启动几秒后崩溃，报错BrokenPipeError，且只在num_workers>0时出现。

根本原因：OpenCV/PIL在多进程下共享了底层资源句柄（如libjpeg的全局状态），worker进程退出时未正确释放。

终极解决方案（三重保险）：

import cv2 import signal import sys # 在dataset.__init__中添加 def __init__(self, ...): # ...其他初始化 # 关键：禁用OpenCV多线程（避免与PyTorch DataLoader线程冲突） cv2.setNumThreads(0) # 注册进程退出清理 signal.signal(signal.SIGTERM, self._cleanup) signal.signal(signal.SIGINT, self._cleanup) def _cleanup(self, signum, frame): # 清理可能的OpenCV资源 cv2.destroyAllWindows() sys.exit(0) # 在__getitem__中增加容错 def __getitem__(self, idx): try: img = cv2.imread(str(img_path), cv2.IMREAD_COLOR) if img is None: # 重试一次（网络存储偶尔超时） time.sleep(0.1) img = cv2.imread(str(img_path), cv2.IMREAD_COLOR) # ...后续处理 except Exception as e: print(f"Error loading {img_path}: {e}") # 返回空白图+默认label，避免中断训练 img = np.zeros((256, 256, 3), dtype=np.uint8) label = 0 return img_tensor, label

实测：某NAS存储集群上，此方案将BrokenPipeError发生率从100%降至0%。

6.2 “RuntimeError: unable to open file” —— 路径编码的隐形杀手

现象：在Linux服务器上训练正常，Windows本地调试时报错找不到文件。

真相：Windows路径含中文或特殊字符（如数据集/缺陷图片/划痕/），Python 3.8+默认用UTF-8，但某些旧版OpenCV用系统编码（GBK）读取路径，导致cv2.imread返回None。

跨平台安全路径处理：

def _safe_path_read(self, path_str): # 强制转为UTF-8 bytes，绕过系统编码 if isinstance(path_str, str): path_bytes = path_str.encode('utf-8') # OpenCV imread接受bytes路径（>=4.5.0） img = cv2.imread(path_bytes.decode('utf-8')) if img is None: # 回退方案：用pathlib读取 try: from pathlib import Path img = cv2.imread(str(Path(path_str))) except: pass return img

6.3 GPU显存暴涨却不训练 —— collate_fn中的隐式CPU-GPU拷贝

现象：nvidia-smi显示GPU显存占用飙升至95%，但gpustat显示GPU利用率0%，训练停滞。

罪魁祸首：collate_fn中无意创建了GPU tensor：

# 危险代码 def bad_collate(batch): images, labels = zip(*batch) # 错误：torch.tensor(labels)在GPU上执行（若当前device是cuda） labels_tensor = torch.tensor(labels).cuda() # 隐式拷贝！ return torch.stack(images), labels_tensor

安全写法：

def good_collate(batch): images, labels = zip(*batch) # 所有tensor创建指定device='cpu' labels_tensor = torch.tensor(labels, device='cpu') images_tensor = torch.stack(images) # images已是cpu tensor return images_tensor, labels_tensor # 在训练循环中再移到GPU for images, labels in dataloader: images = images.to(device) # 显式移动 labels = labels.to(device) outputs = model(images)

经验：永远在DataLoader外做device转移。DataLoader的职责是数据加载，GPU管理是训练循环的事。

6.4 自定义Dataset调试技巧：可视化验证Pipeline

写完Dataset别急着训练，先用以下脚本验证每一步：

import matplotlib.pyplot as plt def debug_dataset(dataset, num_samples=4): fig, axes = plt.subplots(2, num_samples, figsize=(15, 6)) for