
vit-pytorch
用纯 Transformer 做图像分类,代码简洁易上手,快速复现 SOTA 效果。
项目预览
项目数据
项目介绍
Star 增长趋势
本站每日记录该项目 GitHub 星标变化,观测越久曲线越完整
数据来源:GitHub API 逐日快照 · 收录于 2026-08-10 · 当前星标 ★ 25560
使用教程
环境要求
- 可用的 Python 环境(README 示例为 Python 代码)
- 已安装 PyTorch(示例代码中有 import torch)
- 可访问 PyPI 的 pip 网络环境
安装与启动步骤
-
1安装 vit-pytorch
使用 pip 从 PyPI 安装官方发布的 vit-pytorch 包,建议在虚拟环境中执行。
pip install vit-pytorch -
2确认 PyTorch 环境
README 示例依赖 torch,安装前请确保本机已能正常 import torch(README 未给出 PyTorch 安装命令)。
-
3编写 ViT 示例脚本
把 README 的 Usage 示例保存为 demo.py,按需修改 image_size、patch_size、num_classes 等参数。
import torch from vit_pytorch import ViT v = ViT( image_size = 256, patch_size = 32, num_classes = 1000, dim = 1024, depth = 6, heads = 16, mlp_dim = 2048, dropout = 0.1, emb_dropout = 0.1 ) img = torch.randn(1, 3, 256, 256) preds = v(img) # (1, 1000) -
4运行示例脚本
在保存 demo.py 的目录下执行脚本,观察是否成功输出形状为 (1, 1000) 的预测张量。
python demo.py
关键配置
| 配置项 | 必填 | 说明 | 示例 |
|---|---|---|---|
image_size | 是 | 输入图像的长宽尺寸 | 256 |
patch_size | 是 | 切分图像时每个 patch 的边长 | 32 |
num_classes | 是 | 图像分类的类别数量,决定输出维度 | 1000 |
dim | 否 | Transformer 隐层特征维度 | 1024 |
depth | 否 | Transformer 编码器的层数 | 6 |
heads | 否 | 多头注意力机制的头数 | 16 |
如何确认成功
运行 demo.py 后无报错,且 preds 的形状输出为 (1, 1000),与 num_classes 一致。
常见问题
Q:没有安装 PyTorch 可以直接用吗?
A:不可以。README 示例第一行就 import torch,需先自行准备好可用的 PyTorch 环境,再安装 vit-pytorch。
Q:如何修改处理不同尺寸的图像?
A:调整 ViT 初始化时的 image_size 参数,示例里为 256,输入张量形状需与之一致,如 (1, 3, 256, 256)。
Q:输出类别数由哪个参数决定?
A:由 num_classes 决定,示例设为 1000,因此输出 preds 的形状为 (1, 1000)。
Q:除了基础 ViT,还支持哪些变体?
A:README 目录列出了 DeiT、CaiT、DeepViT、CrossViT、MaxViT、MobileViT 等多种变体,可查看对应章节。
注意事项
- README 只提供 pip 安装与推理示例,未给出训练、数据集和预训练权重说明。
- 示例中的参数值仅作演示,实际使用时需按自己的图像尺寸和类别数调整。
- 更多模型变体的用法请参考 README 目录中对应的小节。
核心亮点
- 代码简洁模块化,每个组件独立清晰,便于学习和二次开发
- 支持多种 ViT 变体(DeiT、CaiT、DeepViT 等),覆盖前沿研究
- 文档和示例丰富,附带论文链接和详细注释,上手门槛低
不足之处
- 仅提供模型实现,不包含训练脚本和预训练权重,需自行训练
- 对大规模数据集训练效率较低,需配合分布式等优化手段
适用场景
- 学术研究:快速实验 ViT 架构变体,验证新想法
- 教学演示:讲解 Transformer 在视觉任务中的原理
- 模型对比:在自定义数据集上对比 ViT 与 CNN 的性能
替代项目
timm、huggingface/transformers、pytorch-image-models
同类项目推荐
ImageIndexer
开源
本地跑大模型,自动给图片打标签,搜图快人一步
Creates an index of images, queries a local LLM and adds tags to the image metadata
Occular-ocr
开源
俄语文档扫描件一键转文字,快且准
Fast and accurate OCR for Russian documents
gImageReader
开源
桌面点点鼠标,把图片和扫描件变成可复制文字
A Gtk/Qt front-end to tesseract-ocr.
PhysInOne
开源
让视频生成懂物理规律,动态重建更真实
PhysInOne in Python (CVPR 2026)
EmerNeRF
开源
自监督分解动态场景,轻松重建三维世界
PyTorch Implementation of EmerNeRF: Emergent Spatial-Temporal Scene Decomposition vi···
OpenCVTutorials
开源
中文 OpenCV 教程,从入门到实战,边看边跑代码
OpenCV-Python4.1 中文文档
MLOps_VideoAnomalyDetection
开源
在Azure上快速落地视频异常检测模型,从训练到部署全自动
Operationalize a video anomaly detection model with Azure ML
Awesome-Interaction-Aware-Trajec···
开源
一页看全轨迹预测前沿论文,省去大海捞针
A selection of state-of-the-art research materials on trajectory prediction
EMLight
开源
单图秒出光照环境贴图,让虚拟物体无缝融入真实场景
[AAAI 2021] EMLight: Lighting Estimation via Spherical Distribution Approximation, [···