近一兩年,Transformer 跨界 CV 任務不再是什么新鮮事了,
自 2020 年 10 月谷歌提出 Vision Transformer (ViT) 以來,各式各樣視覺 Transformer 開始在影像合成、點云處理、視覺 - 語言建模等領域大顯身手,
之后,在 PyTorch 中實作 Vision Transformer 成為了研究熱點,GitHub 中也出現了很多優秀的專案,今天要介紹的就是其中之一,
該專案名為vit-pytorch,它是一個 Vision Transformer 實作,展示了一種在 PyTorch 中僅使用單個 transformer 編碼器來實作視覺分類 SOTA 結果的簡單方法,
專案當前的 star 量已經達到了 7.5k,創建者為 Phil Wang,他在 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, DistillWrapper
teacher = 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 實作感興趣的讀者可以參閱原專案,
轉載請註明出處,本文鏈接:https://www.uj5u.com/qita/400406.html
標籤:其他
