做图像分类这个方向我始终觉得最难的不是那些花哨的网络结构而是怎么把一个看似简单的二分类任务从头到尾跑通并且跑到满意。很多人拿到一批图片数据第一反应就是上预训练模型、调超参结果模型还没写两行先被数据目录、尺寸不一致、标签格式这些基础问题折腾到放弃。这篇内容我把 PyTorch 里的图像二分类从数据组织、预处理、模型设计、训练循环、评估到单张图片预测完整拆一遍每一个环节都给出可以直接抄走的思路与代码。如果你正在准备自己的第一个图像分类项目或者你只需要在两类图片之间做一个判断那这篇分享就是为你准备的。它不会让你立刻调出极致效果但能帮你把标准流程跑通并且知道每一步为什么这么做。我还会把训练过程中遇到过的常见坑和排查思路一并写出来方便你遇到问题时快速定位。1. 图像二分类项目整体拆解1.1 二分类的本质与五段流程所谓图像二分类就是给模型一张图片让它输出“属于类别 A”还是“属于类别 B”。听起来很简单但要把这个任务做扎实背后至少涉及五段流程数据准备、数据预处理、模型设计、训练评估、推理部署。这几段并不是顺序走完就算结束实际项目里经常需要在任意一段反复折返。为什么强调拆解因为很多新手遇到问题习惯性“到处调参”最后根本不知道是数据坏了还是模型坏了。把这个流程拆开之后你可以先确认数据能正确读到再确认模型能跑通一次 forward然后才去关心训练效果。这样每一步都有明确的验收标准排查起来会快很多。二分类本身也是很多更复杂视觉任务的起点把这套流程吃透之后多分类、细粒度识别、目标检测里的数据加载和训练思路都能顺势迁移。1.2 为什么用 PyTorch 而不是别的工具我自己的体验是PyTorch 最让人舒服的地方在于动态图和调试的自由度。模型定义、训练循环都跟写普通 Python 代码一样print 可以打在任意一行断点想加就加。对于图像二分类这种中小型任务这种灵活性比所谓的“性能上限”更重要。另外生态是真的方便。torchvision里自带常用的网络结构、预训练权重和基础数据集读取工具torch.utils.data也把 DataLoader、Sampler 这些组件准备得很完整。就算你不会写底层的东西也可以先把整个流程跑起来再一点点深入。社区资料也很多搜到的问题基本能直接参考。这套组合在图像二分类上的影响范围不只是简单的分类器训练像工业质检、医学影像初步筛查、人脸活体判断这类场景底层的图像二分类套路几乎都是同一套逻辑。2. 数据准备与预处理影响最终效果的第一步2.1 目录结构直接决定代码写法用 PyTorch 自带的ImageFolder读取分类数据时目录结构非常关键。最简单的结构是data/ train/ cat/ dog/ val/ cat/ dog/你只需要把每张图片放进对应的类别文件夹ImageFolder会自动按子目录名生成标签。它还会把类别名按字母顺序映射成整数编号比如 cat 对应 0、dog 对应 1。使用方式非常直接from torchvision import datasets train_dataset datasets.ImageFolder(data/train, transformtrain_transform) print(train_dataset.class_to_idx)这样读出来的 dataset每次用下标访问会返回(image_tensor, label)的元组。直接丢给DataLoader之后训练循环里就能用常见的for images, labels in loader取数据。如果你手头的数据是散落在同一个文件夹里的首先需要自己划分。我一般会把每个类别的文件按比例随机复制到 train 和 val 两个目录比例按数据量来通常 8:2 或者 7:3。我习惯在划分的时候顺便打印每个目录的文件数确保没有类被漏掉。2.2 训练集和验证集的 transform 不能一样数据预处理是二分类项目里最容易被忽略、也最能拉开差距的地方。我这里给出一个比较通用的配置from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_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]) ])训练集用RandomResizedCrop(224)相当于随机裁一块 224x224 的区域并缩放到统一大小这是很常见的数据增强。RandomHorizontalFlip()按照 50% 概率随机水平翻转等于免费扩充了一倍数据。验证集和预测阶段不需要随机增强所以使用Resize(256)加CenterCrop(224)保证每张图看到的主要区域都在。Normalize这一步很多人会漏。它的作用是将像素从 0 到 255 的范围转换到模型更适应的分布。上面的 mean/std 是从大规模通用图像数据上统计出来的如果打算使用 torchvision 提供的预训练模型就继续沿用这套参数如果是自己从零训练的小网络也可以自定义但为了省事我还是建议先沿用这套。需要特别提醒的是预测或验证时的预处理必须和训练时保持一致尤其是尺寸和 Normalize 参数。如果训练时输入是 224x224预测时却直接塞一个原始大图进去模型效果大概率会崩。2.3 类别不平衡时先别急着换模型二分类里正负样本经常相差很大比如异常检测任务里负样本可能只占 1%。如果直接用原始数据训练模型很容易把所有图片都判成多数类。最简单的做法有两个一个是给少数类更高采样权重一个是给损失函数加类别权重。使用 PyTorch 原生组件可以用WeightedRandomSamplerfrom torch.utils.data import WeightedRandomSampler # 给每个样本一个权重少数类权重更大 weights [] for image, label in train_dataset.samples: if label 0: weights.append(0.5) else: weights.append(2.0) sampler WeightedRandomSampler(weights, num_sampleslen(train_dataset), replacementTrue) train_loader DataLoader(train_dataset, batch_size32, samplersampler)注意设置了sampler之后就不能再设置shuffleTrue因为 sampler 本身就决定了取数据的顺序。这个方案不用改网络结构在训练初期比较容易上手。如果你不想引入采样器也可以在CrossEntropyLoss里传weight参数直接偏向少数类。3. 模型设计先跑通再追求高性能3.1 自己写一个简单的 CNN 网络对于二分类你不需要一开始就上几十层的网络。我建议先用一个小型卷积网络把整个流程跑通等能快速迭代了再换成预训练模型。下面这个网络比较简单但结构完整import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d((1, 1)) ) self.fc nn.Linear(128, num_classes) def forward(self, x): x self.features(x) x x.flatten(1) return self.fc(x)卷积层使用kernel_size3, padding1是为了让特征图尺寸不缩小太多配合MaxPool2d(2)逐步降低分辨率。最后用AdaptiveAvgPool2d((1, 1))把任意尺寸的特征图池化成 1x1这样全连接层的输入维度就固定是 128不需要手算。这个网络虽然简单但在普通二分类数据上经常能给出一个不错的基线。3.2 最后一层到底用一个输出还是两个输出二分类任务最后有两种常见做法很多新手容易混淆。第一种是最后一层输出 2 个神经元配合CrossEntropyLoss。第二种是最后一层只输出 1 个神经元配合BCEWithLogitsLoss。两者都能用但我个人更推荐第一种因为它和常规的多分类写法一致以后要是扩展成多分类直接改num_classes就行。输出方案模型最后一层损失函数预测方式两神经元Linear(..., 2)CrossEntropyLosssoftmax 后取类别1概率或 argmax单神经元Linear(..., 1)BCEWithLogitsLosssigmoid 后大于 0.5 判为类别1如果你用BCEWithLogitsLoss不要在模型外面再手动做 sigmoid因为这个损失函数本身是在数值稳定的 logits 上计算的。手动加 sigmoid 后再算损失虽然也能跑但梯度在极端情况下容易出现问题。预测时倒是可以用torch.sigmoid把 logits 转成 0 到 1 之间的概率。3.3 使用 torchvision 预训练模型做迁移学习当数据量不大时从零训练一个 CNN 很容易过拟合。更稳妥的做法是加载已经在通用图片数据上预训练过的模型把最后一层全连接换掉让模型复用已经学到的浅层特征这在数据量有限时非常管用。from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) num_ftrs model.fc.in_features model.fc nn.Linear(num_ftrs, 2)如果你使用的是旧版 torchvisionweights参数可能不存在那时候可以写pretrainedTrue。换成预训练模型之后训练策略也要相应调整。比较稳的流程是先把model.fc之前的所有层冻结只训练新的全连接层跑几个 epoch 之后再解冻全部参数用更小的学习率微调。for param in model.parameters(): param.requires_grad False # 替换后的 fc 层默认 requires_gradTrue所以只更新 fc optimizer torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3)这个做法既省显存又能避免随机初始化的分类层在初期把预训练特征破坏掉。等分类头表现稳定后再允许骨干网络微调效果通常更好。4. 训练配置与核心循环把模型真正训起来4.1 超参数学习率、批次大小和训练轮数学习率是最需要盯的超参数之一。使用 Adam 或 AdamW 优化器时我的默认选择是从零训练小 CNN 用 1e-3微调预训练模型用 1e-4。这不是绝对的但作为起点很安全。batch size 一般用 16、32、64 这种 2 的幂具体多少由显存和图片尺寸决定。训练轮数不要一开始就写死 100 轮最好每隔一段时间看一眼验证集指标用早停来定。为什么迁移学习用更小的学习率因为预训练权重已经处于一个比较好的解附近步长太大会把学到的特征覆盖掉。而小 CNN 是一张白纸稍微大一点的学习率能让它更快收敛。还有一个比较实用的经验如果训练 loss 下降很慢不要急着调网络结构先调学习率。4.2 一个可以直接用的训练循环训练循环的关键点有三个model.train()、model.eval()、optimizer.zero_grad()。我封装了两个函数一个负责训练一个负责验证。import torch def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in loader: images images.to(device) labels labels.to(device) optimizer.zero_grad() logits model(images) loss criterion(logits, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) pred logits.argmax(dim1) correct (pred labels).sum().item() total images.size(0) return total_loss / total, correct / total torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() total_loss 0.0 correct 0 total 0 for images, labels in loader: images images.to(device) labels labels.to(device) logits model(images) loss criterion(logits, labels) total_loss loss.item() * images.size(0) pred logits.argmax(dim1) correct (pred labels).sum().item() total images.size(0) return total_loss / total, correct / total训练时一定要调用model.train()验证时一定要调用model.eval()。这一点看起来基础但影响很大BatchNorm和Dropout在两种模式下的行为完全不同。如果验证时忘了切到 eval验证集上的指标会忽高忽低特别容易被误导。主循环里记录训练和验证的 loss、accuracy并保存验证集上性能最好的模型权重best_acc 0.0 for epoch in range(epochs): train_loss, train_acc train_one_epoch( model, train_loader, criterion, optimizer, device ) val_loss, val_acc evaluate(model, val_loader, criterion, device) print(fEpoch {epoch1}: train_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 torch.save(model.state_dict(), best_model.pth)保存state_dict而不是整个模型是因为这样更轻便也更容易在后续加载时把结构定义和权重分开管理。如果你中间需要恢复训练可以把优化器 state 也一起存下来但只做推理的话保存state_dict就够了。4.3 评估指标不能只盯着准确率准确率看着直观但在二分类任务里经常不够用。比如正样本占 99% 的数据集模型全部预测成负样本也有 99% 准确率但显然没解决问题。所以我会在验证阶段额外记录混淆矩阵、精确率、召回率和 AUC。在 PyTorch 中收集概率和标签然后交给sklearn计算from sklearn.metrics import confusion_matrix, roc_auc_score y_true [] y_score [] def collect_preds(model, loader, device): model.eval() with torch.no_grad(): for images, labels in loader: images images.to(device) logits model(images) prob torch.softmax(logits, dim1)[:, 1] y_true.extend(labels.cpu().tolist()) y_score.extend(prob.cpu().tolist()) # 调用 collect_preds 后计算 auc roc_auc_score(y_true, y_score) cm confusion_matrix(y_true, [1 if s 0.5 else 0 for s in y_score])AUC 的好处是不受分类阈值影响它直接衡量模型把正样本排在负样本前面的能力。如果业务上“宁可误报也不能漏报”那你就需要关注召回率而不仅仅是 AUC 或准确率。把这些指标记录下来比只记录 loss 有价值得多。5. 预测与部署训练完之后怎么用起来5.1 单张图片预测的完整代码模型训练完接下来就是输入一张新图片得到预测结果。这一步要非常小心预处理的一致性。下面是单张图片预测的示例from PIL import Image def predict_one(model, img_path, transform, device): model.eval() img Image.open(img_path).convert(RGB) x transform(img) x x.unsqueeze(0).to(device) with torch.no_grad(): logits model(x) prob torch.softmax(logits, dim1) return prob[0, 1].item()这里有几个细节很容易出错。第一convert(RGB)是为了统一图片通道数防止某些图片是灰度图或带透明通道导致模型输入尺寸对不上。第二transform必须使用和验证集一样的val_transform不能带随机翻转或随机裁剪。第三图片经过 transform 后得到的是[C, H, W]的张量模型期望输入[N, C, H, W]所以一定要unsqueeze(0)补上 batch 维。5.2 训练与预测的预处理为什么必须一致我见过不少人在训练阶段用RandomResizedCrop预测阶段却直接把图片缩放到 224x224 塞进去结果准确率明显下降。原因不在于随机裁剪本身而在于模型在训练时见过的是带随机裁剪变化的输入预测时却要面对全新的尺寸分布。验证集用Resize(256) CenterCrop(224)是为了模拟预测时“只看图片中间区域”的状态。另外一个容易被忽略的是Normalize。训练用的 mean/std 如果和验证、预测不一致模型输出概率会全接近 0.5 甚至完全乱猜。这类问题通常不会报错只会让你觉得模型效果很差排查起来很痛苦。我自己在调灰度图输入时就踩过一次类似的坑所以现在每次写预测函数第一件事就是检查 transform 是不是和验证集完全一致。5.3 批量推理与设备管理如果你需要一次预测很多张图片可以先把若干单张图转成张量然后用torch.cat合成一个 batch一次性送入模型。这样比一张张调用更高效在 GPU 上尤其明显。def predict_batch(model, img_paths, transform, device): model.eval() tensors [] for path in img_paths: img Image.open(path).convert(RGB) tensors.append(transform(img).unsqueeze(0)) batch torch.cat(tensors, dim0).to(device) with torch.no_grad(): logits model(batch) probs torch.softmax(logits, dim1)[:, 1] return probs.cpu().numpy()使用 GPU 时注意模型和数据都要放到同一个 device 上。通常写一个device torch.device(cuda if torch.cuda.is_available() else cpu)然后统一调用.to(device)这样代码在 CPU 和 GPU 上都能跑。训练期间如果显存不够优先减小 batch size而不是换更大的模型。6. 常见问题与排查技巧6.1 训练 loss 不降或变成 nan如果模型跑了很多个 epoch 却不收敛我先按下面几步排查。先确认标签是否真的是从 0 开始的整数因为CrossEntropyLoss的 target 不接受字符串或负数。再检查模型输出神经元数量是否和类别数一致。然后看学习率学习率太大会导致 loss 跳过最优点表现为 loss 剧烈抖动或者直接变成 nan学习率太小又会收敛极慢。最后一步是检查预处理图片没有归一化或者混入了损坏文件也可能导致训练异常。遇到 loss 变成 nan一个快速定位的小技巧是把 batch size 设为 1然后用同样的流程跑几个 step如果在单样本上就 nan大概率是数据或网络层的问题而不是优化器的问题。6.2 过拟合训练集很好验证集很差这是二分类小数据任务里最常见的情况。原理解释起来很直观模型容量大于数据信息量它就开始背训练样本而不是学泛化规律。我一般优先做三件事增加数据增强、换成预训练模型、用早停。数据增强可以加RandomRotation、ColorJitter只要不会改变图片原本的类别语义就行。预训练模型加冻结策略能显著降低过拟合因为可训练参数变少了。早停则是在验证集指标连续多轮不提升时直接停止训练避免后期纯粹记忆训练集。还有一个容易被低估的细节数据划分时必须保证训练集和验证集来自同分布。比如图片是按时间或者设备分组的如果随机划分不当验证集和训练集有可能出现重叠或者分布偏差会严重干扰评估。6.3 DataLoader 加载数据时总报错新手最容易遇到的两个问题一是 Windows 下设置num_workers0报错二是返回的 label 不是预期的类型。解决办法是Windows 下可以把num_workers先设为 0或者把训练逻辑封装进if __name__ __main__:里label 需要在 Dataset 内部显式转成整数。图片损坏导致的报错也很常见。比如有些图片本身是坏的PIL 打不开一旦采样到它就会在__getitem__里崩溃。我的处理办法是在预处理阶段写一个简单的脚本遍历所有图片遇到打不开的就把路径打印出来并移到另一个目录。这个步骤虽然枯燥但能节省后面排查的大量时间。6.4 显存不够和训练太慢如果 GPU 显存不够第一步是减小batch_size。也可以适当降低输入分辨率比如把图片从 224x224 改成 160x160但最终效果需要自己试。更重要的是训练循环里一定记得用model.train()和torch.no_grad()验证阶段使用torch.no_grad()可以省掉梯度缓存。如果你的显卡支持自动混合精度可以试试 PyTorch 的torch.cuda.amp训练速度通常有提升显存占用也会下降。但这不是必须的老显卡不一定支持代码也会更复杂一些。训练太慢的时候另一个思路是换更小的模型比如MobileNetV3而不是一味加显卡。6.5 每次训练结果都不一样怎么办深度学习训练有一定的随机性。为了尽量可复现在训练脚本开头固定随机种子import random import numpy as np import torch def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)固定种子之后同一个显存环境下的结果通常会比较接近但受 CUDA 底层算法影响不一定能完全一致。这个问题不需要过度追求我更在意的是把实验配置固定下来比如 transform、学习率、epochs、数据划分版本每次都记录下来。这样下次即使结果有些偏差也能知道是哪一步改动引起的。最后说一个我自己的习惯。每次做图像二分类我不急着调最好的模型而是先拿一个小网络把整个流程跑通精确到训练曲线、验证指标、预测函数都能正常打印结果然后再换预训练模型。这个顺序帮我省掉了大量无效调试。项目到了后期保存模型和日志的习惯比调参本身更值钱。希望这篇分享能让你少踩几个我踩过的坑。