GitHub 7.5k star量，各种视觉Transformer的PyTorch实现合集整理好了

MLNLP(机器学习算法与自然语言处理)社区是国内外最大的自然语言处理社区之一，汇聚超过50w订阅者，受众覆盖国内外NLP硕博生、高校老师以及企业研究人员。

社区的愿景是促进国内外自然语言处理，机器学习学术界、产业界和广大爱好者之间的交流和进步。

来源 | 机器之心

近一两年，Transformer 跨界 CV 任务不再是什么新鲜事了。

自 2020 年 10 月谷歌提出 Vision Transformer (ViT) 以来，各式各样视觉 Transformer 开始在图像合成、点云处理、视觉 - 语言建模等领域大显身手。

之后，在 PyTorch 中实现 Vision Transformer 成为了研究热点。GitHub 中也出现了很多优秀的项目，今天要介绍的就是其中之一。

该项目名为「vit-pytorch」，它是一个 Vision Transformer 实现，展示了一种在 PyTorch 中仅使用单个 transformer 编码器来实现视觉分类 SOTA 结果的简单方法。

项目当前的 star 量已经达到了 7.5k，创建者为 Phil Wang，ta 在 GitHub 上有 147 个资源库。

项目地址：https://github.com/lucidrains/vit-pytorch

项目作者还提供了一段动图展示：

项目介绍

首先来看 Vision Transformer-PyTorch 的安装、使用、参数、蒸馏等步骤。

第一步是安装：

$ pip install vit-pytorch

第二步是使用：

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)

第三步是所需参数，包括如下：

image_size：图像大小
patch_size：patch 数量
num_classes：分类类别的数量
dim：线性变换 nn.Linear(..., dim) 后输出张量的最后维
depth：Transformer 块的数量
heads：多头注意力层中头的数量
mlp_dim：MLP（前馈）层的维数
channels：图像通道的数量
dropout：Dropout rate
emb_dropout：嵌入 dropout rate
……

最后是蒸馏，采用的流程出自 Facebook AI 和索邦大学的论文《Training data-efficient image transformers & distillation through attention》。

论文地址：https://arxiv.org/pdf/2012.12877.pdf

从 ResNet50（或任何教师网络）蒸馏到 vision transformer 的代码如下：

import torchfrom torchvision.models import resnet50from vit_pytorch.distill import DistillableViT, DistillWrapperteacher = resnet50(pretrained = True)
v = DistillableViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 8,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)
distiller = DistillWrapper(
    student = v,
    teacher = teacher,
    temperature = 3,           # temperature of distillationalpha = 0.5,               # trade between main loss and distillation losshard = False               # whether to use soft or hard distillation
)
img = torch.randn(2, 3, 256, 256)labels = torch.randint(0, 1000, (2,))
loss = distiller(img, labels)loss.backward()
# after lots of training above ...pred = v(img) # (2, 1000)

除了 Vision Transformer 之外，该项目还提供了 Deep ViT、CaiT、Token-to-Token ViT、PiT 等其他 ViT 变体模型的 PyTorch 实现。

对 ViT 模型 PyTorch 实现感兴趣的读者可以参阅原项目。

技术交流群邀请函

△长按添加小助手

扫描二维码添加小助手微信

请备注：姓名-学校/公司-研究方向（如：小张-哈工大-对话系统）即可申请加入自然语言处理/Pytorch等技术交流群

关于我们

MLNLP(机器学习算法与自然语言处理) 社区是由国内外自然语言处理学者联合构建的民间学术社区，目前已经发展为国内外最大的自然语言处理社区之一，汇聚超过50w订阅者，旗下包括万人顶会交流群、AI臻选汇、AI英才汇以及AI学术汇等知名品牌，旨在促进机器学习，自然语言处理学术界、产业界和广大爱好者之间的进步。

社区可以为相关从业者的深造、就业及研究等方面提供开放交流平台。欢迎大家关注和加入我们。