我先说个结论这篇我拖了挺久才动笔。网上一搜pytorch训练自己的数据集出来的大多是互相抄的demo要么只讲分类要么跳过了数据处理直接上模型等你照着敲完换上自己的数据就各种报错。所以这次我按自己实际跑完项目的顺序把从原始图片到模型部署的整个链路完整串一遍不搞花活每一步都给能直接抄的代码和踩过的坑。这篇文章适合谁刚学完PyTorch基础、想拿自己的图片数据跑通一个完整训练流程的人。不管你是做分类、检测还是分割前半段的数据处理和Dataset定义思路是通用的。我会以图像分类为主示例其他任务的区别我会在对应位置指出来。整个项目我用的环境是PyTorch 1.13、Python 3.9、单张消费级显卡显存8GCPU也能跑就是慢一些后面会说到怎么调整。1. 项目拆解与整体设计思路1.1 自定义数据集训练的完整链路训练自己的数据集和用官方现成数据集最大的区别是什么官方数据集已经帮你把图片、标签、划分都做好了你拿到手直接torchvision.datasets.ImageFolder一加载就能开训。但自己的数据往往是这样的图片散落在好几个文件夹里命名混乱图片尺寸不一致有横有竖标签信息可能在文件名里也可能在一个Excel/CSV里正负样本数量极不平衡有些图片是坏的读不出来。所以完整版流程应该是数据准备与清洗数据集划分训练/验证/测试自定义Dataset类与DataLoader模型设计与预训练权重加载训练循环与调参模型评估与保存加载模型做推理每一步都有独立的坑。比如数据划分这步很多人偷懒不划分直接全量训练最后告诉我训练loss很低但测试效果差得离谱——这是典型的没做验证集导致的过拟合。所以我把每一步都当作独立章节讲别嫌啰嗦这些都是真金白银换来的。1.2 技术选型与方案考量核心选型如下环节选型理由框架PyTorch动态图调试方便社区生态好数据加载torchvision.datasets DataLoader与框架集成度高多进程加载模型ResNet18/34 预训练权重小数据集上迁移学习效果远好于随机初始化数据增强torchvision.transforms轻量、够用避免额外依赖训练可视化TensorBoard监控loss和准确率曲线为什么不选Keras或Paddle不是不能用而是PyTorch在这类数据比较乱、需要自定义逻辑的场景下更灵活。比如我后面要写的自定义Dataset类想怎么读数据就怎么读想做什么预处理就在__getitem__里加没有框架约束。模型选ResNet是综合考虑的。ResNet18在消费级显卡上训练速度快显存占用低而预训练权重的泛化能力在小数据集上表现远好于从零训练。如果你的数据量很大每类几千张可以尝试ResNet50或EfficientNet但训练时间会成倍增加。1.3 示例任务设定猫狗二分类为了让全文不飘在抽象层面我固定一个具体任务猫狗二分类。数据集结构如下data/ ├── train/ │ ├── cat/ │ │ ├── cat_001.jpg │ │ ├── cat_002.jpg │ │ └── ... │ └── dog/ │ ├── dog_001.jpg │ └── ... └── test/ ├── cat/ └── dog/这套结构其实就是一个天然的ImageFolder格式PyTorch可以自动把文件夹名当作类别标签。后续我会在这个基础上加难度比如把标签放到CSV里怎么办、类别不均衡怎么办。2. 数据准备与预处理2.1 目录组织结构与标签映射先别急着写代码。数据这块我的建议是动手之前先把目录结构定下来因为后续所有代码都依赖这个结构。最省事的方案直接用ImageFolder格式。每个类别一个文件夹文件夹名就是类别名data/ ├── train/ │ ├── cat/ │ ├── dog/ ├── val/ │ ├── cat/ │ ├── dog/ └── test/ ├── cat/ ├── dog/这样加载的时候直接from torchvision import datasets train_dataset datasets.ImageFolder(rootdata/train) print(train_dataset.class_to_idx) # {cat: 0, dog: 1}ImageFolder会自动把cat映射到0dog映射到1顺序按字母排。这里有个坑如果类别名是class_1、class_2这种字符序和逻辑序不一致的名字映射会乱。所以我建议类别文件夹直接用有意义的英文单词不要用数字开头。如果你的数据已经在别的结构里比如所有图片在一个文件夹里标签在CSV里那就需要写转换脚本。这类情况我会在进阶处理里单独讲。2.2 训练集/验证集/测试集划分划分比例是我在实践中踩过坑的地方。早期做项目时图省事只分训练集和测试集结果调参时没有独立验证集每次改完超参数都用测试集试最后测试集变成了训练集的一部分指标虚高。正确的做法是分成三份训练集用于更新模型参数一般占80%左右验证集用于调超参、选模型占10%-15%测试集只在最终评估时用一次占10%左右。划分代码用torch.utils.data.random_split就行但注意先打乱from torch.utils.data import random_split dataset datasets.ImageFolder(rootdata/all) train_size int(0.8 * len(dataset)) val_size int(0.1 * len(dataset)) test_size len(dataset) - train_size - val_size train_ds, val_ds, test_ds random_split( dataset, [train_size, val_size, test_size], generatortorch.Generator().manual_seed(42) )这里用manual_seed(42)固定随机种子保证每次运行划分结果一致。不然每次跑代码数据划分都不同前后对比实验就没意义了。2.3 数据清洗与异常处理这一步看起来不起眼但真能救命。我接过一个项目2000张图里有30多张是损坏的或根本不是目标物体训练的时候没报错但loss一直在高位抖动。后来排查发现是那批脏数据在搞鬼。数据清洗最简单的方法写个脚本把所有图片统一读一遍检查尺寸和通道数。from PIL import Image import os bad_images [] for root, dirs, files in os.walk(data): for f in files: if not f.lower().endswith((.jpg, .jpeg, .png)): continue path os.path.join(root, f) try: img Image.open(path) img.load() if img.mode not in (RGB, L): bad_images.append((path, mode img.mode)) except Exception as e: bad_images.append((path, str(e))) print(f发现 {len(bad_images)} 个异常文件) for p, reason in bad_images[:20]: print(p, reason)还有一类要特别注意虽然文件夹叫cat但里面混了几张狗图。这个脚本检查不出来只能靠肉眼或者聚类方法排查。项目规模不大的时候我建议把所有图片按文件夹拼成九宫格图看一眼丑是丑点但心里踏实。2.4 尺寸统一与归一化参数计算深度学习模型要求输入尺寸固定我一般用224x224ResNet的标准输入。尺寸统一有两种思路直接resize或者先resize到稍大尺寸再中心裁剪。前者简单但会拉伸变形后者保留更多信息。我推荐第二种from torchvision import transforms train_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) test_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里的mean和std是ImageNet数据集的统计值用预训练模型时这个值必须保持和预训练时一致不能随便改。如果是从零训练应该用自己的数据重新计算均值和标准差import torch from torchvision import datasets, transforms temp_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor() ]) ds datasets.ImageFolder(rootdata/train, transformtemp_transform) loader torch.utils.data.DataLoader(ds, batch_size64, shuffleFalse) mean torch.zeros(3) std torch.zeros(3) n 0 for images, _ in loader: batch_samples images.size(0) images images.view(batch_samples, images.size(1), -1) mean images.mean(2).sum(0) std images.std(2).sum(0) n batch_samples mean / n std / n print(f自定义均值: {mean}, 标准差: {std})注意算统计量的时候不要开数据增强比如随机翻转不然统计值会偏移。3. 自定义Dataset与DataLoader3.1 为什么用ImageFolder就够了如果你的数据已经整理成每个类别一个文件夹的结构ImageFolder就是你最省事的方案不需要写自定义Dataset。它底层做的事情是扫描目录建立文件名到标签的映射__getitem__里读图、做transform、返回图片张量, 标签。但一旦出现下面任一情况ImageFolder就不够用了标签不在文件夹名里而是在CSV/Excel表格中一个样本对应多张图比如视频帧序列需要返回图片路径或额外信息用于调试标签不是简单的类别索引而是检测框坐标或分割掩码。这时候就需要自定义Dataset类。我建议即使ImageFolder能用也花10分钟把自定义Dataset类写了因为后面做数据增强、调试、扩展都方便。3.2 自定义Dataset代码与飞书import os import pandas as pd from PIL import Image from torch.utils.data import Dataset class CustomImageDataset(Dataset): def __init__(self, root_dir, label_fileNone, transformNone): self.root_dir root_dir self.transform transform if label_file is not None: # 从CSV读取标签 self.df pd.read_csv(label_file) self.image_paths self.df[image_path].tolist() self.labels self.df[label].tolist() else: # 从目录结构读取标签ImageFolder风格 self.classes sorted(os.listdir(root_dir)) self.class_to_idx {cls: i for i, cls in enumerate(self.classes)} self.image_paths [] self.labels [] for cls in self.classes: cls_dir os.path.join(root_dir, cls) for fname in os.listdir(cls_dir): if fname.lower().endswith((.jpg, .jpeg, .png)): self.image_paths.append(os.path.join(cls_dir, fname)) self.labels.append(self.class_to_idx[cls]) def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img_path self.image_paths[idx] image Image.open(img_path).convert(RGB) label self.labels[idx] if self.transform: image self.transform(image) return image, label这个类是核心。几个关键点__len__返回数据集大小DataLoader需要知道总样本数__getitem__支持索引访问每次返回一个图片, 标签对打开图片时统一.convert(RGB)把灰度图、RGBA图都转成三通道避免shape不一致报错transform在__getitem__里做意味着每次epoch数据增强都会重新做一遍。注意__getitem__里千万不要用cv2.imread配合transforms.ToTensor()因为cv2读出来是BGR通道顺序而torchvision的模型都按RGB预训练颜色会错乱。要么用 PIL要么用cv2后手动cv2.cvtColor(img, cv2.COLOR_BGR2RGB)。3.3 DataLoader参数配置与坑from torch.utils.data import DataLoader train_loader DataLoader( train_ds, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue if torch.cuda.is_available() else False, drop_lastFalse ) val_loader DataLoader( val_ds, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue if torch.cuda.is_available() else False )参数逐一说明batch_size显存不够就调小可以先用默认值32跑OOM再降到16或8shuffle训练集必须True打乱顺序避免模型学到样本顺序的偏见。验证集和测试集必须Falsenum_workers数据加载的并行进程数。Windows上设成0最稳Linux可以设成CPU核心数。Mac上偶尔有spawn问题报错就调成0pin_memory用GPU时设True把数据固定在锁页内存传输到显存更快drop_last当最后一个batch样本数小于batch_size时是否丢弃。训练时一般False就行但某些BatchNorm层对batch大小敏感如果最后一个batch只有1-2张图可能导致bn报错这时候设True更稳。还有一个Windows专属大坑DataLoader配合自定义Dataset时如果把代码写在Jupyter Notebook里且没有if __name__ __main__保护多进程加载会递归报错。解决办法是把代码存成.py文件运行或者num_workers0。3.4 数据增强的取舍数据增强的本质是扩充样本分布我见过不少人把增强策略堆到天花板随机旋转、缩放、裁剪、颜色抖动、遮挡全上。结果训练集准确率反而下降因为增强过强导致学习难度过大。我的经验是分情况数据量增强策略每类500张轻中度随机水平翻转、随机小角度旋转、颜色微调每类500-2000张增强力度加大RandomResizedCrop、随机旋转±15°、亮度对比度微调每类5000张可以不用太多增强重点防过拟合用Dropout二分类猫狗这种任务随机水平翻转就够了猫狗左右翻转不影响语义。但有些任务不能随便翻转比如数字识别6翻转变成9、文字OCR文字镜像没意义这类任务要谨慎。还有一个被忽视的点测试时不要用数据增强只用中心裁剪。否则你测试集出来的指标每次都不一样没法复现。4. 模型构建从零训练还是迁移学习4.1 两种方案的对比与选择刚入门时我也干过从零搭VGG、ResNet的蠢事结果在小数据集上训练几十个epoch准确率卡在60%上下后来换了迁移学习5个epoch就上90%。这个差距不是模型不行而是数据量撑不起随机初始化的深层网络。方案训练时间小数据集效果代码量从零训练长差过拟合严重多迁移学习特征提取短好少迁移学习微调全网络中等最好少个人建议数据量小于1万张优先迁移学习。如果数据来源于真实场景且和ImageNet分布差异大比如是医学影像而不是自然图片可以冻结前几层只微调后面层。4.2 迁移学习代码实现import torchvision.models as models import torch.nn as nn def get_model(num_classes, pretrainedTrue): model models.resnet18(pretrainedpretrained) # 替换最后一层全连接适配自己的类别数 in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.5), nn.Linear(in_features, num_classes) ) return model model get_model(num_classes2, pretrainedTrue)两个关键点第一pretrainedTrue需要联网下载权重。如果在受限网络环境可以提前下载好然后用如下方式加载import torch state_dict torch.load(resnet18-5c106cde.pth, map_locationcpu) model models.resnet18(pretrainedFalse) model.load_state_dict(state_dict)第二model.fc被我替换成了Dropout Linear的组合。加一个Dropout是防止最后一层过拟合这里0.5是常见配置如果数据很少0.3更稳。如果任务不是分类而是回归比如预测坐标、年龄等把最后一层输出的激活函数去掉nn.Linear(in_features, 1)就行。4.3 最后一层初始化技巧替换了全连接层之后新层是随机初始化的而前面的预训练层已经收敛。训练初期梯度由新层主导如果新层初始化不好可能引起loss剧烈震荡。解决办法是给新层做一次Xavier初始化def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) nn.init.constant_(m.bias, 0) model.fc.apply(init_weights)这一步不是必须的PyTorch默认的初始化也能跑但我发现加上之后训练前期更平稳前10个epoch的loss曲线不再大幅抖动。特别是从零训练时这个初始化方式对收敛速度有明显帮助。4.4 损失函数与优化器配置import torch.nn as nn import torch.optim as optim criterion nn.CrossEntropyLoss() # 优化器 optimizer optim.SGD(model.parameters(), lr0.001, momentum0.9, weight_decay1e-4)CrossEntropyLoss自带Softmax操作所以模型最后一层不需要额外加Softmax。这是新手很容易搞混的点——如果你在模型里加了LogSoftmax又用CrossEntropyLoss就会得到诡异的结果。优化器这块我默认推荐SGD momentum理由是泛化性能通常比Adam好而且学习率策略成熟。Adam收敛快但容易收敛到尖锐极小值测试集表现不稳定。当然这只是经验之谈如果数据比较难拟合Adam可以先用来快速找个好的起点再换SGD精调。4.5 训练设备与显存管理device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)这是最基础的写法。如果你显存不大小于8G几个建议batch_size从32降到16或8使用混合精度训练AMP显存直接减半速度还快。PyTorch 1.6内置支持from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()AMP在原理解释上是前向计算用FP16提速梯度计算用FP32保证精度scaler防止小梯度被FP16下溢吃掉。我用一次之后就没关过算是白捡的优化。5. 训练循环与可视化5.1 标准训练循环模板很多教程把训练代码写得支离破碎我给出一个能直接复制保存的完整模板。这个模板我在多个项目上复用改改模型和数据加载就能跑。def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc def validate(model, val_loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() return running_loss / total, correct / total需要解释两个关键点第一model.train()和model.eval()必须成对出现。前者的作用是开启BatchNorm的统计量更新和Dropout后者关闭它们。很多人忘了切回train模式导致验证时结果忽高忽低或者在训练时用eval模式BN统计量不更新Loss异常。第二验证时用with torch.no_grad()包裹不构建计算图。如果不加验证阶段也会保存每个中间变量的梯度信息显存直接被撑爆。新手几乎都会踩这个坑。total用labels.size(0)而不是labels.shape[0]写法上等价但前者可读性更好。分类问题的total是和images.size(0)相同的因为每个样本一个标签。5.2 完整训练流程与模型保存import copy best_acc 0.0 num_epochs 30 for epoch in range(num_epochs): train_loss, train_acc train_one_epoch( model, train_loader, criterion, optimizer, device ) val_loss, val_acc validate(model, val_loader, criterion, device) print(fEpoch [{epoch1}/{num_epochs}] fTrain Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f} | fVal Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}) # 保存最佳模型 if val_acc best_acc: best_acc val_acc best_model_wts copy.deepcopy(model.state_dict()) torch.save(model.state_dict(), best_model.pth) # 训练结束后加载最佳权重 model.load_state_dict(best_model_wts) print(fBest val Acc: {best_acc:.4f})保存模型的判断依据用验证集准确率而不是训练集。训练集准确率100%只能代表模型在背答案验证集准确率才是真实泛化能力的近似。有几点经验每个epoch保存一次的话训练30个epoch会产生30个文件没必要。只保存最优和最后一轮就行torch.save保存的是state_dict也就是权重参数不是整个模型对象。加载的时候需要先有模型结构再load_state_dictcopy.deepcopy是为了避免后续训练时引用被修改。如果直接best_model_wts model.state_dict()后面模型更新时这个引用也会跟着变最佳权重就被覆盖了。5.3 学习率调整策略固定学习率训练30个epoch到后期会出现loss在某个值附近震荡的现象。这是学习率太大导致的需要按epoch衰减。我常用的策略是StepLR和ReduceLROnPlateau二选一from torch.optim import lr_scheduler # 方案一每10个epoch学习率乘以0.1 scheduler lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) # 方案二验证集loss不下降就降学习率 scheduler lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience3 ) # 训练循环内调用方式不同 # StepLR: 每个epoch结束后调用 scheduler.step() # ReduceLROnPlateau: 每个epoch结束后传入验证loss scheduler.step(val_loss)StepLR的问题是它不看模型实际表现到点就降可能降太早或太晚。ReduceLROnPlateau则更智能验证loss连续3个epoch不降学习率就减半我实际用下来效果更好。学习率初始值的经验微调预训练模型用0.001从零训练用0.01到0.05Adam用0.001以下。初始学习率设得过大前几个epoch loss会出现NaN或巨大跳变这时候不要慌先调小学习率再试。5.4 TensorBoard可视化训练时只看终端输出的loss数值很难发现过拟合的苗头。TensorBoard能画出曲线一眼看出训练集和验证集的gap。from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/experiment_1) # 训练循环内 writer.add_scalar(Loss/train, train_loss, epoch) writer.add_scalar(Loss/val, val_loss, epoch) writer.add_scalar(Acc/train, train_acc, epoch) writer.add_scalar(Acc/val, val_acc, epoch) # 还可以记录模型结构 dummy_input torch.randn(1, 3, 224, 224).to(device) writer.add_graph(model, dummy_input)启动可视化tensorboard --logdirruns --port6006浏览器打开localhost:6006。这个工具胜在轻量不用额外装什么环境PyTorch自带。5.5 断点续训与早停训练中途断了怎么办特别是数据量大、训练几十个小时的场景。我习惯每个epoch保存一次完整的训练状态torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), best_acc: best_acc, }, checkpoint.pth)恢复时checkpoint torch.load(checkpoint.pth) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) scheduler.load_state_dict(checkpoint[scheduler_state_dict]) start_epoch checkpoint[epoch] 1这里面有个细节优化器的state_dict里保存了动量等历史信息如果不恢复断点续训后优化器相当于重新初始化前几个epoch效果会退步。所以完整保存优化器状态是必须的。早停Early Stopping是防止过拟合最简单的手段验证集loss连续N个epoch不下降就提前终止训练。我一般设patience10配合checkpoint保存最佳模型即使提前停了也能回到最佳状态。6. 模型评估与推理部署6.1 验证集评估指标的选择准确率Accuracy是最直观的指标但它有一个致命问题类别不平衡时完全失真。比如99%的猫1%的狗模型全预测猫准确率就有99%但这个模型毫无用处。所以除了准确率我还看这几个指标混淆矩阵Confusion Matrix清楚地看出模型把哪一类混淆了。from sklearn.metrics import confusion_matrix, classification_report, accuracy_score def evaluate_model(model, loader, device): model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) report classification_report( all_labels, all_preds, target_names[cat, dog] ) print(report) print(Confusion Matrix:) print(cm) return accuracy_score(all_labels, all_preds)输出分类报告里有Precision、Recall、F1-score。在样本不均衡时F1-score尤其是macro-F1比准确率更能反映模型真实水平。6.2 模型权重转化与CPU推理训练完的模型要落地到实际场景网页后端、手机App、嵌入式设备核心问题就是把PyTorch模型转成可部署的格式。先看PyTorch自带的推理方式适用于服务端import torch import torchvision.models as models # 先构建模型结构 device torch.device(cpu) model models.resnet18(pretrainedFalse) num_classes 2 model.fc torch.nn.Linear(model.fc.in_features, num_classes) # 加载权重 model.load_state_dict(torch.load(best_model.pth, map_locationcpu)) model.eval()注意map_locationcpu即使模型是在GPU上训练的也能在CPU上加载。如果报错键值对不匹配检查一下是不是保存的时候带了module.前缀用了DataParallel训练。推理时记得关闭梯度from PIL import Image from torchvision import transforms def predict(image_path, model, devicecpu): transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(image_path).convert(RGB) img_tensor transform(img).unsqueeze(0).to(device) model.to(device) model.eval() with torch.no_grad(): outputs model(img_tensor) probs torch.softmax(outputs, dim1) _, predicted torch.max(probs, 1) return predicted.item(), probs[0].tolist() class_idx {0: cat, 1: dog} pred_class, probs predict(test_images/dog_01.jpg, model) print(f预测类别: {class_idx[pred_class]}, 概率: {probs})推理有个小细节unsqueeze(0)给图片加一个batch维度因为模型期望输入是[B, C, H, W]四维张量单张图只有[C, H, W]三维。6.3 部署时的几大坑推理时的预处理必须和训练时完全一致。训练用了Resize(256) CenterCrop(224)推理就不能只用Resize((224, 224))否则图片的缩放比例不同模型输入分布偏移准确率直接掉几个点。不要在预测函数里加model.train()。推理时dropout会被错误地启用输出结果随机。如果CPU推理速度太慢比如ResNet50一张图要100ms以上可以考虑用ONNX导出dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}})ONNX Runtime推理一般比PyTorch原生CPU推理快两倍以上而且不依赖PyTorch环境。7. 常见问题与排查实录7.1 数据维度不匹配报错报错示例RuntimeError: Expected input batch_size (32) to match target batch_size (1)。这个问题的根源是模型输入和标签的batch维度不一致。排查步骤打印images.shape和labels.shape确认是[32, 3, 224, 224]和[32]检查是不是在某个地方对images做了切片或reshape检查自定义Dataset的__getitem__返回值——如果返回了(image, label, extra)而训练循环只接收两个变量就会错位。还有一个常见情况image image.unsqueeze(0)在Dataset里给图片加了batch维度3张图就变成[3, 1, 3, 224, 224]。正确做法是单张图放进Dataset返回时保持[C, H, W]batch维度由DataLoader自动拼接。7.2 Loss不下降的排查Loss不下降是训练中最磨人的问题。我的排查顺序先用少量数据比如一个batch拟合试试。如果模型连几个样本都学不进去那大概率是代码bug而不是调参问题。检查标签是否错位。如果数据加载时图片和标签不对应loss会在某个值附近震荡且准确率接近随机猜二分类约50%。学习率是否过大或过小。loss爆炸或NaN一般是学习率太大loss纹丝不动可能是学习率太小或者梯度消失。检查模型是否进入train模式。忘了model.train()时BN和Dropout行为不对loss可能下不去。最后一层初始化是否合理。一键定位的方法打印前10个batch的outputs和labels看输出的数值范围是否正常。比如二分类输出层如果是0附近的随机值softmax后概率接近0.5说明初始化正常。7.3 显存溢出OOMRuntimeError: CUDA out of memory是新手必经之路。最快的解决办法batch_size从32降到16、8、4每次减半图片尺寸从224降到160或128前提是模型输入尺寸可调关闭验证时的梯度计算确保用torch.no_grad()用AMP混合精度显存占用减少约40%。如果上面全试过还OOM检查是不是代码里存了过多的中间变量比如把每个batch的输出都append到list里没有释放。一个实用技巧在DataLoader里加上persistent_workersTrue配合num_workers0时可以减少数据加载的显存碎片Windows需要multiprocessing_context配合Linux直接用。7.4 过拟合的判断与应对训练集loss持续下降但验证集loss开始上升这就是过拟合。某次实验训练到第8个epoch训练集准确率95%验证集还是80%而且开始往下跌典型的过拟合。应对手段按优先级排序增加数据增强强度先试简单方案加入Dropout在分类层之前插入减小模型容量ResNet50换成ResNet18提前停止训练监控验证集loss增加weight_decay从1e-4调到1e-3还有一个很多人忽略的方案收集更多数据或者用半监督的思路把无标签数据和伪标签加进训练集。7.5 多GPU训练时的数据并行如果你的机器有多张显卡# 在模型定义之后 if torch.cuda.device_count() 1: model nn.DataParallel(model) # 保存时要取module torch.save(model.module.state_dict(), best_model.pth)加载时就要先定义单卡模型结构再load。如果不小心保存了DataParallel的完整state_dict键名带module.前缀加载到单卡模型时会报键值不匹配。解决办法state_dict torch.load(best_model.pth) new_state_dict {k.replace(module., ): v for k, v in state_dict.items()} model.load_state_dict(new_state_dict)7.6 结果不可复现如果每次运行结果都不同先不要怀疑代码。PyTorch的随机性来自三处数据加载顺序、数据增强、模型初始化。要固定结果import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False注意cudnn.deterministic True会让部分卷积层选择确定性算法训练速度会变慢但换来可复现性实验对比时值得。8. 写在最后的经验心得跑完整个流程我最大的体会是PyTorch训练自定义数据集真正麻烦的不是模型代码而是数据。从目录组织、清洗、划分到Dataset实现、增强策略每一步都决定了最终模型的上限。模型只是把数据里的规律提取出来数据质量不行再好的网络也白搭。第二个体会是一定要把验证集和测试集的边界划清楚。验证集用于日常调参测试集只在最终评估时碰一次。这是很多项目最后指标虚高的根源。第三个建议是训练过程的记录习惯尽早养成。每改一个参数记录下对应的超参数、loss曲线、最终准确率。过两个月回头看这些都是宝贵的实验日志。我就是因为早期不做记录重复调了很多次一样的参。最后分享一个调试小技巧不要一开始就在全量数据上训练。先用200张图跑通完整流程每个环节都验证OK了再上全量数据。全量数据训练一次可能几小时用200张图验证流程只要几分钟这笔时间账怎么算都划算。如果你在跑的过程中遇到没提到的报错不要慌。PyTorch的错误信息虽然长但关键信息一般在最后一行。把最后一行报错贴到搜索引擎基本都能找到答案。多跑几次踩过的坑都会变成经验。