# =================================================== # Vision Transformer（ViT）の実装 # =================================================== import torch import torch.nn as nn import torch.nn.functional as F # =================================================== # 1. パッチ埋め込み層 # =================================================== class PatchEmbedding(nn.Module): “”” 画像をパッチに分割して埋め込みベクトルに変換 処理の流れ: 1. 画像を16×16のパッチに分割 2. 各パッチを768次元ベクトルに変換（線形変換） “”” def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768): “”” Args: img_size: 入力画像のサイズ（正方形を想定） patch_size: パッチのサイズ（16×16が標準） in_channels: 入力チャンネル数（RGBなら3） embed_dim: 埋め込み次元（ViT-Baseは768） “”” super().__init__() self.img_size = img_size self.patch_size = patch_size # パッチ数を計算: (224÷16) × (224÷16) = 14×14 = 196 self.num_patches = (img_size // patch_size) ** 2 # 線形変換を畳み込みで実装 # kernel_size=stride=patch_size により、重なりなしでパッチを処理 # これは各パッチを展開して線形変換するのと数学的に同じ self.proj = nn.Conv2d( in_channels, # 入力: 3チャンネル（RGB） embed_dim, # 出力: 768チャンネル kernel_size=patch_size, # 16×16のカーネル stride=patch_size # 16ピクセルずつ移動（重なりなし） ) def forward(self, x): “”” Args: x: 入力画像 (batch_size, 3, 224, 224) Returns: パッチ埋め込み (batch_size, num_patches, embed_dim) = (batch_size, 196, 768) “”” # 畳み込みでパッチ埋め込み # (B, 3, 224, 224) → (B, 768, 14, 14) x = self.proj(x) # 空間次元を1次元に展開 # (B, 768, 14, 14) → (B, 768, 196) x = x.flatten(2) # 次元を入れ替えてTransformerの入力形式に # (B, 768, 196) → (B, 196, 768) x = x.transpose(1, 2) return x

# =================================================== # 2. Multi-Head Self-Attention # =================================================== class MultiHeadSelfAttention(nn.Module): “”” Multi-Head Self-Attention 処理の流れ: 1. 入力からQ（Query）, K（Key）, V（Value）を計算 2. Q×K^Tで各トークン間のAttentionスコアを計算 3. softmaxで正規化 4. Attention重みでVを重み付け和 5. 複数ヘッドの結果を結合 “”” def __init__(self, embed_dim=768, num_heads=12, dropout=0.0): “”” Args: embed_dim: 埋め込み次元（768） num_heads: Attentionヘッド数（12） dropout: ドロップアウト率 “”” super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads # 各ヘッドの次元: 768 ÷ 12 = 64 self.head_dim = embed_dim // num_heads # embed_dimがnum_headsで割り切れることを確認 assert embed_dim % num_heads == 0, \ f”embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})” # Q, K, V を一度に計算する線形層 # 768次元 → 768×3 = 2304次元 # これを後でQ, K, Vに分割 self.qkv = nn.Linear(embed_dim, embed_dim * 3) # 出力の線形変換 self.proj = nn.Linear(embed_dim, embed_dim) # ドロップアウト self.attn_dropout = nn.Dropout(dropout) self.proj_dropout = nn.Dropout(dropout) # スケーリング係数: 1/√(head_dim) = 1/√64 = 0.125 self.scale = self.head_dim ** -0.5 def forward(self, x): “”” Args: x: 入力 (batch_size, num_tokens, embed_dim) = (B, 197, 768) Returns: 出力 (batch_size, num_tokens, embed_dim) “”” B, N, C = x.shape # B=batch, N=197, C=768 # Q, K, V を計算 # (B, N, C) → (B, N, 3*C) → (B, N, 3, num_heads, head_dim) qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) # 次元を入れ替え: (3, B, num_heads, N, head_dim) qkv = qkv.permute(2, 0, 3, 1, 4) # Q, K, V に分割（それぞれ (B, num_heads, N, head_dim)） q, k, v = qkv[0], qkv[1], qkv[2] # Attentionスコアを計算 # Q × K^T: (B, h, N, d) × (B, h, d, N) = (B, h, N, N) # スケーリングで値が大きくなりすぎるのを防ぐ attn = (q @ k.transpose(-2, -1)) * self.scale # softmaxで正規化（各行の合計が1になる） attn = attn.softmax(dim=-1) # ドロップアウト（正則化） attn = self.attn_dropout(attn) # Attention重みでVを重み付け和 # (B, h, N, N) × (B, h, N, d) = (B, h, N, d) x = attn @ v # ヘッドを結合 # (B, h, N, d) → (B, N, h, d) → (B, N, C) x = x.transpose(1, 2).reshape(B, N, C) # 出力の線形変換 x = self.proj(x) x = self.proj_dropout(x) return x

# =================================================== # 3. MLP（Feed-Forward Network） # =================================================== class MLP(nn.Module): “”” Feed-Forward Network 構造: Linear(D, 4D) → GELU → Dropout → Linear(4D, D) → Dropout 例: 768 → 3072 → 768 “”” def __init__(self, in_features, hidden_features=None, dropout=0.0): “”” Args: in_features: 入力次元（768） hidden_features: 隠れ層の次元（デフォルト: 入力の4倍 = 3072） dropout: ドロップアウト率 “”” super().__init__() # 隠れ層の次元（指定がなければ入力の4倍） hidden_features = hidden_features or in_features * 4 # 第1線形層: 768 → 3072 self.fc1 = nn.Linear(in_features, hidden_features) # 活性化関数: GELU（ReLUより滑らか） # GELU(x) = x × Φ(x) where Φ is the CDF of standard normal self.act = nn.GELU() # 第2線形層: 3072 → 768 self.fc2 = nn.Linear(hidden_features, in_features) # ドロップアウト self.dropout = nn.Dropout(dropout) def forward(self, x): “”” Args: x: (batch_size, num_tokens, embed_dim) Returns: (batch_size, num_tokens, embed_dim) “”” x = self.fc1(x) # 768 → 3072 x = self.act(x) # GELU活性化 x = self.dropout(x) # ドロップアウト x = self.fc2(x) # 3072 → 768 x = self.dropout(x) # ドロップアウト return x

# =================================================== # 4. Transformer Block（Encoder Layer） # =================================================== class TransformerBlock(nn.Module): “”” Transformer Encoder Block 構造: Input │ ├──→ LayerNorm → MHSA ─┐ │ │ └──────────────────────┼──→ Add（残差接続） │ ├──→ LayerNorm → MLP ─┐ │ │ └──────────────────────┼──→ Add（残差接続） │ Output “”” def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout=0.0): “”” Args: embed_dim: 埋め込み次元（768） num_heads: Attentionヘッド数（12） mlp_ratio: MLPの拡大率（4倍） dropout: ドロップアウト率 “”” super().__init__() # Layer Normalization 1（Attention前） self.norm1 = nn.LayerNorm(embed_dim) # Multi-Head Self-Attention self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout) # Layer Normalization 2（MLP前） self.norm2 = nn.LayerNorm(embed_dim) # MLP（Feed-Forward Network） mlp_hidden_dim = int(embed_dim * mlp_ratio) # 768 × 4 = 3072 self.mlp = MLP(embed_dim, mlp_hidden_dim, dropout) def forward(self, x): “”” Args: x: (batch_size, num_tokens, embed_dim) Returns: (batch_size, num_tokens, embed_dim) “”” # Attention + 残差接続 # Pre-Norm: 正規化を先に行う（オリジナルTransformerはPost-Norm） x = x + self.attn(self.norm1(x)) # MLP + 残差接続 x = x + self.mlp(self.norm2(x)) return x

# =================================================== # 5. Vision Transformer（完全版） # =================================================== class VisionTransformer(nn.Module): “”” Vision Transformer (ViT) 全体の流れ: 1. 画像をパッチに分割して埋め込み 2. [CLS]トークンを追加 3. 位置埋め込みを追加 4. Transformer Encoderで処理（12層） 5. [CLS]トークンの出力を分類ヘッドに入力 “”” def __init__( self, img_size=224, # 入力画像サイズ patch_size=16, # パッチサイズ in_channels=3, # 入力チャンネル数（RGB） num_classes=1000, # 分類クラス数（ImageNet） embed_dim=768, # 埋め込み次元 depth=12, # Transformer層数 num_heads=12, # Attentionヘッド数 mlp_ratio=4.0, # MLP拡大率 dropout=0.1 # ドロップアウト率 ): super().__init__() # パッチ埋め込み self.patch_embed = PatchEmbedding( img_size, patch_size, in_channels, embed_dim ) num_patches = self.patch_embed.num_patches # 196 # [CLS]トークン（学習可能なパラメータ） # 形状: (1, 1, 768) → バッチサイズ分複製して使用 self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) # 位置埋め込み（学習可能なパラメータ） # 形状: (1, 197, 768) → [CLS] + 196パッチ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) # ドロップアウト（位置埋め込み後） self.pos_drop = nn.Dropout(dropout) # Transformer Encoder（12層のブロック） self.blocks = nn.ModuleList([ TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout) for _ in range(depth) ]) # 最終的なLayer Normalization self.norm = nn.LayerNorm(embed_dim) # 分類ヘッド: 768 → 1000クラス self.head = nn.Linear(embed_dim, num_classes) # 重みの初期化 self._init_weights() def _init_weights(self): “””重みの初期化””” # 位置埋め込みと[CLS]トークンの初期化 # truncated normal: 正規分布から外れ値を切り捨て nn.init.trunc_normal_(self.pos_embed, std=0.02) nn.init.trunc_normal_(self.cls_token, std=0.02) # 他のパラメータの初期化 self.apply(self._init_module_weights) def _init_module_weights(self, m): “””モジュールごとの重み初期化””” if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward(self, x): “”” Args: x: 入力画像 (batch_size, 3, 224, 224) Returns: クラス確率 (batch_size, num_classes) “”” B = x.shape[0] # バッチサイズ # 1. パッチ埋め込み # (B, 3, 224, 224) → (B, 196, 768) x = self.patch_embed(x) # 2. [CLS]トークンを追加 # (1, 1, 768) → (B, 1, 768) cls_tokens = self.cls_token.expand(B, -1, -1) # (B, 1, 768) + (B, 196, 768) → (B, 197, 768) x = torch.cat((cls_tokens, x), dim=1) # 3. 位置埋め込みを追加 # (B, 197, 768) + (1, 197, 768) → (B, 197, 768) x = x + self.pos_embed x = self.pos_drop(x) # 4. Transformer Encoder（12層） for block in self.blocks: x = block(x) # 5. Layer Normalization x = self.norm(x) # 6. [CLS]トークンの出力を取得 # (B, 197, 768) → (B, 768) cls_token_final = x[:, 0] # 7. 分類ヘッド # (B, 768) → (B, num_classes) out = self.head(cls_token_final) return out

# =================================================== # 6. 使用例 # =================================================== # ViT-Base/16 モデルを作成 model = VisionTransformer( img_size=224, patch_size=16, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0 ) # パラメータ数を確認 total_params = sum(p.numel() for p in model.parameters()) print(f”Total parameters: {total_params:,}”)

# 推論テスト # ダミーの入力画像を作成 x = torch.randn(2, 3, 224, 224) # バッチサイズ2 # 推論 model.eval() # 評価モード with torch.no_grad(): output = model(x) print(f”Input shape: {x.shape}”) print(f”Output shape: {output.shape}”)

# =================================================== # 7. 事前学習済みモデルの使用（timmライブラリ） # =================================================== # timmのインストール（Google Colabの場合） # !pip install timm import timm # 利用可能なViTモデルを確認 vit_models = timm.list_models(‘vit*’, pretrained=True) print(“利用可能なViTモデル（一部）:”) for model_name in vit_models[:10]: print(f” {model_name}”)

# 事前学習済みViT-Base/16をロード model_pretrained = timm.create_model( ‘vit_base_patch16_224’, # モデル名 pretrained=True # ImageNetで事前学習済み ) model_pretrained.eval() # モデル情報を確認 print(f”Model: vit_base_patch16_224″) print(f”Number of parameters: {sum(p.numel() for p in model_pretrained.parameters()):,}”)

# 推論の実行 import torch from PIL import Image import requests from io import BytesIO # サンプル画像をダウンロード（例: 猫の画像） url = “https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg” response = requests.get(url) img = Image.open(BytesIO(response.content)).convert(‘RGB’) # 前処理 # timmの設定に合わせた前処理を取得 data_config = timm.data.resolve_model_data_config(model_pretrained) transforms = timm.data.create_transform(**data_config, is_training=False) # 画像を前処理 img_tensor = transforms(img).unsqueeze(0) # バッチ次元を追加 # 推論 with torch.no_grad(): output = model_pretrained(img_tensor) probabilities = torch.softmax(output, dim=1) # Top-5予測を取得 top5_prob, top5_idx = torch.topk(probabilities, 5) # ImageNetのクラス名を取得（簡易版） # 実際にはimagenet_classes.txtなどから読み込む print(“\nTop-5 predictions:”) for i in range(5): print(f” Class {top5_idx[0][i].item()}: {top5_prob[0][i].item()*100:.2f}%”)