Vision Transformer：图像怎样变成视觉 Token

N = (H / P) × (W / P)

python -m pip install "torch>=2.3,<3"

import torch
from torch import nn


class PatchEmbedding(nn.Module):
    def __init__(self, image_size=32, patch_size=8, channels=3, dim=64):
        super().__init__()
        self.grid = image_size // patch_size
        self.projection = nn.Conv2d(channels, dim, kernel_size=patch_size, stride=patch_size)
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.position = nn.Parameter(torch.zeros(1, self.grid**2 + 1, dim))

    def forward(self, images):
        patches = self.projection(images)       # B, D, grid, grid
        tokens = patches.flatten(2).transpose(1, 2)  # B, N, D
        cls = self.cls.expand(images.size(0), -1, -1)
        return torch.cat([cls, tokens], dim=1) + self.position


images = torch.randn(2, 3, 32, 32)
tokens = PatchEmbedding()(images)
print(images.shape)  # [2, 3, 32, 32]
print(tokens.shape)  # [2, 17, 64]