VIT vision transformer pytorch代码复现

最新推荐文章于 2024-08-17 18:22:32 发布

haohulala

最新推荐文章于 2024-08-17 18:22:32 发布

阅读量5.3k

点赞数 9

分类专栏：计算机视觉

本文链接：https://blog.csdn.net/haohulala/article/details/115747883

版权

这篇博客介绍了一种将Transformer机制应用于计算机视觉，尤其是图片分类的方法。通过将图片切分成小块并利用TransformerEncoder进行处理，实现了对图像的有效理解。关键组件包括多头自注意力（MSA）模块和多层感知机（MLP），最终通过平均池化和全连接层进行分类。代码示例展示了模型的具体实现，包括自注意力计算和TransformerEncoder层的构建。

摘要由CSDN通过智能技术生成

这篇论文将transformer机制运用到计算机视觉领域（主要是进行了图片分类），并且取得了不错的效果

其实整体思路挺简单的，就是将是图片拆分成很多小块，然后将小块排列成矩阵送入transformer encoder模块中计算，具体的计算过程如下图所示

我主要分享一下代码

import torch
import torch.nn as nn
import math

class MLP(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super(MLP, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, input):
        output = self.net(input)
        return output

class MSA(nn.Module):
    "