vit-pytorch

用纯 Transformer 做图像分类,代码简洁易上手,快速复现 SOTA 效果。

开源 🏷️完全免费 🔧 仅开源项目,无在线服务,需自行部署 计算机视觉 🍴 复刻 3495 🐛 待处理 153 👀 关注 155
GitHub ↗
GitHub 星标 ★ 25560
维护状态 活跃
是否开源 开源
定价模式 🏷️完全免费 开源项目,MIT许可证,完全免费,无付费版本。

项目预览

vit-pytorch 项目界面预览图
图片来源:GitHub 官方仓库名片

项目数据

分类计算机视觉
开发团队lucidrains
定价模式🏷️完全免费
价格说明开源项目,MIT许可证,完全免费,无付费版本。
访问状态🔧 仅开源项目,无在线服务,需自行部署
是否开源是
开源协议MIT
主要语言Python
技术栈/模型artificial-intelligence,attention-mechanism,computer-vision,image-classification,transformers
GitHub 星标★ 25560
上线时间2020-10-03 00:00:00
代码更新2026-11-07
仓库体积11.0 MB
本站数据更新2026-11-08
维护状态活跃
收录时间2026-08-10
浏览次数8

项目介绍

vit-pytorch 是一个基于 PyTorch 实现的 Vision Transformer(ViT)模型库,由 Phil Wang(lucidrains)开发。它用简洁的代码实现了 ViT 论文中的核心思想,即通过纯 Transformer 编码器处理图像分类任务,无需卷积神经网络。该项目解决了传统 CNN 在图像分类中依赖局部感受野的问题,提供了一种全局注意力机制来捕捉图像中的长距离依赖关系。核心能力包括:完整的 ViT 模型实现、支持多种变体(如 DeiT、CaiT、DeepViT 等)、灵活的 patch 嵌入和位置编码、易于扩展的模块化设计。代码风格清晰,注释详尽,适合研究者和开发者快速理解 Transformer 在视觉领域的应用。项目在 GitHub 上拥有超过 2.5 万星标,是计算机视觉领域最流行的 ViT 参考实现之一。

Star 增长趋势

本站每日记录该项目 GitHub 星标变化,观测越久曲线越完整

数据来源:GitHub API 逐日快照 · 收录于 2026-08-10 · 当前星标 ★ 25560

使用教程

难度:入门 约 10 分钟 部署方式:库/依赖 4 步

环境要求

  • 可用的 Python 环境(README 示例为 Python 代码)
  • 已安装 PyTorch(示例代码中有 import torch)
  • 可访问 PyPI 的 pip 网络环境

安装与启动步骤

  1. 1安装 vit-pytorch

    使用 pip 从 PyPI 安装官方发布的 vit-pytorch 包,建议在虚拟环境中执行。

    pip install vit-pytorch
  2. 2确认 PyTorch 环境

    README 示例依赖 torch,安装前请确保本机已能正常 import torch(README 未给出 PyTorch 安装命令)。

  3. 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. 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

上一篇:mediapipe

下一篇:vision

同类项目推荐

ImageIndexer 开源

本地跑大模型,自动给图片打标签,搜图快人一步

Creates an index of images, queries a local LLM and adds tags to the image metadata

★ 399 2026-08-09
Occular-ocr 开源

俄语文档扫描件一键转文字,快且准

Fast and accurate OCR for Russian documents

★ 74 2026-10-18
gImageReader 开源

桌面点点鼠标,把图片和扫描件变成可复制文字

A Gtk/Qt front-end to tesseract-ocr.

★ 2022 2026-10-27
PhysInOne 开源

让视频生成懂物理规律,动态重建更真实

PhysInOne in Python (CVPR 2026)

★ 118 2026-08-09
EmerNeRF 开源

自监督分解动态场景,轻松重建三维世界

PyTorch Implementation of EmerNeRF: Emergent Spatial-Temporal Scene Decomposition vi···

★ 641 2026-08-09
OpenCVTutorials 开源

中文 OpenCV 教程,从入门到实战,边看边跑代码

OpenCV-Python4.1 中文文档

★ 1440 2026-08-20
MLOps_VideoAnomalyDetection 开源

在Azure上快速落地视频异常检测模型,从训练到部署全自动

Operationalize a video anomaly detection model with Azure ML

★ 138 2026-08-09
EMLight 开源

单图秒出光照环境贴图,让虚拟物体无缝融入真实场景

[AAAI 2021] EMLight: Lighting Estimation via Spherical Distribution Approximation, [···

★ 172 2026-08-09