四种基本形状数据集实战：从构建、增强到部署的避坑指南

简介这份数据集面向深度学习入门者、计算机视觉学习者及需要图形分类实验的开发者提供星形、圆形、正方形、三角形四种基本形状的标注图像可用于图像分类、卷积神经网络训练、数据增强与模型评估等基础实验场景。压缩包共收录14973个文件其中14970个为png格式的200×200像素图片另有3个py脚本便于直接读取、预处理或搭建训练流程整体包体约21.57MB体量轻便、易于下载与本地部署。目前已有293人学习下载适合作为课堂作业、课程设计或算法验证的练手素材。数据集覆盖四种形状类别样本数量充足能支撑从数据加载、模型搭建到准确率对比的完整实验链路帮助读者快速验证分类网络效果也可用于测试不同优化器、学习率与网络深度对识别精度的影响是入门视觉任务的实用起点。1. 四种基本形状数据集从“玩具级”到“产线级”的落差在哪星形、圆形、正方形、三角形这四个词摆在一起很多人第一反应是“太简单了随便训个模型都能跑满”。我最初也这么想直到把它接到一条真实的零件分拣线上——现场光照不均、形状有旋转、边缘还有毛刺模型准确率直接从实验室的 99% 掉到 71%。四种基本形状的图片数据集表面看是入门级分类任务实际是检验数据工程、增强策略和部署鲁棒性的绝佳试金石。它解决的核心问题不是“能不能分类”而是“在受控条件下把分类做到可复现、可迁移、可解释”。适合谁刚接触视觉分类想跑通全流程的新手以及需要一个小而干净的数据集来验证新增强策略、新网络结构或新部署链路的老手。这一章先把边界划清楚它简单但简单不等于没有坑。2. 数据集构建与格式选型为什么我坚持从 200 张起步2.1 四种形状的类间差异与类内陷阱星形、圆形、正方形、三角形的类间差异在理想情况下是几何拓扑级别的圆形没有角点正方形有四个直角三角形有三个锐角星形有十个交替的凹凸。但真实数据里类内差异会吃掉这些优势。圆形在透视下会变成椭圆正方形旋转 45 度后和菱形难以区分星形的凹角在低分辨率下会糊成圆形三角形的钝角版本和扇形边界模糊。我一般会先做一次类内多样性盘点而不是急着写 DataLoader。具体做法是每个类别至少覆盖旋转0 到 360 度步长 15 度、尺度占画幅 20% 到 80%、背景纯色、纹理、渐变、边缘锐利、抗锯齿、轻微模糊四个维度。如果某个类别在某个维度上缺失训练出来的模型就会在那个维度上翻车。常见做法是每个类别先收集 200 张四个类共 800 张作为基线。这个量级在迁移学习下足够跑出 95% 以上的验证准确率同时小到可以在单卡上几分钟跑完一轮方便快速迭代增强策略。不要一上来就堆到几万张那样调参周期会拖垮你的验证节奏。2.2 目录结构与标注格式ImageFolder 还是 COCO四种基本形状是典型的单标签分类任务最省事的格式是 ImageFolder 式的按类分目录dataset/ ├── train/ │ ├── circle/ │ ├── square/ │ ├── star/ │ └── triangle/ ├── val/ │ ├── circle/ │ ├── square/ │ ├── star/ │ └── triangle/ └── test/ ├── circle/ ├── square/ ├── star/ └── triangle/这种结构的好处是 torchvision 的ImageFolder和 TensorFlow 的image_dataset_from_directory都能直接读零解析成本。但如果你后续要做形状检测画框而不是分类就得换成 COCO 或 YOLO 格式每个形状带边界框标注。我一般会同时保留两份分类用 ImageFolder检测用 COCO。转换脚本不复杂但要注意边界框的坐标归一化和类别 ID 映射。下面是一个把 ImageFolder 转成 COCO 风格 JSON 的最小脚本import json import os from PIL import Image # 类别映射顺序固定避免训练时标签错位 CLASS_MAP {circle: 0, square: 1, star: 2, triangle: 3} def build_coco(root_dir, split): images, annotations [], [] ann_id 1 img_id 1 split_dir os.path.join(root_dir, split) for cls_name, cls_id in CLASS_MAP.items(): cls_dir os.path.join(split_dir, cls_name) if not os.path.isdir(cls_dir): continue for fname in os.listdir(cls_dir): if not fname.lower().endswith((.png, .jpg, .jpeg)): continue path os.path.join(cls_dir, fname) w, h Image.open(path).size images.append({id: img_id, file_name: path, width: w, height: h}) # 分类任务转检测时默认整图作为一个框后续可替换为真实标注 annotations.append({ id: ann_id, image_id: img_id, category_id: cls_id, bbox: [0, 0, w, h], area: w * h, iscrowd: 0 }) ann_id 1 img_id 1 return {images: images, annotations: annotations, categories: [{id: v, name: k} for k, v in CLASS_MAP.items()]} if __name__ __main__: coco build_coco(./dataset, train) with open(train_coco.json, w) as f: json.dump(coco, f)逻辑说明脚本遍历每个类别目录读取图片尺寸生成 COCO 标准的 images 和 annotations 两个数组。参数说明CLASS_MAP的顺序一旦确定就不能改否则训练时标签和模型输出会对不上bbox这里默认整图是因为分类转检测时如果没有真实框标注整图框只能作为占位不能直接用于检测训练。如果你有真实标注把bbox替换成标注值即可。提示类别 ID 从 0 开始还是从 1 开始不同框架要求不同。COCO 官方类别 ID 从 1 开始但很多 PyTorch 实现从 0 开始。转换后务必用一小批数据验证标签是否对齐。2.3 训练集、验证集、测试集的划分比例与随机种子四种形状数据集因为类间差异大划分比例不需要太激进。我一般用 70/15/15即训练 70%、验证 15%、测试 15%。如果每个类只有 200 张训练集 140 张、验证 30 张、测试 30 张。验证集用来调增强策略和早停测试集只在最后跑一次避免信息泄漏。随机种子必须固定。我见过太多人因为划分时没固定种子导致两次实验的验证集不同指标波动 3 到 5 个百分点还以为是模型改了有效。下面是一个带固定种子的划分脚本import random import os import shutil SEED 42 random.seed(SEED) def split_dataset(src_dir, dst_dir, ratios(0.7, 0.15, 0.15)): classes [circle, square, star, triangle] for cls in classes: files sorted(os.listdir(os.path.join(src_dir, cls))) random.shuffle(files) n len(files) n_train int(n * ratios[0]) n_val int(n * ratios[1]) splits { train: files[:n_train], val: files[n_train:n_train n_val], test: files[n_train n_val:] } for split, flist in splits.items(): out_dir os.path.join(dst_dir, split, cls) os.makedirs(out_dir, exist_okTrue) for f in flist: shutil.copy(os.path.join(src_dir, cls, f), os.path.join(out_dir, f)) if __name__ __main__: split_dataset(./raw, ./dataset)逻辑说明先对每个类别的文件列表排序再打乱保证可复现然后按比例切分并复制到目标目录。参数说明SEED固定后每次运行划分结果一致ratios三个值之和必须为 1如果数据量很小可以调成 80/10/10但验证集不能少于每类 20 张否则指标抖动太大。3. 增强策略与基线模型把 800 张用出 8000 张的效果3.1 几何增强的边界旋转、翻转、裁剪的合理范围四种形状数据集的增强核心是几何增强。但几何增强有边界不是随便转。圆形旋转任意角度都不变所以对圆形做旋转增强是无效的甚至有害——它会让模型学到“圆形可以出现在任何角度”这种废话浪费容量。正方形旋转 90 度不变但旋转 45 度就变成菱形如果测试集里没有菱形模型会把菱形误判成正方形或圆形。三角形旋转 180 度后方向相反星形旋转 72 度后自重合。我一般会按类别设置增强范围圆形只做平移和缩放不做旋转正方形做 90 度整数倍旋转三角形做 0 到 360 度全旋转星形做 72 度整数倍旋转。翻转方面圆形和正方形水平垂直翻转都不变三角形和星形翻转后会改变手性如果测试集里没有翻转样本训练时就不要加翻转。下面是一个按类别配置增强的 PyTorch 示例import torchvision.transforms as T # 按类别定义增强避免无效增强浪费模型容量 def get_transform(cls_name, is_trainTrue): if not is_train: return T.Compose([T.Resize((128, 128)), T.ToTensor()]) if cls_name circle: # 圆形旋转不变只做平移缩放 return T.Compose([ T.Resize((128, 128)), T.RandomAffine(degrees0, translate(0.1, 0.1), scale(0.8, 1.2)), T.ToTensor() ]) elif cls_name square: # 正方形只做 90 度整数倍旋转 return T.Compose([ T.Resize((128, 128)), T.RandomApply([T.RandomRotation((90, 90))], p0.5), T.RandomAffine(degrees0, translate(0.1, 0.1), scale(0.8, 1.2)), T.ToTensor() ]) elif cls_name triangle: # 三角形全旋转 return T.Compose([ T.Resize((128, 128)), T.RandomRotation(180), T.RandomAffine(degrees0, translate(0.1, 0.1), scale(0.8, 1.2)), T.ToTensor() ]) else: # star # 星形 72 度整数倍旋转 return T.Compose([ T.Resize((128, 128)), T.RandomApply([T.RandomRotation((72, 72))], p0.5), T.RandomAffine(degrees0, translate(0.1, 0.1), scale(0.8, 1.2)), T.ToTensor() ])逻辑说明每个类别走不同的增强分支圆形不做旋转正方形和星形做对称旋转三角形做全旋转。参数说明translate(0.1, 0.1)表示水平和垂直各平移 10%scale(0.8, 1.2)表示缩放 80% 到 120%这两个范围对四种形状都安全RandomRotation(180)表示在 -180 到 180 度之间随机旋转对三角形适用。注意增强策略必须和测试集分布匹配。如果测试集里只有正向三角形训练时做全旋转会让模型对方向不敏感反而降低正向样本的准确率。先看测试集再定增强。3.2 从 ResNet18 到轻量 CNN基线模型怎么选四种形状数据集上ResNet18 从零训练就能到 95% 以上但参数量 11M推理延迟在边缘设备上不够看。如果目标是部署到产线我一般会先跑一个轻量 CNN 作为基线再和 ResNet18 对比。轻量 CNN 的结构可以很简单三个卷积块每个块是 Conv-BN-ReLU-MaxPool最后全局平均池化加全连接。下面是一个最小实现import torch.nn as nn class ShapeCNN(nn.Module): def __init__(self, num_classes4): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d(1) ) self.classifier nn.Linear(128, num_classes) def forward(self, x): x self.features(x) x x.flatten(1) return self.classifier(x)逻辑说明三个卷积块逐步把空间尺寸从 128 降到 1通道数从 3 升到 128最后全连接输出 4 类。参数说明AdaptiveAvgPool2d(1)把任意空间尺寸压成 1x1这样输入图片尺寸可以灵活调整num_classes4对应四种形状。这个模型参数量约 0.1M在 CPU 上单张推理不到 5ms适合产线部署。训练时用交叉熵损失Adam 优化器学习率 1e-3batch size 32跑 50 个 epoch。如果验证准确率在 10 个 epoch 内没提升就降低学习率到 1e-4。我一般会同时跑 ResNet18 和 ShapeCNN如果 ShapeCNN 能达到 ResNet18 的 98% 准确率就选 ShapeCNN 部署。3.3 训练循环与早停别让模型在验证集上过拟合训练循环里最容易翻车的是早停策略。四种形状数据集小模型很容易在验证集上过拟合。我一般用验证损失而不是验证准确率作为早停指标因为准确率在 95% 以上时抖动大损失更平滑。import torch from torch.utils.data import DataLoader def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) def evaluate(model, loader, criterion, device): model.eval() total_loss, correct, total 0, 0, 0 with torch.no_grad(): for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) outputs model(imgs) loss criterion(outputs, labels) total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() total labels.size(0) return total_loss / len(loader), correct / total # 早停验证损失连续 8 个 epoch 不降就停 best_loss float(inf) patience, wait 8, 0 for epoch in range(50): train_loss train_one_epoch(model, train_loader, optimizer, criterion, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) if val_loss best_loss: best_loss val_loss wait 0 torch.save(model.state_dict(), best.pth) else: wait 1 if wait patience: print(fEarly stop at epoch {epoch}) break逻辑说明每个 epoch 先训练再验证验证损失创新低就保存模型并重置等待计数否则等待计数加一超过耐心值就停止。参数说明patience8表示连续 8 个 epoch 验证损失不降就停这个值对 800 张数据集体量合适best.pth保存的是验证损失最低的模型不是最后一个 epoch 的模型。4. 避坑与排查四种形状数据集上最容易翻车的 5 个点4.1 现象验证准确率 99%测试准确率 70%原因验证集和测试集划分时没有固定随机种子或者验证集被重复用于调参导致信息泄漏。更隐蔽的原因是增强策略在验证集上也生效了验证集分布和测试集不一致。解决固定随机种子验证集只做 Resize 和 ToTensor不做任何随机增强。测试集只在最后跑一次跑完之前不要看测试指标。4.2 现象圆形被大量误判为星形原因星形的凹角在低分辨率下被平滑掉轮廓接近圆形。如果训练集里星形图片分辨率普遍低于圆形模型会学到“低分辨率等于星形”这种伪特征。解决统一所有图片的输入分辨率至少 128x128。如果原始图片分辨率差异大先做一次分辨率归一化再送入训练。另外检查星形类别的边缘是否过于模糊必要时做锐化增强。4.3 现象训练损失震荡不收敛原因学习率太大或者 batch size 太小导致梯度噪声大。四种形状数据集小batch size 32 时每个 batch 只有 8 张每类梯度方差大。解决把学习率从 1e-3 降到 1e-4或者把 batch size 提到 64。如果显存不够用梯度累积模拟大 batch。另外检查数据归一化的均值和方差是否用了 ImageNet 的默认值如果数据集背景和 ImageNet 差异大最好重新计算。4.4 现象模型对旋转后的正方形识别率骤降原因训练时正方形只做了 90 度整数倍旋转但测试集里出现了 45 度旋转的正方形菱形。模型没见过这个角度把它判成了圆形或三角形。解决先确认测试集里是否有非 90 度倍数的正方形。如果有训练时必须加入对应角度的旋转增强。如果没有测试集就不应该出现这种样本检查测试集构建流程是否混入了异常数据。4.5 现象部署后推理结果和训练时不一致原因训练时用了 PIL 读图加 ToTensor部署时用了 OpenCV 读图通道顺序从 RGB 变成 BGR或者归一化参数不一致。这种问题在四种形状数据集上尤其明显因为形状颜色单一通道错位后颜色变化不大但模型看到的输入分布已经变了。解决训练和部署用同一套预处理代码或者至少用同一张图片对比两边预处理后的张量是否一致。我一般会写一个preprocess.py训练和部署都 import 这个模块避免手写两套。5. 进阶技巧用混淆矩阵和 t-SNE 把四种形状的边界看清楚训练完模型准确率只是一个数字。要真正理解模型学到了什么我一般会做两件事混淆矩阵和 t-SNE 可视化。混淆矩阵能告诉你哪两个类别最容易混。四种形状数据集上最常见的混淆是圆形和星形、正方形和三角形。如果混淆矩阵显示圆形和星形互混严重说明模型没有学到星形的凹角特征可能增强策略里星形的旋转范围不对或者分辨率不够。t-SNE 把最后一层特征降到二维看四个类别的聚类是否分开。如果圆形和星形在 t-SNE 上重叠说明特征空间里这两个类没有分离需要检查数据标注是否有误或者增加一个专门区分凹角的特征通道。下面是一个混淆矩阵和 t-SNE 的代码示例import numpy as np import matplotlib.pyplot as plt from sklearn.manifold import TSNE from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay # 收集测试集上的预测和特征 model.eval() all_preds, all_labels, all_feats [], [], [] with torch.no_grad(): for imgs, labels in test_loader: imgs imgs.to(device) feats model.features(imgs).flatten(1) # 取卷积特征 outputs model.classifier(feats) all_preds.extend(outputs.argmax(1).cpu().numpy()) all_labels.extend(labels.numpy()) all_feats.append(feats.cpu().numpy()) all_feats np.concatenate(all_feats, axis0) # 混淆矩阵 cm confusion_matrix(all_labels, all_preds) ConfusionMatrixDisplay(cm, display_labels[circle, square, star, triangle]).plot() plt.savefig(confusion_matrix.png) # t-SNE tsne TSNE(n_components2, random_state42, perplexity15) feats_2d tsne.fit_transform(all_feats) plt.figure() for cls_id, cls_name in enumerate([circle, square, star, triangle]): idx np.array(all_labels) cls_id plt.scatter(feats_2d[idx, 0], feats_2d[idx, 1], labelcls_name, s10) plt.legend() plt.savefig(tsne.png)逻辑说明先跑一遍测试集收集预测标签、真实标签和卷积层特征然后用 sklearn 的混淆矩阵和 t-SNE 分别可视化。参数说明perplexity15对 800 张数据集体量合适太大或太小都会让 t-SNE 图失去结构random_state42固定后每次图一致方便对比不同模型。看混淆矩阵时如果某一类的召回率低于 90%先看它的训练样本数是否偏少再看增强是否过度。看 t-SNE 时如果四个类聚成四团但有两团挨得很近说明模型能分但边界模糊可以尝试加一个对比损失或者中心损失来拉开距离。我自己的习惯是每次改完增强策略或模型结构先跑混淆矩阵确认没有类别被系统性误判再看 t-SNE确认特征空间没有塌缩。这两个图比准确率更能告诉我模型到底学到了什么。希望帮到你。本文还有配套的精品资源点击获取