ViT（Vision Transformer）简介

最新推荐文章于 2026-06-26 20:17:11 发布

原创

最新推荐文章于 2026-06-26 20:17:11 发布 · 510 阅读

标签

#transformer #深度学习 #人工智能

收录于

以下是一个简化的ViT（Vision Transformer）模型的实现代码示例。ViT模型用于图像分类任务，通过将图像分割成小块（patches），然后将每个小块视为一个序列输入到Transformer模型中。

import torch
import torch.nn as nn
from einops import rearrange

class PatchEmbedding(nn.Module):
    def __init__(self, in_channels=3, patch_size=16, emb_size=768, img_size=224):
        super().__init__()
        self.patch_size = patch_size
        self.proj = nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size)
        self.cls_token = nn.Parameter(torch.randn(1, 1, emb_size))
        self.pos_embedding = nn.Parameter(torch.randn((img_size // patch_size) ** 2 + 1, emb_size))

    def forward(self, x):
        B, C, H