简介面向软件工程、人工智能、通信工程、自动化等计算机相关专业的毕业设计与课程设计这份基于Python的卷积神经网络CNN图像分类系统资源包提供完整源码、训练好的模型和说明文档既能用于项目答辩演示也适合初学者进阶学习网络搭建与调参。压缩包内含22个文件以13个Python源码文件为主覆盖模型定义、训练矩阵、图形界面入口并配有数据集、说明文档、配置及可视化静态文件整体约62KB结构清晰便于二次开发。资源中收录了LeNet-5、AlexNet、GoogLeNet、ResNet等经典网络结构并涉及TensorFlow与PyTorch两种实现思路可帮助对比主流框架下不同网络训练效果。目前已有130人学习下载代码经过运行验证可直接用作毕业设计、课程设计或初期立项演示的可靠参照。1. 用Python卷积神经网络CNN做一个图像分类系统这个项目包能解决什么问题从网盘拖下一个写着“毕业设计 基于Python卷积神经网络CNN的图像分类系统源码模型说明文档全部数据资料.zip”的压缩包第一件事不是急着解压运行而是先想清楚怎么把它变成能写进论文、能现场演示的东西。这类项目在计算机视觉里非常成熟用Python搭一个卷积神经网络CNN对图片做分类猫狗识别、垃圾分拣、森林图像分类都在同一个框架里跑。对毕业设计和刚入门深度学习的人来说它最大的价值是提供一套能跑的模板——数据、模型、训练、验证从头串到尾少踩很多弯路。这篇笔记就按读源码、装环境、处理数据、调参、避坑、验证六步把这个zip包拆开讲清楚。2. 跑通CNN图像分类源码环境、依赖与项目结构都要动手查一遍拿到一个图像分类项目包最忌讳的行为是直接跑python train.py然后盯着屏幕等报错。项目能不能跑通很大程度上取决于你的运行环境跟作者写代码时的环境差了多远。先花十分钟把目录结构、依赖清单、权重文件格式看清楚比盲目调包要节省一整个晚上。2.1 从zip包里的文件布局判断项目是哪个技术栈解压之后先不要运行任何脚本用文件管理器或tree命令把整个目录看一眼。一个典型的毕业设计级CNN图像分类包结构大致是下面这样的project/ ├── train.py # 训练入口负责数据加载、模型训练、保存权重 ├── test.py # 测试/预测入口加载权重并对图片分类 ├── models.py # 网络结构定义可能是自写CNN也可能是预训练模型封装 ├── utils.py # 数据读取、日志、评估等工具函数 ├── requirements.txt # Python依赖清单 ├── best_model.pth # 训练好的模型权重 ├── data/ │ ├── train/ # 按类别分好的训练图片 │ ├── val/ # 验证集 │ └── test/ # 测试集 └── docs/ └── 说明文档.docx/pdf这个布局判断标准很通用图片数据按类别子文件夹存放的说明数据加载大概率用的是torchvision.datasets.ImageFolder能直接看到.pth或.pt结尾的权重文件说明用的框架是PyTorch。文件/目录你该重点看什么常见坑requirements.txttorch、torchvision版本号版本对不上会直接启动失败models.py网络类名、forward里的层顺序换GPU后要改设备变量test.py图像预处理是否与训练一致尺寸不一致会报shape错误best_model.pth保存的是state_dict还是整个模型加载方式写法不同data/类别文件夹命名不要有中文和空格docs/说明文档数据集来源、训练参数、运行步骤很多人把它当装饰其实它最值钱如果压缩包里没有requirements.txt就靠import语句反推依赖。常见的是torch、torchvision、numpy、Pillow、opencv-python、scikit-learn、tqdm、matplotlib这八件套。先花五分钟把这些确认了再动手装环境不然中途发现少一个包前面配的环境全得回滚。2.2 用Python创建虚拟环境装依赖版本与CUDA的坑Python入门阶段最劝退的就是依赖混乱。系统里既有Python 3.8又有3.10全局装了一堆包再往里面堆深度学习依赖迟早把环境搞坏。常见做法是用conda或venv给这个项目单独开一个虚拟环境我一般习惯用conda# 创建一个干净的Python 3.9环境3.8/3.10也常见先看代码里语法兼容性 conda create -n cnn_class python3.9 -y conda activate cnn_class # CPU版和GPU版二选一GPU版需要先确认NVIDIA驱动版本 # 有独立显卡且显存4G的装CUDA 11.8版本没有显卡就装CPU版 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install -r requirements.txt这里有几个参数要解释清楚。cu118表示CUDA 11.8对应的PyTorch版本如果是新一点的显卡驱动也可以用cu121但要注意跟本机驱动匹配。判断方法是在终端执行nvidia-smi看右上角CUDA Version驱动支持的版本只要不低于PyTorch要求的版本就能用。没有NVIDIA显卡的机器直接装CPU版torch即可别硬凑GPU慢一点但能跑。下载慢的话把pip源换成内网镜像比如-i https://pypi.tuna.tsinghua.edu.cn/simple只对本次安装生效不影响全局配置。装完依赖后执行python -c import torch; print(torch.__version__)验证安装结果。输出版本号不代表就能用还要看torch.cuda.is_available()如果是GPU环境返回False大概率是CUDA版本和驱动不匹配需要换对应版本重新装。2.3 先用现成模型做单张预测确认权重文件没坏环境就绪后不要立刻训模型那样既慢又不知道会在哪一步断。先把项目里现成的权重文件跑起来对一张测试图片做预测确认这条链路是通的。这一步本质上是给后面的所有改动找一个参照点。import torch from PIL import Image from torchvision import transforms # 这里的预处理需要和训练时保持一致否则预测结果会明显变差 transform transforms.Compose([ transforms.Resize((224, 224)), # 统一到模型输入尺寸 transforms.ToTensor(), # 转成张量像素值从0-255变为0-1 transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), # ImageNet统计量 ]) # map_locationcpu 保证机器没GPU时也能加载 model torch.load(best_model.pth, map_locationcpu) model.eval() # 切到推理模式关闭Dropout和BN的batch维度统计 img Image.open(data/test/sample.jpg).convert(RGB) inputs transform(img).unsqueeze(0) # 单张图也要加batch维形状变成[1,3,224,224] with torch.no_grad(): logits model(inputs) prob torch.softmax(logits, dim1) # 把输出转成概率分布 print(prob)这中间有四个参数最容易出问题。第一个是map_locationcpu在GPU上训练的权重保存了cuda:0的设备信息CPU机器不写这个参数会报错找不到设备。第二个是model.eval()漏了它带Dropout或BatchNorm的模型预测结果每次都不一样很多人误以为模型坏了其实只是没切模式。第三个是transform里的Normalize它用的均值和标准差不是随便填的ImageNet预训练模型基本都用那四组数字自定义模型如果训练时用的是别的统计量这里要跟着改。第四个是unsqueeze(0)模型期望输入是四维张量[N, C, H, W]单张图片只有三维必须补上batch维度。跑通之后去做一个实验分别预测三张不同类别图片把概率最高的三个类打印出来看是不是符合直觉。如果输出结果和真实标签完全对不上先怀疑预处理不一致再怀疑标签文件class_names.txt的顺序问题最后才怀疑权重文件损坏。权重文件损坏的概率其实很低被预处理坑的概率却非常高。提示如果项目里的test.py已经写得很完整直接用它跑效果更好。自己写预测脚本的用途是验证环境而不是替代项目里的测试代码。3. 把图片数据喂给CNN数据划分、尺寸和归一化的落地实现数据是图像分类项目里最容易被低估的部分。很多同学拿到zip后先看网络结构觉得那个最“像技术”但实际决定模型上限的往往是数据怎么组织、怎么增强、怎么平衡。数据没处理好再深的网络也白搭。3.1 把图片整理成ImageFolder能直接读的目录torchvision.datasets.ImageFolder要求数据按“类别名/图片文件”的层级存放比如data/train/cat/xx.jpg、data/train/dog/xx.jpg。如果压缩包里已经分好了直接用没有分好或者想重新划分用脚本批量处理最省事。写一个按比例划分数据集的脚本import pathlib import random import shutil # 固定随机种子保证每次划分结果一致论文里能写清比例 random.seed(42) src pathlib.Path(data/images) # 原始图片总目录里面是按类别分的子目录 train pathlib.Path(data/train) val pathlib.Path(data/val) test pathlib.Path(data/test) # 每个类别目录下按 8:1:1 划分训练/验证/测试 for cls in src.iterdir(): if not cls.is_dir(): continue # 跳过非目录文件 imgs list(cls.glob(*.jpg)) list(cls.glob(*.jpeg)) list(cls.glob(*.png)) random.shuffle(imgs) n len(imgs) n_train, n_val int(n * 0.8), int(n * 0.1) for i, im in enumerate(imgs): if i n_train: dest train / cls.name elif i n_train n_val: dest val / cls.name else: dest test / cls.name dest.mkdir(parentsTrue, exist_okTrue) shutil.copy(im, dest / im.name) # 用copy保留原始数据便于反悔 print(划分完成)三个参数需要注意。random.seed(42)是保证可复现的关键深度学习项目里玄学问题特别多随机种子固定后至少能让同一份数据每次跑出一样的结果排错时少一个变量。划分比例0.8/0.1/0.1是毕业设计里比较稳妥的默认值数据量在两万张以下时验证集太少准确率数字会忽高忽低建议验证集至少留几百张能说明问题的图。用shutil.copy而不是shutil.move因为移动会破坏原始数据集后续想换比例重新划分还得重新找源文件等于删了自己的后悔药。划分完成后数一数每个类别的图片数量。如果某个类别只有几十张另一个类别有几千张这就是类别不均衡后面训练时要专门处理。另外目录不要起中文名ImageFolder对中文路径的兼容性时好时坏一次报错卡半小时不值得。3.2 变换参数决定模型看到什么Resize、Crop与Normalize数据变换是图像分类里门槛最低但效果最明显的一环。它决定模型训练时看到的是什么尺寸、什么分布、什么颜色范围的图片。训练和验证用的变换通常是分开的训练集多一些随机性验证集保持稳定from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放到224x224 transforms.RandomHorizontalFlip(), # 随机水平翻转相当于数据增强 transforms.ToTensor(), # 像素从0-255转到0-1 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)不是随便选个尺寸224是ImageNet时代定下来的标准输入尺寸绝大多数预训练模型都以它为基准。Resize(256) CenterCrop(224)是验证集的标准组合先把图放大到256再裁224比直接Resize到224保留更多中心区域的细节。RandomHorizontalFlip对大多数自然图片有效但要注意有方向性的数据不能加比如车牌识别、文字方向确定的路标分类水平翻转会把语义弄反。Normalize那四个数字是ImageNet数据集的全局统计量只要用了PyTorch官方预训练模型这套数值基本是通用的。如果是完全自己训练的小模型理论上应该用自己数据集的均值和标准差但实际中直接用ImageNet统计量也不会差太多省一步算一步。这里最容易踩的坑是训练和测试的transform不一致或者测试时忘了加Normalize导致同样的图在训练时和预测时的数值分布完全不同。3.3 类别不均衡时怎么喂数据加权采样与过采样图像分类里类别不均衡是最常见的翻车点。比如森林图像分类项目林地照片可能有八万张火烧迹地只有八百张模型会倾向于把一切预测成林地因为这样准确率也能到90%以上。解决思路有重采样和重加权两条路最常见且省内存的是WeightedRandomSamplerimport torch from torch.utils.data import WeightedRandomSampler # dataset.samples 是 (图片路径, 类别id) 的列表 samples dataset.samples labels [s[1] for s in samples] # 统计每个类别样本数并按倒数赋权少数类权重更高 counts torch.bincount(torch.tensor(labels)) weights 1.0 / counts.float() # 给每个样本配一个权重 sample_weights torch.tensor([weights[l] for l in labels]) sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader torch.utils.data.DataLoader( dataset, batch_size32, samplersampler, num_workers4)两个参数要理解清楚。replacementTrue表示同一张图在同一个epoch里可能被抽到多次这对样本很少的类别是好事相当于隐式过采样。num_sampleslen(sample_weights)指定一个epoch里采样的总次数保持和原始样本量一致这样epoch的概念不用重新定义。这套方案比直接复制少数类图片更省内存因为不需要真的把图片复制多份只是采样时重复选中磁盘上没有冗余数据。不过它有个副作用少数类被反复看到容易过拟合。所以在做加权采样时最好搭配前面的数据增强手段让模型每次看到的少数类样本都不完全一样。训练结束后用原始分布的数据集做验证和测试不要用采样后的数据否则评估出的准确率会虚高。4. 训练一个能交差的CNN分类模型网络结构、学习率与batch size怎么配数据准备好环境跑通接下来才是真正动手训练模型。从零手写CNN、用预训练模型微调、直接用开源权重做迁移学习三条路各有各的适用场景。这一章先把选型逻辑说清楚再给出训练主循环和参数怎么定。4.1 先想清楚选哪条路自写网络、PyTorch迁移学习还是开源模型毕业设计项目包里常见的网络结构有三种手写的小型CNN、基于ResNet等预训练模型的迁移学习、直接套用完整开源模型。很多入门者倾向于认为自己从头写一个CNN才有技术含量其实在几万张图片的数据量下手写模型的上限普遍低于迁移学习。方案参数量收敛速度CPU能否跑合适数据量自写3层CNN小约几十万快能但慢几千张以下微调ResNet18中等约1100万快只需训练最后一层时更快勉强能建议GPU几千到几万张完整训练ResNet50大约2500万慢基本不能几十万张以上如果你的数据量在两千到两万张之间微调ResNet18是最稳的选择。它的预训练权重在ImageNet上见过大量真实图片底层的边缘、纹理特征直接拿来用你只需要把最后一层全连接换成自己类别数训练后面的分类头即可。如果要自己写网络一个能出结果的最小结构是三层卷积池化加一层全连接。通道数逐层翻倍是经验性的常规做法因为越深的层提取的特征越抽象需要更多通道来承载信息。import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), # 输入RGB三通道 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), # 不管输入多大池化成1x1 ) self.classifier nn.Linear(128, num_classes) def forward(self, x): x self.features(x) return self.classifier(x.view(x.size(0), -1))这段代码里有几个参数设计不是随便写的。kernel_size3是卷积核的标准选择奇数核能确保padding后尺寸不变堆两层就能扩大感受野。padding1配合stride1让输出尺寸跟输入一致不用手动计算降维。AdaptiveAvgPool2d(1)是通用结构它把任意尺寸的输入特征图池化成1x1替换了原先要手算的view操作这样网络对输入尺寸的容忍度更高。但说句实在话毕业设计能用迁移学习解决的没必要硬从零训。微调ResNet的效果更稳定、论文里更好解释而且调试成本低一个量级。这不算偷懒工作里也都是这么干的。4.2 训练主循环里的几个关键参数不能乱拍训练脚本是整个项目里含金量最高的一环也是参数最多的一环。下面这段是PyTorch训练一个微调ResNet18的最小主循环import torch from torchvision import models, datasets, transforms device torch.device(cuda if torch.cuda.is_available() else cpu) # 用ImageNet预训练权重只替换最后一层全连接 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc torch.nn.Linear(model.fc.in_features, num_classes) # in_features512 model model.to(device) criterion torch.nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.001, momentum0.9) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.1, patience3) epochs 30 for epoch in range(epochs): model.train() # 切到训练模式 running_loss 0.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) val_loss evaluate(model, val_loader) # 单独写的验证函数 scheduler.step(val_loss) print(fepoch {epoch1}, loss {running_loss/len(train_loader.dataset):.4f})参数选择逻辑要讲透。lr0.001是微调场景下的安全起点比这个再大容易把预训练权重冲坏如果是从零训练小网络可以开到0.01。optimizer用SGD加momentum0.9还是Adam取决于你的数据量。SGD收敛略慢但最终效果通常更好Adam在前期loss降得快但后期容易在最优解附近震荡数据量小的时候用SGD更不容易出过拟合。batch_size受显存制约常见的32和64是默认值调高能加快训练但占用显存线性增长4G显存跑ResNet18时32已经接近上限爆显存会直接报CUDA out of memory。scheduler用ReduceLROnPlateau比固定学习率衰减更省心。它的含义是验证loss连续3个epoch不下降就把学习率乘0.1。这一个参数能解决大量“训练到一半loss降不下去”的问题代价只是多等几个epoch。epoch设为30也是经验值小数据集上ResNet微调20轮左右就能看到收敛趋势不用一开始就写100。提示训练不是跑完30轮就万事大吉。每轮保存一次日志方便后面画曲线也方便判断过拟合从第几个epoch开始。4.3 训练过程中盯住哪些指标loss和准确率训练时盯着四个数看训练集loss、验证集loss、训练集准确率、验证集准确率。只看loss不看准确率会被loss数值欺骗只看准确率不看loss过拟合早期信号容易漏掉。保存模型的最佳时机不是最后一个epoch而是验证集准确率最高的时候best_acc 0.0 for epoch in range(epochs): # 训练和验证代码省略 if val_acc best_acc: best_acc val_acc # 只保存state_dict比保存整个model对象更稳 torch.save(model.state_dict(), best_model.pth)state_dict是模型的权重字典不包含网络结构定义加载时需要先实例化模型再load_state_dict。这看起来麻烦一点但好处是换机器、换Python版本时不受序列化差异影响也避免把设备信息一起打包进去。训练结束后把best_model.pth和最后一轮的权重分开存放论文里报告指标时用best的复现时用last的两条线互相验证。训练日志建议每轮至少打印一次格式固定下来epoch、train_loss、val_loss、val_acc。后期分析时能一眼看出过拟合的起点。另外在训练开始前固定随机种子torch.manual_seed(42)否则两次训练的结果差异会被误认为是代码改动引起的而实际上只是随机性在作怪。5. 训练与部署CNN的常见问题与排查五个易踩的坑训练图像分类模型时大部分时间不是花在写代码上而是花在排查那些看起来莫名其妙的现象上。下面这五个问题我见过太多次每一条都按“现象→原因→解决”写按这个顺序排查能少走很多弯路。5.1 训练卡住不动loss不降与loss为NaN的排查问题一训练了好几轮loss一直在一个常数附近波动几乎没有下降。现象epoch从1跑到10loss徘徊在0.69附近预测结果清一色指向同一个类别准确率稳定在样本比例最大的那个类上。原因这通常是类别不均衡引发的“懒惰解”模型发现把所有图判成同一类就能拿到最低loss于是停在了这个没有意义的解上。另一个常见原因是网络结构里最后一层全连接没有被正确初始化输出分布让梯度计算失效。解决先用混淆矩阵确认模型是否全部预测为多数类如果是给数据加载器加上第3章的WeightedRandomSampler或者给CrossEntropyLoss传入weight参数给少数类更高的loss系数。如果加权后依然不降把学习率从0.001降到0.0001有时只是步子太大跳过了最优区域。问题二loss爆成NaN然后一直保持NaN。现象某次训练过程中loss突然变成nan之后所有轮次都是NaN即使调小学习率再跑也复现。原因学习率过大导致梯度爆炸或者输入图片里存在全黑图、损坏图像素值出现异常再或者Normalize的std设了0导致除零操作。解决先把学习率降到1e-4再试一次。如果还是NaN在DataLoader里加一个过滤逻辑跳过无法解码或尺寸异常的图片文件。最后检查预处理管线确认ToTensor和Normalize确实被执行了而不是被某个条件分支跳过去。修复后不要试图“救”那个NaN模型它已经废了直接从头重跑。5.2 验证指标异常过拟合、模型加载失败和尺寸不一致问题三训练集准确率95%验证集只有60%。现象训练曲线里train_acc一路上升val_acc升到某个点后开始下降两者差距越拉越大。原因这是典型的过拟合。模型把训练样本的特征背下来了尤其当数据量只有几千张、而网络参数量有上千万时这种情况几乎不可避免。解决第一选择是加数据增强RandomHorizontalFlip、RandomRotation(10)、ColorJitter组合起来一张图变出多种形态变相扩大训练集。第二选择是提前停止val_acc不再上升时停止训练用best_model.pth而不是最后一个epoch的权重。第三选择是减小模型容量把ResNet18换成层数更少的网络或者在全连接层后加Dropout(0.5)。问题四换电脑加载模型报Missing key(s) in state_dict。现象在训练机器上无损加载的权重拿到另一台机器或者改了几行代码之后load_state_dict直接抛异常提示某些key缺失。原因保存和加载时的模型结构不一致。最常见的情况是训练时用了DataParallel自动在state_dict的key前加了module.前缀或者模型定义中途改动了层数、层名。解决加载时先打印model.load_state_dict(torch.load(best_model.pth), strictFalse)警告里缺失的key名。如果多出来module.前缀用一个循环去掉前缀再加载。如果是自己改了网络结构把models.py恢复到保存权重时的版本。尽量不要跨结构传权重这是很多人反复踩的坑。问题五批量预测时输入尺寸不一致报错。现象单张图预测正常对测试集循环预测时某些图片报shape错误程序中断。原因测试集里的图片尺寸五花八门模型在训练时见过固定224x224的输入测试时却进来一张1920x1080的图没走Resize。解决把预处理统一成一个函数对每张图强制先Resize再CenterCrop不要在每个批次里单独缩放。另外检查test_dataset的transform是否和val_transform完全一致漏掉一个就可能导致尺寸或数值分布偏差。6. 验证和导出用混淆矩阵、分类报告与ONNX收尾训练结束不代表项目完成还得能证明模型确实好用、能脱离训练环境跑起来。这一章讲最后的验证和导出。6.1 用classification_report看类别级的precision、recall与f1准确率是个很粗糙的指标类别不均衡时尤其容易骗人。比如多数类占90%模型全猜多数类也有90%准确率。要看模型真正的短板得看每个类别独立的precision、recall和f1。from sklearn.metrics import classification_report, confusion_matrix y_true, y_pred [], [] for images, labels in test_loader: images images.to(device) with torch.no_grad(): out model(images) y_true.extend(labels.numpy()) y_pred.extend(out.argmax(dim1).cpu().numpy()) # 打印每个类别的precision/recall/f1以及总体数据 print(classification_report(y_true, y_pred, target_namesclass_names)) print(confusion_matrix(y_true, y_pred))classification_report里有两列最关键。recall低的类别说明这个类的图大量被漏检precision低的类别说明大量其他类被误判成它。比如森林图像分类里如果“火烧迹地”的recall只有50%说明有一半的火烧迹地被认成林地这个项目就没法上线。再看混淆矩阵哪两个类别互相混淆最多优先想办法补数据或者加针对性增强。6.2 导出TorchScript或ONNX脱离训练代码做预测训练用的Python代码带着大量依赖演示时如果现场没有PyTorch环境就很尴尬。提前把模型导出成TorchScript或ONNX可以在没有训练框架的环境里做推理model.eval() dummy_input torch.randn(1, 3, 224, 224) # 固定形状的示例输入 # TorchScript导出适合PyTorch生态内使用 traced torch.jit.trace(model, dummy_input) traced.save(model_exported.pt) # ONNX导出适合跨语言跨平台部署 torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})dynamic_axes是导出ONNX时的关键参数它声明batch维度是可变的这样导出后既能预测单张图也能批量预测。设置后外部调用时不用刻意把batch固定为1。TorchScript导出的dummy_input形状要和训练一致否则trace出来的图结构是错的这是导出环节最隐蔽的坑。我现在的习惯是拿到任何CNN图像分类项目包第一件事永远不是跑训练而是先跑通预测脚本、确认权重能加载、环境没问题再谈改代码。顺序反了容易把环境问题误判成代码问题在错误方向上周旋半天。导出的模型文件放一份到演示目录答辩现场用不带PyTorch的机器也能跑出结果这个习惯帮我避开过不少次现场翻车。希望帮到你。本文还有配套的精品资源点击获取