简介这是一套基于PyTorch实现的猫狗图像二分类完整项目主要面向深度学习初学者、计算机视觉学习者也适用于课程设计、毕业设计或算法竞赛热身。资源覆盖从数据准备、模型训练到验证测试的全流程修改config.py即可匹配CPU或CUDA环境通过main.py的train/val/test三个命令分别完成训练、验证和测试便于快速复现基础分类效果。包体共647个文件其中604张jpg图片构成主要数据另有11个py源码、7个pth预训练/训练权重、2个csv结果文件以及md/txt使用说明文档压缩包约541.86MB目录结构清晰方便按需取用。已有679人学习下载借助该资源可掌握PyTorch图像分类的标准工程结构也能替换数据集做迁移实验是入门计算机视觉的实用模板。1. 猫狗二分类为什么选PyTorch CNN而不是迁移学习全家桶不少刚接触图像的人拿到猫狗分类第一反应是直接调用预训练模型把最后的全连接层换了就跑。这种做法在比赛里能刷分但工程上有个隐患你永远不知道自己的数据到底哪里出了问题也不知道模型对图片尺寸、通道顺序、均值标准差到底有多敏感。这个资源里给的是一套从零搭建的CNN方案数据、模型、训练脚本和config配置都摊开了适合想搞懂卷积神经网络底层原理的人。我的建议是先用这套代码跑通基线再逐步换成ResNet或者EfficientNet这样后续排查问题会顺手很多。整条链路不长但覆盖了DataSet构建、卷积层设计、训练循环和本地模型加载工作量适中。2. 数据管道从图片文件到DataLoader的完整预处理2.1 数据集的目录组织与标签映射拿到压缩包后先别急着写代码看一眼数据目录。这个项目里的图片是散装的类似cat.185.jpg、dog.132.jpg文件名本身就带了标签。常见做法是把数据集按下面的结构重新组织data/ train/ cat/xxx.jpg dog/xxx.jpg val/ cat/xxx.jpg dog/xxx.jpg如果你手头是散文件可以用一段小脚本做划分。注意这里的train/val比例一般取8:2或9:1不要随机乱分最好保证同一只猫或狗的照片不要既出现在训练集又出现在验证集。文件名前缀是cat.或dog.直接按字符串前缀判断即可。如果用torchvision.datasets.ImageFolder这个工具会自动把子目录名当作类别索引顺序按目录名的字母排列所以cat是0dog是1。如果你自定义映射建议在config.py里写死# config.py CLASS_NAMES [cat, dog] CLASS_TO_IDX {name: i for i, name in enumerate(CLASS_NAMES)} IMG_SIZE 224 BATCH_SIZE 32这里CLASS_TO_IDX的作用是给每个目录名分配数字标签模型输出的两个概率值就对应这两个数字。IMG_SIZE决定后续缩放的统一尺寸224是当前主流CNN的语义约定但并不意味着越小越差后面调参数时你可以改成128试试。2.2 自定义Dataset与图像增强策略虽然ImageFolder可以直接用但我更推荐自己写一个Dataset子类这样做有两个好处第一你想看清每一张图被做了什么变换随时可以断点调试第二当数据不是规整目录结构时比如图片和CSV文件混在一起自定义类更容易处理。这个压缩包里还有submission_example.csv和submission.csv前者是提交样例后者是预测结果说明项目本身带有测试集推理的出口。下面是一个典型实现兼容文件名前缀为标签的散图目录import torch from torch.utils.data import Dataset from PIL import Image import os class CatDogDataset(Dataset): def __init__(self, img_dir, transformNone): self.img_dir img_dir self.img_paths [] self.labels [] for fname in os.listdir(img_dir): label 0 if fname.startswith(cat) else 1 self.img_paths.append(os.path.join(img_dir, fname)) self.labels.append(label) self.transform transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): image Image.open(self.img_paths[idx]).convert(RGB) if self.transform: image self.transform(image) label torch.tensor(self.labels[idx], dtypetorch.long) return image, label代码逻辑很直白__len__返回图片总数__getitem__根据索引读图、做变换并返回张量。注意convert(RGB)这一步不能省因为有些图片是灰度图或者带了透明通道不统一转成RGB的话后续torchvision.transforms处理时会报错。标签用torch.long类型是因为nn.CrossEntropyLoss要求目标为长整型索引。这里的transform参数是数据增强和归一化的组合常见做法是from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(10), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])2.3 DataLoader参数对训练速度与收敛的影响构造好数据集后用DataLoader把它包起来。这里有几个参数值得细说from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_sizeconfig.BATCH_SIZE, shuffleTrue, num_workers4, pin_memoryTrue )shuffleTrue只应在训练集上开验证集和测试集必须保持False。num_workers是读取子进程数Windows上设为0最稳妥Linux和Mac可以设成CPU核数的一半。pin_memory在GPU训练时能加速主机到设备的传输CPU训练时开不开无所谓。注意验证集不需要数据增强只做缩放和归一化否则验证指标会忽高忽低。表2-1 常见DataLoader参数推荐值参数训练集验证/测试集说明batch_size3264显存不够就减半shuffleTrueFalse打乱样本顺序避免周期性num_workers42进程数过高会拖慢主线程pin_memoryTrueTrueGPU训练时提升传输效率这里要提一个坑如果你把验证集也开了shuffleTrue每次验证时数据顺序都不一样如果模型还在训练阶段你会发现验证loss来回跳动误以为模型不收敛。另外Resize和RandomCrop的组合比单纯Resize效果更好常见做法是先把图短边缩放到256再随机裁剪224区域相当于做一个弱数据增强。这个项目的config.py里只有IMG_SIZE没有单独的裁剪尺寸你可以自行扩展。3. CNN核心组件卷积、池化与全连接的PyTorch实现3.1 为什么用两层卷积加全连接而不是纯全连接小数据集上全连接网络也能过拟合但泛化能力差。卷积层的价值体现在两点局部连接和权值共享。猫耳朵的纹理、狗鼻子的轮廓这些特征在图像任意位置都可能出现用同一个卷积核去滑一遍全图参数数量比全连接少几个数量级而且对平移有天然容忍度。这个项目选择从零写CNN而不是直接调用ResNet就是为了让你看清每一层张量尺寸是怎么变的。常见做法是采用「卷积BNReLU池化」的堆叠结构。BN的作用是让每层输入分布稳定下来训练时能容忍更大的学习率推理时用的是训练阶段累积的均值方差所以推理前必须调用model.eval()。ReLU相比sigmoid能缓解梯度消失计算也更快。池化层用最大池化它在每个2x2窗口里取最大值相当于把空间尺寸减半同时保留相对显著的特征。下面是这个项目里核心网络结构的一种实现我基于main.py里封装的模块整理了一下import torch.nn as nn class CatDogCNN(nn.Module): def __init__(self, num_classes2): super(CatDogCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(128, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x这段代码里输入是[B, 3, 224, 224]的RGB张量B是batch size。经过第一次卷积后通道数变成32padding1保证输出尺寸不变最大池化后变成112。四轮之后特征图尺寸是[B, 128, 14, 14]。AdaptiveAvgPool2d((1,1))不管前面特图多大都全局平均池化成[B, 128, 1, 1]这层很关键它让网络可以接受任意分辨率输入。3.2 参数量计算与感受野分析表3-1 各层张量尺寸与参数量网络层输出尺寸单层参数量Conv2d(3-32, 3x3)[B, 32, 224, 224]3323*332 896Conv2d(32-64, 3x3)[B, 64, 112, 112]32643*364 18496Conv2d(64-128, 3x3)[B, 128, 56, 56]641283*3128 73856Conv2d(128-128, 3x3)[B, 128, 28, 28]1281283*3128 147584Linear(128-256)[B, 256]128*256256 33024Linear(256-2)[B, 2]256*22 514总参数量大约28万对于一个二分类任务来说绰绰有余。如果用全连接直接处理224x224的图片第一层就会产生百万级参数很容易过拟合。感受野方面两层3x3卷积堆叠等价于一层5x5卷积但参数量只有后者的18/25而且中间有ReLU增加非线性。如果你把卷积核改成5x5会发现训练速度明显变慢但准确率不见得提升所以小数据集上3x3更划算。3.3 前向传播中容易忽略的dtype与通道顺序ToTensor()会把PIL图片从HWC格式转为CHW并且把像素值从0~255缩放到0~1这一步很多人会忽略直接把PIL数组喂给网络结果输出概率要么全0要么全1。另外模型默认输入是浮点张量需要注意torch.float32而标签是torch.long。这两类张量混用不会报错但loss函数内部会对类型做隐式转换建议保持代码里显式保持一致。一个小技巧是在__getitem__里返回(image, label)之后先用next(iter(train_loader))看一眼张量形状images, labels next(iter(train_loader)) print(images.shape, labels.shape, images.dtype)如果你看到images.shape是[32, 3, 224, 224]但labels不是torch.long就需要修正__getitem__里的类型转换。这一步排查能省掉很多训练中期的诡异报错。4. 训练闭环config配置、CPU/CUDA切换与train/val/test4.1 config.py里应该放哪些配置项这个项目的玩法是改config.py来适配本机环境main.py只封装模块并打印参数。这意味着你不需要改业务逻辑只需要把配置项调整好。常见配置项包括设备类型、图片尺寸、学习率、训练轮数、权重保存路径等# config.py import torch DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) DATA_DIR data IMG_SIZE 224 BATCH_SIZE 32 EPOCHS 30 LEARNING_RATE 1e-3 MOMENTUM 0.9 WEIGHT_DECAY 1e-4 SAVE_DIR checkpoints注意torch.device(cuda if torch.cuda.is_available() else cpu)这段是安全处理没装CUDA版PyTorch的机器上torch.cuda.is_available()返回False自动退回CPU。CUDA版安装时常见匹配是Python 3.10、PyTorch 2.8.0、CUDA 12.1这一组合但你自己的环境配对以torch.cuda.is_available()输出为准。4.2 main.py的train/val/test三段式设计训练时使用python main.py train验证用python main.py val测试用python main.py test。这说明main.py里用命令行参数区分了三阶段。下面给出一个清晰的最小实现import argparse import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import transforms import config from model import CatDogCNN from dataset import CatDogDataset def get_transform(augmentFalse): if augment: return transforms.Compose([ transforms.Resize((config.IMG_SIZE, config.IMG_SIZE)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) else: return transforms.Compose([ transforms.Resize((config.IMG_SIZE, config.IMG_SIZE)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) def train(): train_ds CatDogDataset(f{config.DATA_DIR}/train, get_transform(True)) val_ds CatDogDataset(f{config.DATA_DIR}/val, get_transform(False)) train_loader DataLoader(train_ds, batch_sizeconfig.BATCH_SIZE, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_sizeconfig.BATCH_SIZE, shuffleFalse, num_workers2) model CatDogCNN(num_classes2).to(config.DEVICE) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lrconfig.LEARNING_RATE, momentumconfig.MOMENTUM, weight_decayconfig.WEIGHT_DECAY) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) best_acc 0.0 for epoch in range(config.EPOCHS): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(config.DEVICE), labels.to(config.DEVICE) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) val_loss, val_acc evaluate(model, val_loader) scheduler.step() epoch_loss running_loss / len(train_ds) print(fEpoch {epoch1}/{config.EPOCHS} loss{epoch_loss:.4f} val_loss{val_loss:.4f} acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), f{config.SAVE_DIR}/best_model.pth) def evaluate(model, loader): model.eval() criterion nn.CrossEntropyLoss() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(config.DEVICE), labels.to(config.DEVICE) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total def test(): test_ds CatDogDataset(f{config.DATA_DIR}/test, get_transform(False)) test_loader DataLoader(test_ds, batch_sizeconfig.BATCH_SIZE, shuffleFalse) model CatDogCNN(num_classes2).to(config.DEVICE) model.load_state_dict(torch.load(f{config.SAVE_DIR}/best_model.pth, map_locationconfig.DEVICE)) _, acc evaluate(model, test_loader) print(fTest accuracy: {acc:.4f}) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(mode, choices[train, val, test], help运行模式) args parser.parse_args() if args.mode train: train() elif args.mode val: val_ds CatDogDataset(f{config.DATA_DIR}/val, get_transform(False)) val_loader DataLoader(val_ds, batch_sizeconfig.BATCH_SIZE, shuffleFalse) model CatDogCNN(num_classes2).to(config.DEVICE) model.load_state_dict(torch.load(f{config.SAVE_DIR}/best_model.pth, map_locationconfig.DEVICE)) _, acc evaluate(model, val_loader) print(fVal accuracy: {acc:.4f}) elif args.mode test: test()这个代码里有几个细节值得展开。optimizer.zero_grad()必须在backward()之前调用否则梯度会累加。loss.backward()算的是当前batch的梯度optimizer.step()更新参数。running_loss累加时使用了images.size(0)这是为了应对最后一个batch不足BATCH_SIZE的情况如果只累加loss.item()最后平均的loss会偏小。参数方面StepLR每10个epoch把学习率乘以0.5目的是在训练后期缩小步长防止在最优解附近震荡。SGD的momentum取0.9是经验值它能加速收敛并抑制局部震荡。weight_decay是L2正则系数1e-4是中小数据集上的常用起始值。表4-1 训练超参数推荐范围参数推荐范围说明learning_rate1e-3 ~ 1e-4Adam可稍大SGD配动量batch_size16 ~ 64影响BN统计量稳定性epochs20 ~ 50早停看验证集accdropout0.3 ~ 0.5只在全连接层用weight_decay1e-5 ~ 1e-3过拟合时加大4.3 显存不足与CPU训练时的应对策略如果你用CPU训练num_workers建议设为0否则读图进程和主进程之间的数据通信会挤占本来就紧张的算力。显存不足最常见的就是OOM运行时报CUDA out of memory。解决办法按优先级排第一是降低BATCH_SIZE到16或8第二是把IMG_SIZE从224降到160这个项目里因为网络有全连接层AdaptiveAvgPool已经让全连接输入尺寸固定所以降分辨率是安全的第三是关闭pin_memory。还有一个容易被忽视的点验证阶段也要调用model.eval()并包裹torch.no_grad()。前者会关闭Dropout和BatchNorm的统计更新后者不计算梯度图显存占用会小很多。如果验证时忘记切到eval模式每次dropout的随机性会让验证准确率忽高忽低看起来像模型没训练好。5. 部署实战加载本地模型并预测单张猫狗图片5.1 加载权重文件的正确姿势训练结束后生成的best_model.pth里只有参数不含网络结构。加载时必须先实例化模型再调用load_state_dict。这里有一个环境迁移的坑如果训练时用的GPU测试机器上只有CPUtorch.load需要显式指定map_locationimport torch import torch.nn.functional as F from model import CatDogCNN from torchvision import transforms from PIL import Image def predict_image(image_path, weights_pathcheckpoints/best_model.pth): device torch.device(cuda if torch.cuda.is_available() else cpu) model CatDogCNN(num_classes2) state_dict torch.load(weights_path, map_locationdevice) model.load_state_dict(state_dict) model.to(device) model.eval() transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) input_tensor transform(image).unsqueeze(0).to(device) with torch.no_grad(): logits model(input_tensor) probs F.softmax(logits, dim1) predicted_idx torch.argmax(probs, dim1).item() cat_prob probs[0, 0].item() dog_prob probs[0, 1].item() label cat if predicted_idx 0 else dog print(f{image_path}: {label} (cat{cat_prob:.3f}, dog{dog_prob:.3f})) return label, probsunsqueeze(0)很关键因为模型期望的输入是四维张量[B,C,H,W]单张图读进来只有[C,H,W]需要手动加一个batch维度。softmax把logits转成两个类别的概率两个值加起来等于1。测试一个文件时输出0.999/0.001这种结果大概率是置信度偏高这是二分类模型的正常现象不代表一定正确。5.2 提升推理稳定性的进阶技巧第一批量预测时不要逐张调用predict_image而是把所有图片都预处理成张量组成一个[N, 3, H, W]的batch一次性forward这样能充分利用GPU并行能力。第二如果预测结果总是偏向某一类检查一下验证集里两类的样本数是否均衡不均衡时建议在测试时对logits做类频率惩罚。第三如果服务端有多个线程同时推理要确保模型实例是共享的不要在每次请求时重新加载权重否则磁盘IO会成为瓶颈。另一个实用的点是保存权重时保留config.py里用到的IMG_SIZE和归一化参数因为加载模型时必须用与训练完全一致的预处理。常见做法是把这些参数写进权重的附加字段但更简单的是在config.py里固化只从一处取值避免多个脚本各自维护一份参数。如果你要把模型部署到服务端可以把CatDogCNN和预处理封装为一个类对外只暴露predict(image_path)接口内部隐藏设备切换和归一化细节。本文还有配套的精品资源点击获取
PyTorch猫狗分类实战：从零构建CNN训练到部署
简介这是一套基于PyTorch实现的猫狗图像二分类完整项目主要面向深度学习初学者、计算机视觉学习者也适用于课程设计、毕业设计或算法竞赛热身。资源覆盖从数据准备、模型训练到验证测试的全流程修改config.py即可匹配CPU或CUDA环境通过main.py的train/val/test三个命令分别完成训练、验证和测试便于快速复现基础分类效果。包体共647个文件其中604张jpg图片构成主要数据另有11个py源码、7个pth预训练/训练权重、2个csv结果文件以及md/txt使用说明文档压缩包约541.86MB目录结构清晰方便按需取用。已有679人学习下载借助该资源可掌握PyTorch图像分类的标准工程结构也能替换数据集做迁移实验是入门计算机视觉的实用模板。1. 猫狗二分类为什么选PyTorch CNN而不是迁移学习全家桶不少刚接触图像的人拿到猫狗分类第一反应是直接调用预训练模型把最后的全连接层换了就跑。这种做法在比赛里能刷分但工程上有个隐患你永远不知道自己的数据到底哪里出了问题也不知道模型对图片尺寸、通道顺序、均值标准差到底有多敏感。这个资源里给的是一套从零搭建的CNN方案数据、模型、训练脚本和config配置都摊开了适合想搞懂卷积神经网络底层原理的人。我的建议是先用这套代码跑通基线再逐步换成ResNet或者EfficientNet这样后续排查问题会顺手很多。整条链路不长但覆盖了DataSet构建、卷积层设计、训练循环和本地模型加载工作量适中。2. 数据管道从图片文件到DataLoader的完整预处理2.1 数据集的目录组织与标签映射拿到压缩包后先别急着写代码看一眼数据目录。这个项目里的图片是散装的类似cat.185.jpg、dog.132.jpg文件名本身就带了标签。常见做法是把数据集按下面的结构重新组织data/ train/ cat/xxx.jpg dog/xxx.jpg val/ cat/xxx.jpg dog/xxx.jpg如果你手头是散文件可以用一段小脚本做划分。注意这里的train/val比例一般取8:2或9:1不要随机乱分最好保证同一只猫或狗的照片不要既出现在训练集又出现在验证集。文件名前缀是cat.或dog.直接按字符串前缀判断即可。如果用torchvision.datasets.ImageFolder这个工具会自动把子目录名当作类别索引顺序按目录名的字母排列所以cat是0dog是1。如果你自定义映射建议在config.py里写死# config.py CLASS_NAMES [cat, dog] CLASS_TO_IDX {name: i for i, name in enumerate(CLASS_NAMES)} IMG_SIZE 224 BATCH_SIZE 32这里CLASS_TO_IDX的作用是给每个目录名分配数字标签模型输出的两个概率值就对应这两个数字。IMG_SIZE决定后续缩放的统一尺寸224是当前主流CNN的语义约定但并不意味着越小越差后面调参数时你可以改成128试试。2.2 自定义Dataset与图像增强策略虽然ImageFolder可以直接用但我更推荐自己写一个Dataset子类这样做有两个好处第一你想看清每一张图被做了什么变换随时可以断点调试第二当数据不是规整目录结构时比如图片和CSV文件混在一起自定义类更容易处理。这个压缩包里还有submission_example.csv和submission.csv前者是提交样例后者是预测结果说明项目本身带有测试集推理的出口。下面是一个典型实现兼容文件名前缀为标签的散图目录import torch from torch.utils.data import Dataset from PIL import Image import os class CatDogDataset(Dataset): def __init__(self, img_dir, transformNone): self.img_dir img_dir self.img_paths [] self.labels [] for fname in os.listdir(img_dir): label 0 if fname.startswith(cat) else 1 self.img_paths.append(os.path.join(img_dir, fname)) self.labels.append(label) self.transform transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): image Image.open(self.img_paths[idx]).convert(RGB) if self.transform: image self.transform(image) label torch.tensor(self.labels[idx], dtypetorch.long) return image, label代码逻辑很直白__len__返回图片总数__getitem__根据索引读图、做变换并返回张量。注意convert(RGB)这一步不能省因为有些图片是灰度图或者带了透明通道不统一转成RGB的话后续torchvision.transforms处理时会报错。标签用torch.long类型是因为nn.CrossEntropyLoss要求目标为长整型索引。这里的transform参数是数据增强和归一化的组合常见做法是from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(10), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])2.3 DataLoader参数对训练速度与收敛的影响构造好数据集后用DataLoader把它包起来。这里有几个参数值得细说from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_sizeconfig.BATCH_SIZE, shuffleTrue, num_workers4, pin_memoryTrue )shuffleTrue只应在训练集上开验证集和测试集必须保持False。num_workers是读取子进程数Windows上设为0最稳妥Linux和Mac可以设成CPU核数的一半。pin_memory在GPU训练时能加速主机到设备的传输CPU训练时开不开无所谓。注意验证集不需要数据增强只做缩放和归一化否则验证指标会忽高忽低。表2-1 常见DataLoader参数推荐值参数训练集验证/测试集说明batch_size3264显存不够就减半shuffleTrueFalse打乱样本顺序避免周期性num_workers42进程数过高会拖慢主线程pin_memoryTrueTrueGPU训练时提升传输效率这里要提一个坑如果你把验证集也开了shuffleTrue每次验证时数据顺序都不一样如果模型还在训练阶段你会发现验证loss来回跳动误以为模型不收敛。另外Resize和RandomCrop的组合比单纯Resize效果更好常见做法是先把图短边缩放到256再随机裁剪224区域相当于做一个弱数据增强。这个项目的config.py里只有IMG_SIZE没有单独的裁剪尺寸你可以自行扩展。3. CNN核心组件卷积、池化与全连接的PyTorch实现3.1 为什么用两层卷积加全连接而不是纯全连接小数据集上全连接网络也能过拟合但泛化能力差。卷积层的价值体现在两点局部连接和权值共享。猫耳朵的纹理、狗鼻子的轮廓这些特征在图像任意位置都可能出现用同一个卷积核去滑一遍全图参数数量比全连接少几个数量级而且对平移有天然容忍度。这个项目选择从零写CNN而不是直接调用ResNet就是为了让你看清每一层张量尺寸是怎么变的。常见做法是采用「卷积BNReLU池化」的堆叠结构。BN的作用是让每层输入分布稳定下来训练时能容忍更大的学习率推理时用的是训练阶段累积的均值方差所以推理前必须调用model.eval()。ReLU相比sigmoid能缓解梯度消失计算也更快。池化层用最大池化它在每个2x2窗口里取最大值相当于把空间尺寸减半同时保留相对显著的特征。下面是这个项目里核心网络结构的一种实现我基于main.py里封装的模块整理了一下import torch.nn as nn class CatDogCNN(nn.Module): def __init__(self, num_classes2): super(CatDogCNN, self).__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(128, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x这段代码里输入是[B, 3, 224, 224]的RGB张量B是batch size。经过第一次卷积后通道数变成32padding1保证输出尺寸不变最大池化后变成112。四轮之后特征图尺寸是[B, 128, 14, 14]。AdaptiveAvgPool2d((1,1))不管前面特图多大都全局平均池化成[B, 128, 1, 1]这层很关键它让网络可以接受任意分辨率输入。3.2 参数量计算与感受野分析表3-1 各层张量尺寸与参数量网络层输出尺寸单层参数量Conv2d(3-32, 3x3)[B, 32, 224, 224]3323*332 896Conv2d(32-64, 3x3)[B, 64, 112, 112]32643*364 18496Conv2d(64-128, 3x3)[B, 128, 56, 56]641283*3128 73856Conv2d(128-128, 3x3)[B, 128, 28, 28]1281283*3128 147584Linear(128-256)[B, 256]128*256256 33024Linear(256-2)[B, 2]256*22 514总参数量大约28万对于一个二分类任务来说绰绰有余。如果用全连接直接处理224x224的图片第一层就会产生百万级参数很容易过拟合。感受野方面两层3x3卷积堆叠等价于一层5x5卷积但参数量只有后者的18/25而且中间有ReLU增加非线性。如果你把卷积核改成5x5会发现训练速度明显变慢但准确率不见得提升所以小数据集上3x3更划算。3.3 前向传播中容易忽略的dtype与通道顺序ToTensor()会把PIL图片从HWC格式转为CHW并且把像素值从0~255缩放到0~1这一步很多人会忽略直接把PIL数组喂给网络结果输出概率要么全0要么全1。另外模型默认输入是浮点张量需要注意torch.float32而标签是torch.long。这两类张量混用不会报错但loss函数内部会对类型做隐式转换建议保持代码里显式保持一致。一个小技巧是在__getitem__里返回(image, label)之后先用next(iter(train_loader))看一眼张量形状images, labels next(iter(train_loader)) print(images.shape, labels.shape, images.dtype)如果你看到images.shape是[32, 3, 224, 224]但labels不是torch.long就需要修正__getitem__里的类型转换。这一步排查能省掉很多训练中期的诡异报错。4. 训练闭环config配置、CPU/CUDA切换与train/val/test4.1 config.py里应该放哪些配置项这个项目的玩法是改config.py来适配本机环境main.py只封装模块并打印参数。这意味着你不需要改业务逻辑只需要把配置项调整好。常见配置项包括设备类型、图片尺寸、学习率、训练轮数、权重保存路径等# config.py import torch DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) DATA_DIR data IMG_SIZE 224 BATCH_SIZE 32 EPOCHS 30 LEARNING_RATE 1e-3 MOMENTUM 0.9 WEIGHT_DECAY 1e-4 SAVE_DIR checkpoints注意torch.device(cuda if torch.cuda.is_available() else cpu)这段是安全处理没装CUDA版PyTorch的机器上torch.cuda.is_available()返回False自动退回CPU。CUDA版安装时常见匹配是Python 3.10、PyTorch 2.8.0、CUDA 12.1这一组合但你自己的环境配对以torch.cuda.is_available()输出为准。4.2 main.py的train/val/test三段式设计训练时使用python main.py train验证用python main.py val测试用python main.py test。这说明main.py里用命令行参数区分了三阶段。下面给出一个清晰的最小实现import argparse import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import transforms import config from model import CatDogCNN from dataset import CatDogDataset def get_transform(augmentFalse): if augment: return transforms.Compose([ transforms.Resize((config.IMG_SIZE, config.IMG_SIZE)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) else: return transforms.Compose([ transforms.Resize((config.IMG_SIZE, config.IMG_SIZE)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) def train(): train_ds CatDogDataset(f{config.DATA_DIR}/train, get_transform(True)) val_ds CatDogDataset(f{config.DATA_DIR}/val, get_transform(False)) train_loader DataLoader(train_ds, batch_sizeconfig.BATCH_SIZE, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_sizeconfig.BATCH_SIZE, shuffleFalse, num_workers2) model CatDogCNN(num_classes2).to(config.DEVICE) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lrconfig.LEARNING_RATE, momentumconfig.MOMENTUM, weight_decayconfig.WEIGHT_DECAY) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) best_acc 0.0 for epoch in range(config.EPOCHS): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(config.DEVICE), labels.to(config.DEVICE) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) val_loss, val_acc evaluate(model, val_loader) scheduler.step() epoch_loss running_loss / len(train_ds) print(fEpoch {epoch1}/{config.EPOCHS} loss{epoch_loss:.4f} val_loss{val_loss:.4f} acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), f{config.SAVE_DIR}/best_model.pth) def evaluate(model, loader): model.eval() criterion nn.CrossEntropyLoss() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(config.DEVICE), labels.to(config.DEVICE) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total def test(): test_ds CatDogDataset(f{config.DATA_DIR}/test, get_transform(False)) test_loader DataLoader(test_ds, batch_sizeconfig.BATCH_SIZE, shuffleFalse) model CatDogCNN(num_classes2).to(config.DEVICE) model.load_state_dict(torch.load(f{config.SAVE_DIR}/best_model.pth, map_locationconfig.DEVICE)) _, acc evaluate(model, test_loader) print(fTest accuracy: {acc:.4f}) if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(mode, choices[train, val, test], help运行模式) args parser.parse_args() if args.mode train: train() elif args.mode val: val_ds CatDogDataset(f{config.DATA_DIR}/val, get_transform(False)) val_loader DataLoader(val_ds, batch_sizeconfig.BATCH_SIZE, shuffleFalse) model CatDogCNN(num_classes2).to(config.DEVICE) model.load_state_dict(torch.load(f{config.SAVE_DIR}/best_model.pth, map_locationconfig.DEVICE)) _, acc evaluate(model, val_loader) print(fVal accuracy: {acc:.4f}) elif args.mode test: test()这个代码里有几个细节值得展开。optimizer.zero_grad()必须在backward()之前调用否则梯度会累加。loss.backward()算的是当前batch的梯度optimizer.step()更新参数。running_loss累加时使用了images.size(0)这是为了应对最后一个batch不足BATCH_SIZE的情况如果只累加loss.item()最后平均的loss会偏小。参数方面StepLR每10个epoch把学习率乘以0.5目的是在训练后期缩小步长防止在最优解附近震荡。SGD的momentum取0.9是经验值它能加速收敛并抑制局部震荡。weight_decay是L2正则系数1e-4是中小数据集上的常用起始值。表4-1 训练超参数推荐范围参数推荐范围说明learning_rate1e-3 ~ 1e-4Adam可稍大SGD配动量batch_size16 ~ 64影响BN统计量稳定性epochs20 ~ 50早停看验证集accdropout0.3 ~ 0.5只在全连接层用weight_decay1e-5 ~ 1e-3过拟合时加大4.3 显存不足与CPU训练时的应对策略如果你用CPU训练num_workers建议设为0否则读图进程和主进程之间的数据通信会挤占本来就紧张的算力。显存不足最常见的就是OOM运行时报CUDA out of memory。解决办法按优先级排第一是降低BATCH_SIZE到16或8第二是把IMG_SIZE从224降到160这个项目里因为网络有全连接层AdaptiveAvgPool已经让全连接输入尺寸固定所以降分辨率是安全的第三是关闭pin_memory。还有一个容易被忽视的点验证阶段也要调用model.eval()并包裹torch.no_grad()。前者会关闭Dropout和BatchNorm的统计更新后者不计算梯度图显存占用会小很多。如果验证时忘记切到eval模式每次dropout的随机性会让验证准确率忽高忽低看起来像模型没训练好。5. 部署实战加载本地模型并预测单张猫狗图片5.1 加载权重文件的正确姿势训练结束后生成的best_model.pth里只有参数不含网络结构。加载时必须先实例化模型再调用load_state_dict。这里有一个环境迁移的坑如果训练时用的GPU测试机器上只有CPUtorch.load需要显式指定map_locationimport torch import torch.nn.functional as F from model import CatDogCNN from torchvision import transforms from PIL import Image def predict_image(image_path, weights_pathcheckpoints/best_model.pth): device torch.device(cuda if torch.cuda.is_available() else cpu) model CatDogCNN(num_classes2) state_dict torch.load(weights_path, map_locationdevice) model.load_state_dict(state_dict) model.to(device) model.eval() transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) input_tensor transform(image).unsqueeze(0).to(device) with torch.no_grad(): logits model(input_tensor) probs F.softmax(logits, dim1) predicted_idx torch.argmax(probs, dim1).item() cat_prob probs[0, 0].item() dog_prob probs[0, 1].item() label cat if predicted_idx 0 else dog print(f{image_path}: {label} (cat{cat_prob:.3f}, dog{dog_prob:.3f})) return label, probsunsqueeze(0)很关键因为模型期望的输入是四维张量[B,C,H,W]单张图读进来只有[C,H,W]需要手动加一个batch维度。softmax把logits转成两个类别的概率两个值加起来等于1。测试一个文件时输出0.999/0.001这种结果大概率是置信度偏高这是二分类模型的正常现象不代表一定正确。5.2 提升推理稳定性的进阶技巧第一批量预测时不要逐张调用predict_image而是把所有图片都预处理成张量组成一个[N, 3, H, W]的batch一次性forward这样能充分利用GPU并行能力。第二如果预测结果总是偏向某一类检查一下验证集里两类的样本数是否均衡不均衡时建议在测试时对logits做类频率惩罚。第三如果服务端有多个线程同时推理要确保模型实例是共享的不要在每次请求时重新加载权重否则磁盘IO会成为瓶颈。另一个实用的点是保存权重时保留config.py里用到的IMG_SIZE和归一化参数因为加载模型时必须用与训练完全一致的预处理。常见做法是把这些参数写进权重的附加字段但更简单的是在config.py里固化只从一处取值避免多个脚本各自维护一份参数。如果你要把模型部署到服务端可以把CatDogCNN和预处理封装为一个类对外只暴露predict(image_path)接口内部隐藏设备切换和归一化细节。本文还有配套的精品资源点击获取
