简介一套基于 Swin-Transformer 的图像识别完整项目覆盖自定义数据集从采集到训练、推理的全流程。项目内置关键词图像抓取脚本可自动下载网上图片并配套数据清洗与划分工具能够排查损坏图像、生成固定格式的训练集与测试集同时自动生成类别 JSON 文件。训练环节无需修改网络结构仅需调整学习率、轮数等超参数即开箱即用。预测脚本可批量推理 inference 目录下所有图片适合入门 Transformer 视觉任务或快速搭建图像分类流水线的开发者。压缩包共 1375 个文件约 723.38MB以 1200 个 JPEG 图像为主另有 PNG、WebP、BMP 等图片格式以及 14 个 Python 脚本、UI 界面文件、训练权重PTH和 README 说明文档目录结构清晰便于按环节对照使用。目前已有 385 人学习下载是动手实践 Swin-Transformer 的实用参考。 做图像识别项目最怕什么不是模型选不好而是数据半天搞不齐。我最近完整跑了一个基于Swin-Transformer的图像识别项目从“只有几个关键词”到“训练出能用的分类模型”整个过程踩了不少坑也沉淀了一套相对固定的流程。这篇就把整个过程完整写出来把数据集获取、清洗、训练、调参这些环节全拆开讲清楚。适合正在准备做深度学习图像识别项目或者想从CNN切到Transformer架构的同学参考。先说结论Swin-Transformer做图像识别完全够用它不像ViT那样需要海量数据和超长训练时间才能收敛在中等规模数据集上也能稳定出效果。但在动手之前数据和环境这两座大山得先翻过去这也是这篇想重点分享的从关键词到数据集再到训练完成每一步都有可以“抄作业”的细节。1. 为什么选Swin-Transformer作为主干网络1.1 从ResNet到Swin视觉模型的核心变化早几年做图像识别大家的第一反应都是ResNet、VGG这些CNN模型。CNN靠卷积核在局部窗口内提取特征通过堆叠层数扩大感受野。这个思路很成熟但有个天然短板远距离像素之间的依赖关系要经过很多层才能传递长距离建模能力偏弱。Swin-Transformer把NLP领域的自注意力机制搬到了图像上理论上每个位置都能直接和全局其他位置计算关联。但如果真让每对像素都做注意力计算量会爆炸——一张224×224的图有5万多个像素两两计算根本扛不住。Swin的解法是“窗口注意力”只在局部窗口内做自注意力窗口之间再用移位操作交换信息。这样既保住了Transformer的建模能力又把计算复杂度降到了和图像尺寸近似线性的水平。1.2 Swin相比ViT的优势在哪里ViT是直接把整张图切成16×16的patch序列再扔给标准的Transformer encoder。问题在于ViT没有图像固有的层级结构底层特征分辨率不变对小目标不友好训练也需要更大的数据集做支撑。Swin做了两件ViT没做的事第一是金字塔结构通过patch merging逐层下采样得到和CNN类似的多尺度特征图这让它做检测、分割时能直接复用FPN这类成熟的颈部结构第二是移位窗口通过窗口错位形成跨窗口连接等于在不同层之间引入了全局信息交换。实际项目里我拿Swin-T在自建数据集上做了对比实验同样的数据和时间预算下Top-1准确率比ResNet50高出2到3个点。这个差距在真实业务场景里可能就是上线与不上线的区别。1.3 选型时我考虑的开源生态只谈精度不谈工程落地就是耍流氓。Swin-Transformer官方仓库提供了完整的训练脚本、预训练权重和配置文件timm库里也有现成的模型实现HuggingFace transformers同样支持。这就意味着不管是做图像分类、目标检测还是语义分割都能找到一个开箱即用的基线。另外要注意Swin家族里Swin-T、Swin-S、Swin-B分别对应tiny、small、base三个规模。做中小型项目Swin-T是性价比最高的选择数据量大、算力充足再考虑Swin-B。别一上来就上最大模型训练时间翻几倍收益却不一定成正比。2. 训练环境搭建与硬件选型2.1 硬件怎么选显存和训练、推理的差别很多同学纠结显卡显存要怎么选其实要先分清是训练还是推理。推理只需要一次前向传播显存需求小得多训练要保存中间激活值、梯度、优化器状态显存消耗通常是推理的4到6倍。以Swin-T为例输入224×224、batch size为64时训练显存大约需要8到10GB一张RTX 3080或者RTX 4070就能跑。如果用Swin-B同样设置直接吃到20GB以上。所以我的建议是个人学习用一张12GB显存的卡足够batch size不够就用梯度累积如果是公司项目直接上24GB的RTX 3090/4090能少折腾很多事。CPU方面别忽视数据预处理和图像解码一个多核CPU能明显加快数据加载。内存32GB起步训练集图片多的时候DataLoader的prefetch会很吃内存。注意训练和推理的显存估算逻辑完全不同。推理时可以只算模型参数量乘以精度字节数训练时还要额外考虑优化器状态AdamW优化器光状态就要占模型大小的8倍以上。2.2 软件环境安装我用的是PyTorch生态。建议用conda建一个独立环境避免和系统自带环境冲突。关键依赖就这几个conda create -n swin python3.9 -y conda activate swin pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install timm tensorboard matplotlib tqdmtorch和torchvision的版本要和你机器上的CUDA驱动匹配。装完以后用下面这段代码验证一下GPU可用性能打印出显卡型号和CUDA版本才算环境没问题import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))2.3 项目目录与数据格式准备训练前先把目录结构规划好不然数据一多就乱套。我习惯用下面的结构project/ ├── configs/ # 配置文件 ├── data/ # 数据集 │ └── my_dataset/ │ ├── train/ │ │ ├── cat/ │ │ └── dog/ │ └── val/ │ ├── cat/ │ └── dog/ ├── scripts/ # 训练脚本 └── weights/ # 模型权重Swin-Transformer的DataLoader最省事的方式就是按ImageFolder格式读train目录下每个子文件夹代表一个类别图片放里面就行。torchvision的ImageFolder会自动把子文件夹名称映射成类别标签省去手动写label的麻烦。3. 关键词数据集获取与处理的完整流程3.1 用关键词批量获取图片这一节是标题里的重头戏。我这次项目要识别三类目标猫、狗、鸟。最初的输入就只有这“三个关键词”。关键词数据集的意思就是通过关键词在搜索引擎图片库里批量抓图快速建一个初版数据集。批量抓图的思路是用Python写一个爬虫把关键词作为搜索词拿到图片URL列表再下载。实现细节涉及搜索接口的解析不同的搜索引擎返回结构不一样。但整体套路是一致的构造搜索URL、解析返回的图片地址、去重下载。我用的脚本核心逻辑长这样import requests from bs4 import BeautifulSoup def search_images(keyword, num200): # 以必应图片搜索为例构造请求 urls [] for start in range(0, num, 35): url fhttps://www.bing.com/images/search?q{keyword}first{start} resp requests.get(url, headersHEADERS) soup BeautifulSoup(resp.text, html.parser) for a in soup.select(a.thumb): m re.search(rmurl(.?), a.get(href, )) if m: urls.append(m.group(1)) return list(set(urls))下载时务必控制并发数一次最多开5到8个线程加随机延时。千万别暴力抓取否则容易被封IP也给别人服务器添麻烦。图片格式统一转成RGB的JPG遇到gif、webp等格式要转码不然后面DataLoader读图会报错。3.2 清洗是个体力活关键词搜出来的图质量参差不齐。比如搜“猫”可能出现猫的玩偶、猫的卡通画、猫字书法甚至完全不相关的广告图。脏数据直接进训练集模型学到的特征就是歪的。我清洗分三步第一步自动去重。用感知哈希算出每张图片的指纹删除相似度大于0.9的重复图。第二步粗过滤。用预训练的CLIP模型对每张图算一次“图文匹配度”只保留和关键词匹配度高的图片。第三步人工过一遍。做了自动过滤之后几千张图人工翻一遍也就半小时效率可以接受。注意数据集的质量比重更重要。同一类别1000张干净图效果大概率好过10000张混入大量噪声的图。宁可少但要精。3.3 标注与数据集划分分类任务里类别就是文件夹名不需要额外做标注。但这里有个容易忽略的坑类别不平衡。关键词搜图时猫的图最好找鸟的相对少最后我鸟类的图只有猫的三分之一。训练时模型会对猫过拟合对鸟欠拟合。解决办法有两个一是通过旋转、翻转、色彩抖动等方式对少类做数据增强增加有效样本量二是在损失函数里加类别权重给少数的类别更大的梯度权重。我两个方案都用了实际验证下来效果最好的是在采样器上做调整用一个带权重的采样器让每个batch里各类别数量尽量均衡。数据按721的比例划分训练集、验证集和测试集。训练集用来迭代参数验证集用来调超参数和选模型测试集最后只跑一次验证最终效果。这个“三个集合互不干扰”的原则很重要很多人图省事把验证集当测试集反复用等上线才发现泛化能力不行悔之晚矣。4. Swin-Transformer训练脚本解析与超参数调优4.1 数据加载与增强Swin-T标准输入是224×224但直接缩放难免损失细节。我建议训练时用RandomResizedCrop做随机裁剪和缩放配合RandomHorizontalFlip做水平翻转再加Lighting和ColorJitter做颜色扰动。验证集只用Resize和CenterCrop不需要随机增强保证评估的稳定性。from torchvision import datasets, transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.4, 0.4, 0.4), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(data/my_dataset/train, transformtrain_transform)4.2 训练核心参数Swin-Transformer官方在ImageNet上的配置是300轮、学习率5e-4、batch size 128。但那是针对百万级数据量的自建的小数据集照搬这套配置必翻车。我的经验是5000张图的数据集训练轮数50到100轮就够学习率降到1e-4到2e-4之间起步。优化器用AdamW这个不用犹豫。Swin官方用的就是这个它在参数更新时对权重衰减做了修正比传统Adam更好训。权重衰减值一般设0.05太大容易欠拟合太小正则效果弱。学习率策略我强烈推荐cosine退火配合前5到10轮的warmup。warmup的意义在于训练初期模型权重是随机初始化的梯度方向噪声很大直接上大学习率容易不稳定甚至发散。先用小学习率“热身”等loss进入下降通道后再按cosine曲线平滑降低学习率这也是Transformer系模型的通用惯例。optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.05) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxnum_epochs)4.3 评估与权重保存每训练一个epoch就在验证集上跑一次分类准确率。注意要等模型进入eval模式再评估pytorch里是model.eval()还要记得用with torch.no_grad()包住推理过程否则会白占显存。保存权重时别只存模型参数把optimizer的state_dict和epoch数也一起存下来这样断点续训不用重来。torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_acc: best_acc }, fweights/swin_epoch_{epoch}.pth)我实际训练时batch size设为32显存占用约6GB一个epoch跑5000张图大约3分钟。50轮下来不到3小时就能训完这个节奏是很舒服的——白天改配置晚上挂训练第二天起来看结果。5. 视频图像识别的延伸要不要做视频解码5.1 视频识别与单帧识别的区别很多项目做到一半会发现业务方的输入不是图片而是一段视频。这时候会碰到一个灵魂拷问视频图像识别是否要做视频解码答案是要做而且要做得聪明。视频本质是连续的图像帧序列但直接用视频文件喂给模型是行不通的——绝大部分图像识别模型只接受单张图片输入它不知道什么叫帧、什么叫编码。所以视频识别的前置步骤是从视频里抽出若干帧变成普通图片再走图像识别流程。抽帧不是每帧都抽。每帧都识别一是计算量成倍涨二是相邻帧高度相似信息冗余大容易造成结果抖动。常规做法是每秒抽1到2帧或者按场景切换抽帧画面发生明显变化时抽一帧。5.2 帧采样的工程实现我用OpenCV做帧抽取代码非常简洁import cv2 cap cv2.VideoCapture(input.mp4) fps cap.get(cv2.CAP_PROP_FPS) frame_interval int(fps / 1) # 每秒抽1帧 idx 0 while cap.isOpened(): ret, frame cap.read() if not ret: break if idx % frame_interval 0: cv2.imwrite(fframes/frame_{idx:06d}.jpg, frame) idx 1 cap.release()抽出来的帧再挨个送进训练好的Swin模型做推理。如果想提升识别稳定性可以做一个简单的“投票法”同一个目标的连续N帧里哪类被预测得最多就定为最终结果。这个技巧能在不增加模型复杂度的情况下显著降低单帧误判率实测很管用。6. 常见训练难点与排查实录6.1 loss不降怎么办loss迟迟不下降是训练中最容易让人焦虑的问题。我的排查顺序是先看数据随机打印20张训练图片确认图片本身没坏、标签没串再看学习率用pytorch的lr_finder工具或者手动试几个数量级1e-3、1e-4、1e-5观察哪个量级loss有下降趋势最后看代码确认梯度没有写错、backward和step有没有漏写。八成问题出在前两步。如果loss正常下降但acc不动看是不是数据标签分布有问题。比如三分类里有两类图高度相似模型很难区分这时候要么合并类别要么增加细分标注靠调参已经解决不了根本问题。6.2 显存溢出显存溢出的报错信息里只要有“out of memory”就是它了。最简单的应对是调小batch size比如从32调到16或8。但batch太小会导致BN统计量不准训练不稳定。更优雅的解法是梯度累积accumulation_steps 4 # 实际batch 32 / 4 for i, (images, labels) in enumerate(train_loader): outputs model(images) loss criterion(outputs, labels) loss loss / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()这个技巧就是在显存不动的前提下用时间换空间让batch size的“等效值”变大。我用这个方案在8GB显存的卡上跑过Swin-Sbatch size等效到32也没问题。6.3 过拟合自建数据集规模不大过拟合几乎是必经之路。现象是训练集acc一路飙到98%验证集才刚刚过80%。我在这个项目里用了三层防过拟合一是MixUp数据增强按一定比例混合两张训练图及其标签强制模型学习更鲁棒的特征二是加大weight_decay从0.05增到0.1三是早停机制验证集acc连续10轮不再提升就停止训练保存最佳模型权重。三者叠加之后验证集和测试集的acc差距从15个点压缩到了5个点以内可接受的过拟合范围。我在实际项目里的体会是Swin-Transformer这套方案选对了之后真正的瓶颈往往不在模型本身而在数据获取和处理上。关键词数据集的方法能快速验证想法但要想让模型真正落地还得在数据清洗和采样策略上多花心思。最后再分享一个容易被忽略的小技巧训练前最好固定一下随机种子这样每次跑出来的结果是可以复现的排查问题的时候会轻松很多。项目做完之后这套流程我直接复用到了其他领域效果都挺稳你可以放心照着试。本文还有配套的精品资源点击获取