ViT模型详解与Pytorch代码实现

class PatchEmbeddings(nn.Module):
    """
    Convert the image into patches and then project them into a vector space.
    """
    def __init__(self, config):
        super().__init__()
        self.image_size = config["image_size"]
        self.patch_size = config["patch_size"]
        self.num_channels = config[




    
"num_channels"]
        self.hidden_size = config["hidden_size"]
        # Calculate the number of patches from the image size and patch size
        self.num_patches = (self.image_size // self.patch_size) ** 2
        # Create a projection layer to convert the image into patches
        # The layer projects each patch into a vector of size hidden_size
        self.projection = nn.Conv2d(self.num_channels, self.hidden_size, kernel_size=self.patch_size, stride=self.patch_size)
    def forward(self, x):
        # (batch_size, num_channels, image_size, image_size) -> (batch_size, num_patches, hidden_size)
        x = self.projection(x)
        x = x.flatten(2).transpose(1, 2)
        return x

1. class Embeddings ( nn . Module ):

2. def __init__ ( self , config ):

3. super (). __init__ ()

4. self . config = config

5. self . patch_embeddings = PatchEmbeddings ( config )

6. # Create a learnable [CLS] token

7. # Similar to BERT, the [CLS] token is added to the beginning of the input sequence

8. # and is used to classify the entire sequence

9. self . cls_token = nn . Parameter ( torch . randn ( 1 , 1 , config [ "hidden_size" ]))

10. # Create position embeddings for the [CLS] token and the patch embeddings

11. # Add 1 to the sequence length for the [CLS] token

12. self . position_embeddings = \

13. nn . Parameter ( torch . randn ( 1 , self . patch_embeddings . num_patches + 1 , config [ "hidden_size" ]))

14. self . dropout = nn . Dropout ( config [ "hidden_dropout_prob" ])

16. def forward ( self , x ):

17. x = self . patch_embeddings ( x )

18. batch_size , _ , _ = x . size ()

19. # Expand the [CLS] token to the batch size

20. # (1, 1, hidden_size) -> (batch_size, 1, hidden_size)

21. cls_tokens = self . cls_token . expand ( batch_size , - 1 , - 1 )

22. # Concatenate the [CLS] token to the beginning of the input sequence

23. # This results in a sequence length of (num_patches + 1)

24. x = torch . cat (( cls_tokens , x ), dim = 1 )

25. x = x + self . position_embeddings

26. x = self . dropout ( x )

27. return x

1. class AttentionHead ( nn . Module ):

2. """

3. A single attention head.

4. This module is used in the MultiHeadAttention module.

5. """

6. def __init__ ( self , hidden_size , attention_head_size , dropout , bias = True ):

7. super (). __init__ ()

8. self . hidden_size = hidden_size

9. self . attention_head_size = attention_head_size

10. # Create the query, key, and value projection layers

11. self . query = nn . Linear ( hidden_size , attention_head_size , bias = bias )

12. self . key = nn . Linear ( hidden_size , attention_head_size , bias = bias )

13. self . value = nn . Linear ( hidden_size , attention_head_size , bias = bias )

15. self . dropout = nn . Dropout ( dropout )

17. def forward ( self , x ):

18. # Project the input into query, key, and value

19. # The same input is used to generate the query, key, and value,

20. # so it's usually called self-attention.

21. # (batch_size, sequence_length, hidden_size) -> (batch_size, sequence_length, attention_head_size)

22. query = self . query ( x )

23. key = self . key ( x )

24. value = self . value ( x )

25. # Calculate the attention scores

26. # softmax(Q*K.T/sqrt(head_size))*V

27. attention_scores = torch . matmul ( query , key . transpose (- 1 , - 2 ))

28. attention_scores = attention_scores / math . sqrt ( self . attention_head_size )

29. attention_probs = nn . functional . softmax ( attention_scores , dim =- 1 )

30. attention_probs = self . dropout ( attention_probs )

31. # Calculate the attention output

32. attention_output = torch . matmul ( attention_probs , value )

33. return ( attention_output , attention_probs )

class MultiHeadAttention(nn.Module):    """    Multi-head attention module.    This module is used in the TransformerEncoder module.    """
    def __init__(self, config):        super().__init__()        self.hidden_size = config["hidden_size"]        self.num_attention_heads = config["num_attention_heads"]        # The attention head size is the hidden size divided by the number of attention heads        self.attention_head_size = self.hidden_size // self.num_attention_heads        self.all_head_size = self.num_attention_heads * self.attention_head_size        # Whether or not to use bias in the query, key, and value projection layers        self.qkv_bias = config["qkv_bias"]        # Create a list of attention heads        self.heads = nn.ModuleList([])        for _ in range(self.num_attention_heads):            head = AttentionHead(                self.hidden_size,                self.attention_head_size,                config["attention_probs_dropout_prob"],                self.qkv_bias            )            self.heads.append(head)        # Create a linear layer to project the attention output back to the hidden size        # In most cases, all_head_size and hidden_size are the same        self.output_projection = nn.Linear(self.all_head_size, self.hidden_size)        self.output_dropout = nn.Dropout(config["hidden_dropout_prob"])

    def forward(self, x, output_attentions=False):        # Calculate the attention output for each attention head        attention_outputs = [head(x) for head in self.heads]        # Concatenate the attention outputs from each attention head        attention_output = torch.cat([attention_output for attention_output, _ in attention_outputs], dim=-1)        # Project the concatenated attention output back to the hidden size        attention_output = self.output_projection(attention_output)        attention_output = self.output_dropout(attention_output)        # Return the attention output and the attention probabilities (optional)        if not output_attentions:            return (attention_output, None)        else:            attention_probs = torch.stack([attention_probs for _, attention_probs in attention_outputs], dim=1)            return (attention_output, attention_probs)