### Vision Transformer 的代码实现
Vision Transformer (ViT) 是一种基于Transformer架构的模型，用于处理图像数据。下面是一个简化版 ViT 实现的例子：
```python
import torch
from torch import nn, einsum
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
class PreNorm(nn.Module):
def __init__(self, dim, fn):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.fn = fn
def forward(self, x, **kwargs):
return self.fn(self.norm(x), **kwargs)
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.):
super().__init__()
layers = [
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
]
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
class Attention(nn.Module):
def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
super().__init__()
inner_dim = dim_head * heads
project_out = not(heads == 1 and dim_head == dim)
self.heads = heads
self.scale = dim_head ** -0.5
linear_layers = [
('to_qkv', nn.Linear(dim, inner_dim*3, bias=False)),
('to_out_0', nn.Linear(inner_dim, dim)) if project_out else None,
('to_out_1', nn.Dropout(dropout))
]
self.attend = nn.Softmax(dim=-1)
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, dim),
nn.Dropout(dropout)
) if project_out else nn.Identity()
def forward(self, x):
qkv = self.to_qkv(x).chunk(3, dim=-1)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
attn = self.attend(dots)
out = einsum('b h i j, b h j d -> b h i d', attn, v)
out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)[^3]
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
super().__init__()
blocks = []
for _ in range(depth):
attention_block = PreNorm(dim, Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout))
feedforward_block = PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))
block = nn.Sequential(
attention_block,
feedforward_block
)
blocks.append(block)
self.layers = nn.Sequential(*blocks)
def forward(self, x):
return self.layers(x)
class ViT(nn.Module):
def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool='cls', channels=3, dim_head=64, dropout=0., emb_dropout=0.):
super().__init__()
assert image_size % patch_size == 0, "Image dimensions must be divisible by the patch size."
num_patches = (image_size // patch_size) ** 2
patch_dim = channels * patch_size ** 2
self.patch_embedding = nn.Sequential(
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size),
nn.Linear(patch_dim, dim)
)
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
self.dropout = nn.Dropout(emb_dropout)
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
self.pool = pool
self.to_latent = nn.Identity()
self.mlp_head = nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, num_classes)
)
def forward(self, img):
patches = self.patch_embedding(img)
cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=img.shape[0])
tokens = torch.cat((cls_tokens, patches), dim=1)
embeddings = tokens + self.pos_embedding[:, :(tokens.size(1)+1)]
dropped_embeddings = self.dropout(embeddings)
transformed_features = self.transformer(dropped_embeddings)
if self.pool == "mean":
output = reduced_representation.mean(dim=1)
elif self.pool == "cls":
output = transformed_features[:, 0]
final_output = self.to_latent(output)
class_logits = self.mlp_head(final_output)
return class_logits[^1][^2]
```
超级会员免费看
订阅专栏 解锁全文
扫一扫
4532
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
