简介面向深度学习图像分类入门者的数学图案三分类数据集包含正方形、圆形、三角形三类常见几何图形适合用来练习卷积神经网络分类模型、图像预处理与基础评估流程。数据已预分为训练集与测试集其中训练集240张、测试集60张文件夹名称即类别标签解压后可直接接入当前主流深度学习框架省去手动标注和整理目录的额外工作。压缩包内共303个文件主体为301张PNG格式图像另提供1个类别字典文件JSON格式和1个Python可视化脚本整体大小约209KB便于快速查看类别映射与样本示例。目前已有573人学习浏览对刚接触图像分类任务的学生或开发者颇为友好。借助可视化脚本可随机展示各类别样例确认标签顺序是否正确同时目录结构清晰方便自行扩展数据增强、迁移学习等实验是一份轻量实用的入门数据集。1. 图像分类数据集把「正方形、圆形、三角形」当做一个严肃的分类基准把正方形、圆形、三角形做成 3 分类的图像分类数据集听起来像是机器学习的 hello world。但真开工后你会发现这三类数学图案对神经网络的考验比真实照片更“硬”没有纹理、没有语义背景、也没有可信的参照物像素里可用的判别信息几乎全部集中在边缘走向、角点数量和轮廓闭合性上。模型一旦在训练时偷懒——只记住图形面积范围或者依赖背景噪声的统计特征换一批数据立刻现原形。所以这个数据集特别适合干两件事一是作为图像分类数据集的基准线验证从数据生成、目录组织到模型训练、指标分析的一整套流程是否干净二是当作几何特征学习的实验场方便观察 CNN 到底学到了哪些边缘模式。下面就从数据生成开始用 PyTorch 把这套 3 分类流程完整走一遍并给出几个不常见的排查点。2. 数据集构建目录规范、用 OpenCV 生成带扰动的三类数学图案2.1 先定目录train/val/test 三份标签名用英文小写图像分类数据集的第一步不是写模型而是把目录结构定下来。常见做法是每个标签一个子目录分成 train/val/test 三份命名统一为英文小写data/ ├── train/ │ ├── circle/ │ ├── square/ │ └── triangle/ ├── val/ └── test/这样做的好处是 torchvision 的ImageFolder和 Keras 的flow_from_directory都能直接读取不需要自己维护 CSV 索引文件。分配比例建议 7:1.5:1.5类别之间保持等量不重采样。文件全部用 PNG 保存不要用 JPGJPG 的块状压缩会让细线边缘出现毛刺对三角形尖角和正方形直边这类关键判别信息是直接损伤。生成数据时固定一个随机种子。我一般用np.random.default_rng(2025)保证后续任何人重复生成都能得到同一份数据。后面调模型、对比数据增强策略时数据本身必须是确定的。2.2 生成函数用 OpenCV 画三类数学图案并制造有效类内差异合成数据的目标不是“画得好看”而是“制造足够的类内差异”。如果每个三角形都放在图片正中心、大小完全一致模型学到的是“中心有灰色像素块三角形”而不是“有三条闭合直线边三角形”。我见过不少在合成数据上准确率 100%、换到真实拍摄图形就崩的案例根源基本都是生成阶段没有引入随机扰动。下面这个生成函数是完整可跑的用 OpenCV 画图并加入了中心偏移、尺寸随机、颜色随机、实心/空心随机、轮廓线宽随机、整图旋转和高斯噪声import cv2 import numpy as np def draw_any_polygon(img, pts, color, rng): 正方形和三角形共用随机实心或空心轮廓 if rng.random() 0.5: cv2.fillPoly(img, [pts], color, lineTypecv2.LINE_AA) else: cv2.polylines(img, [pts], True, color, thicknessint(rng.integers(1, 4)), lineTypecv2.LINE_AA) def draw_sample(class_name, rng): img np.full((128, 128, 3), 255, dtypenp.uint8) cx int(rng.integers(50, 78)) # 圆心/图形中心 x cy int(rng.integers(50, 78)) # 圆心/图形中心 y r int(rng.integers(24, 40)) # 外接半径 color tuple(int(v) for v in rng.integers(70, 220, size3)) if class_name circle: if rng.random() 0.5: cv2.circle(img, (cx, cy), r, color, -1, lineTypecv2.LINE_AA) else: cv2.circle(img, (cx, cy), r, color, thicknessint(rng.integers(1, 4)), lineTypecv2.LINE_AA) elif class_name square: half int(r * 0.9) # 用半边长控制正方形大小 pts np.array([ [cx - half, cy - half], [cx half, cy - half], [cx half, cy half], [cx - half, cy half] ], dtypenp.int32) draw_any_polygon(img, pts, color, rng) else: # triangle start float(rng.uniform(0, 2 * np.pi)) angles start np.array([0, 2 * np.pi / 3, 4 * np.pi / 3]) pts np.array([ [int(cx r * np.cos(a)), int(cy r * np.sin(a))] for a in angles ], dtypenp.int32) draw_any_polygon(img, pts, color, rng) # 整图旋转避免模型依赖“轴对齐”的边缘方向 M cv2.getRotationMatrix2D((64, 64), float(rng.uniform(-30, 30)), 1.0) img cv2.warpAffine(img, M, (128, 128), borderValue255) # 轻度高斯噪声防止模型直接依赖精确像素值 noise rng.normal(0, 6, img.shape) return np.clip(img noise, 0, 255).astype(np.uint8)几个参数值得单独说。中心坐标cx, cy取 50 到 78配合外接半径最大 39图形不会超出 128x128 画幅如果中心范围太大或者半径太大旋转后三角形顶点会被截断模型会把“一半三角形”当做一种新类别。颜色取 70 到 220 的中等亮度区间避免全黑全白导致归一化后数值分布过偏。实心和空心随机各占一半并随机 1 到 3 像素线宽这一步直接阻止模型通过“图形内部是否封闭”来分类。三角形用外接圆上间隔 120 度的三个点生成是等边三角形旋转后仍是等腰三角形。这里刻意保留规则形状因为不规则三角形的退化形态——比如接近直线的扁三角形——会带来大量低质量样本反而干扰训练。2.3 批量生成脚本与一份必要的简单质量检查生成脚本按 split 和 class 双层循环写文件。建议 train 每类 5000 张val 和 test 每类 750 张总计 19500 张。这个量级在 CPU 上用 OpenCV 生成也就几分钟但足够把后面那个轻量 CNN 训到可用状态。from pathlib import Path rng np.random.default_rng(2025) base Path(data) splits {train: 5000, val: 750, test: 750} classes [circle, square, triangle] for split, count in splits.items(): for cls in classes: out base / split / cls out.mkdir(parentsTrue, exist_okTrue) for i in range(count): img draw_sample(cls, rng) cv2.imwrite(str(out / f{cls}_{i:05d}.png), img) print(generation done)生成完先别急着训练跑一个几十行的质量检查脚本重点看三类样本数量、图像尺寸是否统一、以及有没有图形被边界截断from collections import Counter def inspect(root): files list(Path(root).rglob(*.png)) label_counter Counter(f.parent.name for f in files) sizes Counter() clipped 0 for f in files: img cv2.imread(str(f)) sizes[img.shape[:2]] 1 mask img.min(axis-1) 128 if mask[0].any() or mask[-1].any() or mask[:, 0].any() or mask[:, -1].any(): clipped 1 print(labels:, dict(label_counter)) print(sizes:, dict(sizes)) print(clipped:, clipped) inspect(Path(data))clipped大于 0 说明有图形贴到图像边界多半是中心坐标和半径范围没匹配好。这个问题在训练期几乎不会被发现因为模型照样能学出一个不错的准确率但换到没有截断的真实数据上就会掉点。把这类检查脚本放进数据生成的流水线比训练中途再去翻样本要省事得多。生成参数推荐值作用图像尺寸128x128保留三角形顶点和正方形直边信息中心坐标范围50~78图形不贴边避免截断外接半径24~40图形面积约占 28%~50%过小难分辨旋转范围±30°制造类内姿态变化噪声标准差6轻度扰动防止网络依赖精确像素值3. 模型选型与 PyTorch 训练3 分类轻量 CNN 足够用3.1 为什么先选 CNN而不是传统特征或 ViT面对正方形、圆形、三角形这类数学图案传统 CV 方案其实有非常成熟的路线Canny 边缘检测后配合霍夫直线检测和霍夫圆检测或者直接做轮廓逼近统计顶点个数在这类干净合成图上的准确率可以接近 100%。所以这里选 CNN 不是因为它在基准上一定更强而是因为它学到的边缘滤波器有通用性后续换真实数据时不需要重写特征流程。视觉 TransformerViT在理论上有更强的全局建模能力但在不到两万张的小型合成数据上收敛速度和稳定性通常不如小 CNN。ViT 缺乏 CNN 内置的局部归纳偏置对纯几何线条这种信息密度低的任务需要更大的数据量和更精细的学习率调度才能对齐。正确顺序是先跑通 CNN 基线再把它当对比对象去试 ViT 或预训练 ResNet而不是一开始就上大模型。方案参数量粗估对本数据集的适配度传统形状特征几乎为 0解释性强但抗噪和迁移能力弱小型 CNN约 10 万收敛快、易训练推荐首发预训练 ResNet18约 1100 万精度上限高但简单图上容易过拟合ViT数千万小数据下训练不稳需要额外调参3.2 一个带 BN 和自适应池化的 3 分类 CNN这个模型的参数量在 10 万左右结构是三层卷积加池化最后接自适应平均池化和一个 dropout 分类头import torch from torch import nn class ShapeNet(nn.Module): def __init__(self, num_classes3): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.AdaptiveAvgPool2d(1), ) self.classifier nn.Sequential( nn.Flatten(), nn.Dropout(0.3), nn.Linear(128, num_classes), ) def forward(self, x): return self.classifier(self.features(x))padding1配合 3x3 卷积保持空间分辨率不变每个池化层把尺寸减半128 → 64 → 32 → 16。BatchNorm2d放在卷积和激活之间稳定中间层分布这三分类任务上能明显加快收敛。AdaptiveAvgPool2d的作用是把任意尺寸的特征图压成 1x1这样即使后期把输入分辨率从 128 改成 64分类层的输入维度也不变不需要额外改代码。最后没有接 Softmax。PyTorch 的CrossEntropyLoss内部做了 LogSoftmax 和 NLLLoss 的组合直接输出 logits 即可再接 Softmax 反而会在数值和梯度上多绕一道。3.3 训练代码与超参数表训练部分不写自定义 Dataset直接用ImageFolder配合transformsimport torchvision.transforms as T from torchvision import datasets from torch.utils.data import DataLoader train_tf T.Compose([ T.Resize((128, 128)), T.RandomRotation(15, fill1), T.ToTensor(), T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) val_tf T.Compose([ T.Resize((128, 128)), T.ToTensor(), T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) train_ds datasets.ImageFolder(data/train, transformtrain_tf) val_ds datasets.ImageFolder(data/val, transformval_tf) train_loader DataLoader( train_ds, batch_size64, shuffleTrue, num_workers4, persistent_workersTrue, ) val_loader DataLoader( val_ds, batch_size256, shuffleFalse, num_workers4, ) model ShapeNet() opt torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) sched torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max60) criterion nn.CrossEntropyLoss()训练循环是标准写法每轮先model.train()遍历 train_loader 算 loss、反传、更新结束一个 epoch 后调用sched.step()调整学习率for epoch in range(60): model.train() for x, y in train_loader: x, y x.to(device), y.to(device) loss criterion(model(x), y) opt.zero_grad() loss.backward() opt.step() sched.step() # 这里接 val_loader 评估记录最佳权重RandomRotation(15, fill1)里的fill1容易被忽略。默认情况下旋转会在四角填充 0也就是黑边模型会额外学到“有黑边被旋转过”这个伪特征。这里的 fill1 对应归一化前的白色背景旋转后露出的区域与原图背景一致。超参数方面这套值在近两万张图上基本稳定超参数设置理由输入尺寸128x128保留顶点清晰度batch size64兼顾 BN 统计量与显存优化器AdamW lr3e-4训练平稳默认可用weight_decay1e-4抑制小模型过拟合epochs60配合 cosine 衰减后半程低学习率微调num_workers4本机 CPU 核数的一半避免 IO 瓶颈如果本机没有 GPU这个模型在 CPU 上也能跑完整个训练流程一个 epoch 通常只需要几十秒。先用 CPU 跑通全流程再切换训练平台能省掉大量环境排错时间。4. 验证与分析用混淆矩阵、触边检查和早停定位训练问题4.1 用混淆矩阵看三角形与圆形的真实混淆训练结束先不看准确率直接输出混淆矩阵和分类报告。准确率只能告诉你“好不好”混淆矩阵才能告诉你“错在哪类”from sklearn.metrics import confusion_matrix, classification_report def evaluate(model, loader, devicecpu): model.eval() preds, gts [], [] with torch.no_grad(): for x, y in loader: x x.to(device) preds.extend(model(x).argmax(1).cpu().tolist()) gts.extend(y.tolist()) print(classification_report( gts, preds, target_namesloader.dataset.classes)) print(confusion_matrix(gts, preds))model.eval()这一步必须写。它会把 BatchNorm 层从训练模式切换到推理模式使用累积的 running mean 和 running variance。如果不切同一批数据前后两次评估结果可能都不一样混淆矩阵也就没有参考价值。在这类几何图案数据上最常见的错误是三角形被分到圆形原因是小半径三角形在 128x128 分辨率下顶点被抗锯齿抹平轮廓逼近圆弧。如果混淆矩阵里triangle - circle的格子明显偏高优先提高输入分辨率或增大三角形的最小半径而不是盲目加数据量。4.2 三个容易掩盖问题的训练细节合成数据集的准确率高并不代表模型学到了形状下面是三个高频现象和对应处理症状可能原因处理训练准确率快速到 100%验证 80%数据增强不足模型记住了固定位置加大旋转范围给中心坐标更大随机loss 卡住不降准确率不动学习率设置不当从 3e-4 起手震荡就降到 1e-4验证准确率从第 15 个 epoch 开始回落过拟合提高 weight_decay 到 5e-4加大 Dropout第一种情况在合成数据里特别隐蔽如果中心坐标只在 50 到 78 之间移动模型完全可以靠“图形质心位置”猜类别。你看到准确率很高但换个分布它就失效。排查方法很简单——把训练集里两类样本的质心分布画出来如果三类的质心位置高度重叠说明数据没有问题如果可区分就要扩大中心随机范围。4.3 早停权重保存与回归基线训练 60 个 epoch 不代表第 60 轮的权重最好。更可靠的做法是每个 epoch 结束后计算验证集准确率历史最优时保存权重连续若干轮没有提升就提前停止best_acc, patience, no_improve 0.0, 8, 0 for epoch in range(60): train_one_epoch() acc evaluate_accuracy(val_loader) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_shape_net.pt) no_improve 0 else: no_improve 1 if no_improve patience: print(fearly stop at epoch {epoch}) breakpatience8对小型数据集比较合适太小容易在训练曲线抖动时提前终止太大则失去早停意义。注意保存条件是验证集准确率而不是 loss因为最终考核指标是分类准确率loss 平滑但与分类质量不一定完全同序。5. 进阶技巧用旋转一致性测试检验 CNN 是否学到了形状5.1 把旋转测试写进自动验证正方形、圆形、三角形都是旋转不变的类别旋转 90 度后类别含义完全不变。因此“旋转一致性”是这个数据集最自然的检验指标比单纯看测试集准确率更能暴露模型是否真的学到了形状。常用的做法是把验证集整体旋转若干固定角度后重新推理统计每个角度下的准确率from torchvision.transforms import functional as F def rotation_consistency(model, loader, angles(30, 90, 180, 270), devicecpu): model.eval() for a in angles: correct total 0 with torch.no_grad(): for x, y in loader: x, y x.to(device), y.to(device) xr F.rotate(x, a, fill1.0) pred model(xr).argmax(1) correct (pred y).sum().item() total y.size(0) print(frotation {a:3d} deg: acc{correct / total:.4f}) rotation_consistency(model, val_loader)如果生成阶段已经带 ±30° 旋转模型对 30° 以内的旋转应该几乎不掉点。90°、180°、270° 是生成范围之外的分布外测试最能说明问题如果 90° 旋转后准确率明显下降说明模型实际依赖的是边缘与图像坐标轴的相对角度而不是形状本身。fill1.0是这里最容易踩的坑。F.rotate作用于归一化后的 tensor输入背景值是 1.0所以填充也必须传 1.0。如果保持默认的 0旋转后会生成黑边模型会额外判断“有没有黑边”而不是“是什么形状”。这个测试不用重新训练几十秒就能跑完我一般把它和混淆矩阵一起固化成一个validate.py脚本每次修改数据生成参数或模型结构后都跑一遍。建议把 90° 和 270° 的准确率作为回归基线如果一次改动让这两个角度准确率下降超过 2%说明改动可能只是让模型更适应特定姿态需要重新评估。进阶一点的做法是同时统计实心样本和空心样本的旋转一致性。如果空心样本的 90° 准确率明显低于实心样本说明模型在依靠填充区域做判断这时需要增加空心室内的旋转增广或者把生成阶段空心比例从 0.5 提到 0.7。本文还有配套的精品资源点击获取 用OpenCV和PyTorch生成几何图像数据集并训练CNN
简介面向深度学习图像分类入门者的数学图案三分类数据集包含正方形、圆形、三角形三类常见几何图形适合用来练习卷积神经网络分类模型、图像预处理与基础评估流程。数据已预分为训练集与测试集其中训练集240张、测试集60张文件夹名称即类别标签解压后可直接接入当前主流深度学习框架省去手动标注和整理目录的额外工作。压缩包内共303个文件主体为301张PNG格式图像另提供1个类别字典文件JSON格式和1个Python可视化脚本整体大小约209KB便于快速查看类别映射与样本示例。目前已有573人学习浏览对刚接触图像分类任务的学生或开发者颇为友好。借助可视化脚本可随机展示各类别样例确认标签顺序是否正确同时目录结构清晰方便自行扩展数据增强、迁移学习等实验是一份轻量实用的入门数据集。1. 图像分类数据集把「正方形、圆形、三角形」当做一个严肃的分类基准把正方形、圆形、三角形做成 3 分类的图像分类数据集听起来像是机器学习的 hello world。但真开工后你会发现这三类数学图案对神经网络的考验比真实照片更“硬”没有纹理、没有语义背景、也没有可信的参照物像素里可用的判别信息几乎全部集中在边缘走向、角点数量和轮廓闭合性上。模型一旦在训练时偷懒——只记住图形面积范围或者依赖背景噪声的统计特征换一批数据立刻现原形。所以这个数据集特别适合干两件事一是作为图像分类数据集的基准线验证从数据生成、目录组织到模型训练、指标分析的一整套流程是否干净二是当作几何特征学习的实验场方便观察 CNN 到底学到了哪些边缘模式。下面就从数据生成开始用 PyTorch 把这套 3 分类流程完整走一遍并给出几个不常见的排查点。2. 数据集构建目录规范、用 OpenCV 生成带扰动的三类数学图案2.1 先定目录train/val/test 三份标签名用英文小写图像分类数据集的第一步不是写模型而是把目录结构定下来。常见做法是每个标签一个子目录分成 train/val/test 三份命名统一为英文小写data/ ├── train/ │ ├── circle/ │ ├── square/ │ └── triangle/ ├── val/ └── test/这样做的好处是 torchvision 的ImageFolder和 Keras 的flow_from_directory都能直接读取不需要自己维护 CSV 索引文件。分配比例建议 7:1.5:1.5类别之间保持等量不重采样。文件全部用 PNG 保存不要用 JPGJPG 的块状压缩会让细线边缘出现毛刺对三角形尖角和正方形直边这类关键判别信息是直接损伤。生成数据时固定一个随机种子。我一般用np.random.default_rng(2025)保证后续任何人重复生成都能得到同一份数据。后面调模型、对比数据增强策略时数据本身必须是确定的。2.2 生成函数用 OpenCV 画三类数学图案并制造有效类内差异合成数据的目标不是“画得好看”而是“制造足够的类内差异”。如果每个三角形都放在图片正中心、大小完全一致模型学到的是“中心有灰色像素块三角形”而不是“有三条闭合直线边三角形”。我见过不少在合成数据上准确率 100%、换到真实拍摄图形就崩的案例根源基本都是生成阶段没有引入随机扰动。下面这个生成函数是完整可跑的用 OpenCV 画图并加入了中心偏移、尺寸随机、颜色随机、实心/空心随机、轮廓线宽随机、整图旋转和高斯噪声import cv2 import numpy as np def draw_any_polygon(img, pts, color, rng): 正方形和三角形共用随机实心或空心轮廓 if rng.random() 0.5: cv2.fillPoly(img, [pts], color, lineTypecv2.LINE_AA) else: cv2.polylines(img, [pts], True, color, thicknessint(rng.integers(1, 4)), lineTypecv2.LINE_AA) def draw_sample(class_name, rng): img np.full((128, 128, 3), 255, dtypenp.uint8) cx int(rng.integers(50, 78)) # 圆心/图形中心 x cy int(rng.integers(50, 78)) # 圆心/图形中心 y r int(rng.integers(24, 40)) # 外接半径 color tuple(int(v) for v in rng.integers(70, 220, size3)) if class_name circle: if rng.random() 0.5: cv2.circle(img, (cx, cy), r, color, -1, lineTypecv2.LINE_AA) else: cv2.circle(img, (cx, cy), r, color, thicknessint(rng.integers(1, 4)), lineTypecv2.LINE_AA) elif class_name square: half int(r * 0.9) # 用半边长控制正方形大小 pts np.array([ [cx - half, cy - half], [cx half, cy - half], [cx half, cy half], [cx - half, cy half] ], dtypenp.int32) draw_any_polygon(img, pts, color, rng) else: # triangle start float(rng.uniform(0, 2 * np.pi)) angles start np.array([0, 2 * np.pi / 3, 4 * np.pi / 3]) pts np.array([ [int(cx r * np.cos(a)), int(cy r * np.sin(a))] for a in angles ], dtypenp.int32) draw_any_polygon(img, pts, color, rng) # 整图旋转避免模型依赖“轴对齐”的边缘方向 M cv2.getRotationMatrix2D((64, 64), float(rng.uniform(-30, 30)), 1.0) img cv2.warpAffine(img, M, (128, 128), borderValue255) # 轻度高斯噪声防止模型直接依赖精确像素值 noise rng.normal(0, 6, img.shape) return np.clip(img noise, 0, 255).astype(np.uint8)几个参数值得单独说。中心坐标cx, cy取 50 到 78配合外接半径最大 39图形不会超出 128x128 画幅如果中心范围太大或者半径太大旋转后三角形顶点会被截断模型会把“一半三角形”当做一种新类别。颜色取 70 到 220 的中等亮度区间避免全黑全白导致归一化后数值分布过偏。实心和空心随机各占一半并随机 1 到 3 像素线宽这一步直接阻止模型通过“图形内部是否封闭”来分类。三角形用外接圆上间隔 120 度的三个点生成是等边三角形旋转后仍是等腰三角形。这里刻意保留规则形状因为不规则三角形的退化形态——比如接近直线的扁三角形——会带来大量低质量样本反而干扰训练。2.3 批量生成脚本与一份必要的简单质量检查生成脚本按 split 和 class 双层循环写文件。建议 train 每类 5000 张val 和 test 每类 750 张总计 19500 张。这个量级在 CPU 上用 OpenCV 生成也就几分钟但足够把后面那个轻量 CNN 训到可用状态。from pathlib import Path rng np.random.default_rng(2025) base Path(data) splits {train: 5000, val: 750, test: 750} classes [circle, square, triangle] for split, count in splits.items(): for cls in classes: out base / split / cls out.mkdir(parentsTrue, exist_okTrue) for i in range(count): img draw_sample(cls, rng) cv2.imwrite(str(out / f{cls}_{i:05d}.png), img) print(generation done)生成完先别急着训练跑一个几十行的质量检查脚本重点看三类样本数量、图像尺寸是否统一、以及有没有图形被边界截断from collections import Counter def inspect(root): files list(Path(root).rglob(*.png)) label_counter Counter(f.parent.name for f in files) sizes Counter() clipped 0 for f in files: img cv2.imread(str(f)) sizes[img.shape[:2]] 1 mask img.min(axis-1) 128 if mask[0].any() or mask[-1].any() or mask[:, 0].any() or mask[:, -1].any(): clipped 1 print(labels:, dict(label_counter)) print(sizes:, dict(sizes)) print(clipped:, clipped) inspect(Path(data))clipped大于 0 说明有图形贴到图像边界多半是中心坐标和半径范围没匹配好。这个问题在训练期几乎不会被发现因为模型照样能学出一个不错的准确率但换到没有截断的真实数据上就会掉点。把这类检查脚本放进数据生成的流水线比训练中途再去翻样本要省事得多。生成参数推荐值作用图像尺寸128x128保留三角形顶点和正方形直边信息中心坐标范围50~78图形不贴边避免截断外接半径24~40图形面积约占 28%~50%过小难分辨旋转范围±30°制造类内姿态变化噪声标准差6轻度扰动防止网络依赖精确像素值3. 模型选型与 PyTorch 训练3 分类轻量 CNN 足够用3.1 为什么先选 CNN而不是传统特征或 ViT面对正方形、圆形、三角形这类数学图案传统 CV 方案其实有非常成熟的路线Canny 边缘检测后配合霍夫直线检测和霍夫圆检测或者直接做轮廓逼近统计顶点个数在这类干净合成图上的准确率可以接近 100%。所以这里选 CNN 不是因为它在基准上一定更强而是因为它学到的边缘滤波器有通用性后续换真实数据时不需要重写特征流程。视觉 TransformerViT在理论上有更强的全局建模能力但在不到两万张的小型合成数据上收敛速度和稳定性通常不如小 CNN。ViT 缺乏 CNN 内置的局部归纳偏置对纯几何线条这种信息密度低的任务需要更大的数据量和更精细的学习率调度才能对齐。正确顺序是先跑通 CNN 基线再把它当对比对象去试 ViT 或预训练 ResNet而不是一开始就上大模型。方案参数量粗估对本数据集的适配度传统形状特征几乎为 0解释性强但抗噪和迁移能力弱小型 CNN约 10 万收敛快、易训练推荐首发预训练 ResNet18约 1100 万精度上限高但简单图上容易过拟合ViT数千万小数据下训练不稳需要额外调参3.2 一个带 BN 和自适应池化的 3 分类 CNN这个模型的参数量在 10 万左右结构是三层卷积加池化最后接自适应平均池化和一个 dropout 分类头import torch from torch import nn class ShapeNet(nn.Module): def __init__(self, num_classes3): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.AdaptiveAvgPool2d(1), ) self.classifier nn.Sequential( nn.Flatten(), nn.Dropout(0.3), nn.Linear(128, num_classes), ) def forward(self, x): return self.classifier(self.features(x))padding1配合 3x3 卷积保持空间分辨率不变每个池化层把尺寸减半128 → 64 → 32 → 16。BatchNorm2d放在卷积和激活之间稳定中间层分布这三分类任务上能明显加快收敛。AdaptiveAvgPool2d的作用是把任意尺寸的特征图压成 1x1这样即使后期把输入分辨率从 128 改成 64分类层的输入维度也不变不需要额外改代码。最后没有接 Softmax。PyTorch 的CrossEntropyLoss内部做了 LogSoftmax 和 NLLLoss 的组合直接输出 logits 即可再接 Softmax 反而会在数值和梯度上多绕一道。3.3 训练代码与超参数表训练部分不写自定义 Dataset直接用ImageFolder配合transformsimport torchvision.transforms as T from torchvision import datasets from torch.utils.data import DataLoader train_tf T.Compose([ T.Resize((128, 128)), T.RandomRotation(15, fill1), T.ToTensor(), T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) val_tf T.Compose([ T.Resize((128, 128)), T.ToTensor(), T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]), ]) train_ds datasets.ImageFolder(data/train, transformtrain_tf) val_ds datasets.ImageFolder(data/val, transformval_tf) train_loader DataLoader( train_ds, batch_size64, shuffleTrue, num_workers4, persistent_workersTrue, ) val_loader DataLoader( val_ds, batch_size256, shuffleFalse, num_workers4, ) model ShapeNet() opt torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) sched torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max60) criterion nn.CrossEntropyLoss()训练循环是标准写法每轮先model.train()遍历 train_loader 算 loss、反传、更新结束一个 epoch 后调用sched.step()调整学习率for epoch in range(60): model.train() for x, y in train_loader: x, y x.to(device), y.to(device) loss criterion(model(x), y) opt.zero_grad() loss.backward() opt.step() sched.step() # 这里接 val_loader 评估记录最佳权重RandomRotation(15, fill1)里的fill1容易被忽略。默认情况下旋转会在四角填充 0也就是黑边模型会额外学到“有黑边被旋转过”这个伪特征。这里的 fill1 对应归一化前的白色背景旋转后露出的区域与原图背景一致。超参数方面这套值在近两万张图上基本稳定超参数设置理由输入尺寸128x128保留顶点清晰度batch size64兼顾 BN 统计量与显存优化器AdamW lr3e-4训练平稳默认可用weight_decay1e-4抑制小模型过拟合epochs60配合 cosine 衰减后半程低学习率微调num_workers4本机 CPU 核数的一半避免 IO 瓶颈如果本机没有 GPU这个模型在 CPU 上也能跑完整个训练流程一个 epoch 通常只需要几十秒。先用 CPU 跑通全流程再切换训练平台能省掉大量环境排错时间。4. 验证与分析用混淆矩阵、触边检查和早停定位训练问题4.1 用混淆矩阵看三角形与圆形的真实混淆训练结束先不看准确率直接输出混淆矩阵和分类报告。准确率只能告诉你“好不好”混淆矩阵才能告诉你“错在哪类”from sklearn.metrics import confusion_matrix, classification_report def evaluate(model, loader, devicecpu): model.eval() preds, gts [], [] with torch.no_grad(): for x, y in loader: x x.to(device) preds.extend(model(x).argmax(1).cpu().tolist()) gts.extend(y.tolist()) print(classification_report( gts, preds, target_namesloader.dataset.classes)) print(confusion_matrix(gts, preds))model.eval()这一步必须写。它会把 BatchNorm 层从训练模式切换到推理模式使用累积的 running mean 和 running variance。如果不切同一批数据前后两次评估结果可能都不一样混淆矩阵也就没有参考价值。在这类几何图案数据上最常见的错误是三角形被分到圆形原因是小半径三角形在 128x128 分辨率下顶点被抗锯齿抹平轮廓逼近圆弧。如果混淆矩阵里triangle - circle的格子明显偏高优先提高输入分辨率或增大三角形的最小半径而不是盲目加数据量。4.2 三个容易掩盖问题的训练细节合成数据集的准确率高并不代表模型学到了形状下面是三个高频现象和对应处理症状可能原因处理训练准确率快速到 100%验证 80%数据增强不足模型记住了固定位置加大旋转范围给中心坐标更大随机loss 卡住不降准确率不动学习率设置不当从 3e-4 起手震荡就降到 1e-4验证准确率从第 15 个 epoch 开始回落过拟合提高 weight_decay 到 5e-4加大 Dropout第一种情况在合成数据里特别隐蔽如果中心坐标只在 50 到 78 之间移动模型完全可以靠“图形质心位置”猜类别。你看到准确率很高但换个分布它就失效。排查方法很简单——把训练集里两类样本的质心分布画出来如果三类的质心位置高度重叠说明数据没有问题如果可区分就要扩大中心随机范围。4.3 早停权重保存与回归基线训练 60 个 epoch 不代表第 60 轮的权重最好。更可靠的做法是每个 epoch 结束后计算验证集准确率历史最优时保存权重连续若干轮没有提升就提前停止best_acc, patience, no_improve 0.0, 8, 0 for epoch in range(60): train_one_epoch() acc evaluate_accuracy(val_loader) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_shape_net.pt) no_improve 0 else: no_improve 1 if no_improve patience: print(fearly stop at epoch {epoch}) breakpatience8对小型数据集比较合适太小容易在训练曲线抖动时提前终止太大则失去早停意义。注意保存条件是验证集准确率而不是 loss因为最终考核指标是分类准确率loss 平滑但与分类质量不一定完全同序。5. 进阶技巧用旋转一致性测试检验 CNN 是否学到了形状5.1 把旋转测试写进自动验证正方形、圆形、三角形都是旋转不变的类别旋转 90 度后类别含义完全不变。因此“旋转一致性”是这个数据集最自然的检验指标比单纯看测试集准确率更能暴露模型是否真的学到了形状。常用的做法是把验证集整体旋转若干固定角度后重新推理统计每个角度下的准确率from torchvision.transforms import functional as F def rotation_consistency(model, loader, angles(30, 90, 180, 270), devicecpu): model.eval() for a in angles: correct total 0 with torch.no_grad(): for x, y in loader: x, y x.to(device), y.to(device) xr F.rotate(x, a, fill1.0) pred model(xr).argmax(1) correct (pred y).sum().item() total y.size(0) print(frotation {a:3d} deg: acc{correct / total:.4f}) rotation_consistency(model, val_loader)如果生成阶段已经带 ±30° 旋转模型对 30° 以内的旋转应该几乎不掉点。90°、180°、270° 是生成范围之外的分布外测试最能说明问题如果 90° 旋转后准确率明显下降说明模型实际依赖的是边缘与图像坐标轴的相对角度而不是形状本身。fill1.0是这里最容易踩的坑。F.rotate作用于归一化后的 tensor输入背景值是 1.0所以填充也必须传 1.0。如果保持默认的 0旋转后会生成黑边模型会额外判断“有没有黑边”而不是“是什么形状”。这个测试不用重新训练几十秒就能跑完我一般把它和混淆矩阵一起固化成一个validate.py脚本每次修改数据生成参数或模型结构后都跑一遍。建议把 90° 和 270° 的准确率作为回归基线如果一次改动让这两个角度准确率下降超过 2%说明改动可能只是让模型更适应特定姿态需要重新评估。进阶一点的做法是同时统计实心样本和空心样本的旋转一致性。如果空心样本的 90° 准确率明显低于实心样本说明模型在依靠填充区域做判断这时需要增加空心室内的旋转增广或者把生成阶段空心比例从 0.5 提到 0.7。本文还有配套的精品资源点击获取