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)

class ViTForClassfication(nn.Module):
    """
    The ViT model for classification.
    """

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.image_size = config["image_size"]
        self.hidden_size = config["hidden_size"]
        self.num_classes = config["num_classes"]
        # Create the embedding module
        self.embedding = Embeddings(config)
        # Create the transformer encoder module
        self.encoder = Encoder(config)
        # Create a linear layer to project the encoder's output to the number of classes
        self.classifier = nn.Linear(self.hidden_size, self.num_classes)
        # Initialize the weights
        self.apply(self._init_weights)

    def forward(self, x, output_attentions=False):
        # Calculate the embedding output
        embedding_output = self.embedding(x)
        # Calculate the encoder's output
        encoder_output, all_attentions = self.encoder(embedding_output, output_attentions=output_attentions)
        # Calculate the logits, take the [CLS] token's output as features for classification
        logits = self.classifier(encoder_output[:, 0])
        # Return the logits and the attention probabilities (optional)
        if not output_attentions:
            return (logits, None)
        else:
            return (logits, all_attentions)

https://github.com/lukemelas/PyTorch-Pretrained-ViT/blob/master/pytorch_pretrained_vit/transformer.pyhttps://tintn.github.io/Implementing-Vision-Transformer-from-Scratch/

原价：498

折扣：399

推荐阅读

OpenCV4.8+YOLOv8对象检测C++推理演示

ZXING+OpenCV打造开源条码检测应用

总结 | OpenCV4 Mat操作全接触

三行代码实现 TensorRT8.6 C++ 深度学习模型部署

实战 | YOLOv8+OpenCV 实现DM码定位检测与解析

对象检测边界框损失 – 从IOU到ProbIOU

YOLOv8 OBB实现自定义旋转对象检测

初学者必看 | 学习深度学习的五个误区

YOLOv8自定义数据集训练实现安全帽检测