[入力画像] (H × W × C)
    ↓
[パッチ分割] → N個のパッチ (P × P × C)
    ↓
[パッチ埋め込み] → (N, D)
    ↓
[CLSトークンを先頭に追加] → (N+1, D)
    ↓
[位置埋め込みを加算] → (N+1, D)
    ↓
┌─────────────────────────┐
│ Transformer Encoder ×L  │
│  ├ Layer Normalization  │
│  ├ Multi-Head Self-Attn │
│  ├ 残差接続             │
│  ├ Layer Normalization  │
│  ├ MLP                  │
│  └ 残差接続             │
└─────────────────────────┘
    ↓
[CLSトークンの出力を取得] → (D,)
    ↓
[分類ヘッド (MLP)] → (num_classes,)

import torch
import torch.nn as nn


class PatchEmbedding(nn.Module):
    """画像をパッチに分割して埋め込みベクトルに変換"""
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2

        # 畳み込み層でパッチ分割と線形変換を同時に実行
        self.proj = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        """
        Args:
            x: (batch_size, channels, height, width)
        Returns:
            (batch_size, num_patches, embed_dim)
        """
        # (B, C, H, W) -> (B, D, H/P, W/P)
        x = self.proj(x)
        # (B, D, H/P, W/P) -> (B, D, N) -> (B, N, D)
        x = x.flatten(2).transpose(1, 2)
        return x

class TransformerEncoderBlock(nn.Module):
    """ViTのTransformer Encoderブロック（Pre-Normalization）"""
    def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(
            embed_dim, num_heads, dropout=dropout, batch_first=True
        )
        self.norm2 = nn.LayerNorm(embed_dim)

        mlp_hidden_dim = int(embed_dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, mlp_hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_hidden_dim, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        """
        Args:
            x: (batch_size, seq_len, embed_dim)
        Returns:
            (batch_size, seq_len, embed_dim)
        """
        # Pre-Normalization + Multi-Head Self-Attention + 残差接続
        x_norm = self.norm1(x)
        attn_out, _ = self.attn(x_norm, x_norm, x_norm)
        x = x + attn_out

        # Pre-Normalization + MLP + 残差接続
        x = x + self.mlp(self.norm2(x))
        return x

class VisionTransformer(nn.Module):
    """Vision Transformer (ViT)"""
    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_channels=3,
        num_classes=1000,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4.0,
        dropout=0.0,
    ):
        super().__init__()

        # パッチ埋め込み
        self.patch_embed = PatchEmbedding(
            img_size, patch_size, in_channels, embed_dim
        )
        num_patches = self.patch_embed.num_patches

        # CLSトークン（学習可能）
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

        # 位置埋め込み（学習可能）
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

        # Dropout
        self.pos_drop = nn.Dropout(p=dropout)

        # Transformer Encoderブロックを積み重ねる
        self.blocks = nn.ModuleList([
            TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, dropout)
            for _ in range(depth)
        ])

        # 最終のLayer Normalization
        self.norm = nn.LayerNorm(embed_dim)

        # 分類ヘッド
        self.head = nn.Linear(embed_dim, num_classes)

        # 重みの初期化
        self._init_weights()

    def _init_weights(self):
        # 位置埋め込みとCLSトークンの初期化
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)

        # 線形層とLayerNormの初期化
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.LayerNorm):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        """
        Args:
            x: (batch_size, channels, height, width)
        Returns:
            (batch_size, num_classes)
        """
        batch_size = x.shape[0]

        # パッチ埋め込み: (B, C, H, W) -> (B, N, D)
        x = self.patch_embed(x)

        # CLSトークンを先頭に追加: (B, N, D) -> (B, N+1, D)
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)

        # 位置埋め込みを加算
        x = x + self.pos_embed
        x = self.pos_drop(x)

        # Transformer Encoderブロックを適用
        for block in self.blocks:
            x = block(x)

        # 最終のLayer Normalization
        x = self.norm(x)

        # CLSトークンの出力を取得して分類
        cls_output = x[:, 0]  # (B, D)
        logits = self.head(cls_output)  # (B, num_classes)

        return logits

# モデル作成（ViT-Base/16相当）
model = VisionTransformer(
    img_size=224,
    patch_size=16,
    in_channels=3,
    num_classes=1000,
    embed_dim=768,
    depth=12,
    num_heads=12,
)

# ダミー入力
batch_size = 2
x = torch.randn(batch_size, 3, 224, 224)

# 順伝播
logits = model(x)

print(f"入力形状: {x.shape}")      # (2, 3, 224, 224)
print(f"出力形状: {logits.shape}") # (2, 1000)

# パラメータ数を確認
total_params = sum(p.numel() for p in model.parameters())
print(f"総パラメータ数: {total_params:,}")  # 約86M

import matplotlib.pyplot as plt
import numpy as np


def visualize_attention(model, img, patch_size=16):
    """Attention重みを可視化"""
    model.eval()

    with torch.no_grad():
        # パッチ埋め込み
        x = model.patch_embed(img)
        batch_size = x.shape[0]

        # CLSトークンと位置埋め込み
        cls_tokens = model.cls_token.expand(batch_size, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)
        x = x + model.pos_embed

        # 最後のブロックからAttention重みを取得
        for i, block in enumerate(model.blocks):
            x_norm = block.norm1(x)
            if i == len(model.blocks) - 1:
                # 最後のブロックでAttention重みを取得
                _, attn_weights = block.attn(
                    x_norm, x_norm, x_norm, need_weights=True
                )
            else:
                attn_out, _ = block.attn(x_norm, x_norm, x_norm)
                x = x + attn_out
                x = x + block.mlp(block.norm2(x))

    # CLSトークンから各パッチへのAttention重み
    # attn_weights: (B, num_heads, seq_len, seq_len)
    attn_cls = attn_weights[0, :, 0, 1:]  # (num_heads, num_patches)
    attn_mean = attn_cls.mean(dim=0)  # 全ヘッドの平均

    # パッチグリッドに変形
    num_patches_per_side = int(np.sqrt(attn_mean.shape[0]))
    attn_map = attn_mean.reshape(num_patches_per_side, num_patches_per_side)

    return attn_map.numpy()


# 可視化例（ランダム画像で）
torch.manual_seed(42)
dummy_img = torch.randn(1, 3, 224, 224)

model_small = VisionTransformer(
    img_size=224, patch_size=16, embed_dim=192, depth=4, num_heads=4, num_classes=10
)

attn_map = visualize_attention(model_small, dummy_img)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))

# 入力画像（正規化を戻してRGB表示用に変換）
img_display = dummy_img[0].permute(1, 2, 0).numpy()
img_display = (img_display - img_display.min()) / (img_display.max() - img_display.min())
axes[0].imshow(img_display)
axes[0].set_title("Input Image")
axes[0].axis("off")

# Attention重み
im = axes[1].imshow(attn_map, cmap="viridis")
axes[1].set_title("Attention from CLS token")
axes[1].axis("off")
plt.colorbar(im, ax=axes[1], shrink=0.8)

plt.tight_layout()
plt.show()