ViT_backbone.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from einops.layers.torch import Rearrange

class pre_process(nn.Module):
    def __init__(self, image_size, patch_size, patch_dim, dim):
        super().__init__()
        self.patch_size = patch_size
        self.dim = dim
        self.patch_num = (image_size // patch_size) ** 2
        self.linear_embedding = nn.Linear(patch_dim, dim)

        # Parameter是把这个tensor视为可学习参数
        # 加一个batch = 1的维度是为了方便广播到batch_size上
        self.position_embedding = nn.Parameter(torch.randn(1, self.patch_num + 1, self.dim))
        
        # 每个batch前都要加一个CLS_token
        self.CLS_token = nn.Parameter(torch.randn(1, 1, self.dim))

    def forward(self, x):
        x = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = self.patch_size, p2 = self.patch_size)
        x = self.linear_embedding(x)
        b, l, c = x.shape
        CLS_token = repeat(self.CLS_token, '1 1 d -> b 1 d', b = b)
        x = torch.concat((CLS_token, x), dim = 1)
        x = x + self.position_embedding
        return x

class Multihead_self_attention(nn.Module):
    def __init__(self, heads, head_dim, dim):
        super().__init__()
        self.head_dim = head_dim
        self.heads = heads
        self.inner_dim = self.head_dim * self.heads # 多头自注意力最后的输出维度
        self.scale = self.head_dim ** -0.5 # 1 / sqrt(head_dim)
        self.to_qkv = nn.Linear(dim, self.inner_dim * 3) # 对X做线性映射得到Q K V
        self.to_output = nn.Linear(self.inner_dim, dim) # 相当于Wo的映射作用
        self.norm = nn.LayerNorm(dim)
        self.softmax = nn.Softmax(dim=-1) 

    def forward(self, x):
        x = self.norm(x)
        qkv = self.to_qkv(x).chunk(3, dim = -1) # 返回列表，每个都是(b, l, inner_dim)
        Q, K, V = map(lambda t : rearrange(t, 'b l (h dim) -> b h l dim', dim = self.head_dim), qkv)

        # K.shape = (b, heads, l, head_dim)
        K_T = K.transpose(-1, -2)

        # attention_score.shape = (b, heads, l, l)
        attention_score = Q @ K_T * self.scale

        attention = self.softmax(attention_score)

        # out.shape = (b, heads, l, head_dim)
        out = attention @ V
        out = rearrange(out, 'b h l dim -> b l (h dim)')

        output = self.to_output(out)

        return output

class FeedForward(nn.Module):
    def __init__(self, dim, mlp_dim):
        super().__init__()
        self.fc1 = nn.Linear(dim, mlp_dim)
        self.fc2 = nn.Linear(mlp_dim, dim)
        self.norm = nn.LayerNorm(dim)
    
    def forward(self, x):
        x = self.norm(x)
        x = F.gelu(self.fc1(x))
        x = self.fc2(x)
        return x

class Transformer_block(nn.Module):
    def __init__(self, dim, heads, head_dim, mlp_dim):
        super().__init__()
        self.MHA = Multihead_self_attention(heads, head_dim, dim)
        self.Feedforward = FeedForward(dim, mlp_dim)

    def forward(self, x):
        x = self.MHA(x) + x
        x = self.Feedforward(x) + x
        return x

class ViT(nn.Module):
    def __init__(self, image_size, channels, patch_size, dim, heads, head_dim, mlp_dim, depth, num_class):
        super().__init__()
        self.to_patch_embedding = pre_process(
            image_size=image_size, 
            patch_size=patch_size, 
            patch_dim=channels * patch_size ** 2, 
            dim=dim
        )
        self.transformer = nn.ModuleList([
            Transformer_block(dim, heads, head_dim, mlp_dim)
            for _ in range(depth)
        ])

        self.MLP_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_class)
        )

        self.softmax = nn.Softmax(dim = -1)
    
    def forward(self, x):
        x = self.to_patch_embedding(x)

        for block in self.transformer:
            x = block(x)
        
        CLS_token = x[:, 0, :]
        out = self.softmax(self.MLP_head(CLS_token))
        return out