Transformer フレームワークに基づく Segmenter モデル | 学習

自分で作った小さなモデルです — 大したものではありません。

はじめに & プロジェクト紹介

このプロジェクトには半学期近くかかりました。最初はチュートリアルを見ながらTransformerの各コンポーネント（自己注意機構、Embedding、位置エンコーディングなど）を学んでいましたが、それらを組み合わせる方法がわからなかったんです。この機会を活かして、誰でも作れる小さなモデルに統合しました。

初心者レベルですが、それでも結構満足しています。

個人リポジトリ：cateude/transformer_segment

自分で良いと思う点：

ドキュメント習慣がまあまあ（CLAUDE.mdを作っただけだけど）
PASCAL VOC 2012を使用 — データソースがしっかりしていて、train/val を明確に分離（テストセットにトレーニングセットのパラメータを混ぜてない — それはダメだからね）
インターフェースが整然としており、パラメータ制御可能、ログ対応 — ハードコードされたスクリプトより追跡しやすい
デバイス処理が適切 — まあこんなところ

しかし欠点も結構あります：

バリデーションは最初の20バッチのみ（すみません、私のPCではこれが限界です）
データパイプラインは理想的な入力を前提としている — エッジケースを考慮していなかった（PS：私のミス）
Config ですべてのパラメータを管理しているが、他のコードブロックでもクラス定義内にパラメータがあるため、チューニングがやや面倒

プロジェクトのフレームワーク

プロジェクト詳細

モデルのブランチ構築：ViT と Decoder

まず ViT と decoder の2つのファイルを構築 — 一つはエンコード用、もう一つはデコード用。

def __init__(self, num_patches: int = 576, embed_dim: int = 384, depth: int = 8,
             num_heads: int = 6, patch_size: int = 16, pretrained: bool = True):
    super().__init__()
    if embed_dim not in _TIMM_NAME:
        raise ValueError(f"embed_dim={embed_dim} に対応するtimmモデルがありません。サポート: {list(_TIMM_NAME)}")
    image_size = int(num_patches ** 0.5) * patch_size

    self.backbone = timm.create_model(
        _TIMM_NAME[embed_dim],
        pretrained=pretrained,
        num_classes=0,
        img_size=image_size,
    )
    self.num_prefix_tokens = self.backbone.num_prefix_tokens

def forward(self, x):
    tokens = self.backbone.forward_features(x)
    return tokens[:, self.num_prefix_tokens:, :]

PS：トレーニング効果のため、やむを得ず timm モジュールを導入しました — しくしく！

モデルの組み合わせ：model.py

model.py でこれら2つを接続し、学習可能な温度パラメータを追加。

class Segmenter(nn.Module):
    def __init__(self, num_classes = 21, image_size = 384, patch_size = 16,
                 embed_dim = 768, encoder_depth = 12, encoder_heads = 12,
                 decoder_depth = 4, decoder_heads = 12):
        super().__init__()
        self.num_classes = num_classes
        self.image_size = image_size
        self.patch_size = patch_size
        num_patches = (image_size // patch_size) ** 2
        self.encoder = VIT(num_patches = num_patches, embed_dim = embed_dim, depth = encoder_depth, num_heads = encoder_heads, patch_size = patch_size)
        self.decoder = MaskTransformer(embed_dim = embed_dim, num_heads = decoder_heads, num_classes = num_classes, depth = decoder_depth)
        self.logit_scale = nn.Parameter(torch.log(torch.tensor(1.0 / 0.07)))

    def forward(self, x):
        B = x.shape[0]
        patch_tokens = self.encoder(x)
        cls_embedding = self.decoder(patch_tokens)

        patch_tokens = F.normalize(patch_tokens, dim = -1)
        cls_embedding = F.normalize(cls_embedding, dim = -1)
        logit_scale = self.logit_scale.exp().clamp(max = 100.0)
        logits = logit_scale * (patch_tokens @ cls_embedding.transpose(1, 2))
        logits = logits.permute(0, 2, 1)

        H_patch = W_patch = x.shape[2] // self.patch_size
        logits = logits.reshape(B, self.num_classes, H_patch, W_patch)

        logits = F.interpolate(logits, size = (self.image_size, self.image_size), mode = 'bilinear', align_corners = False)

        return logits

データ処理：Dataset と Transforms

データセットは訓練セットとテストセットを分離し、画像読み込みを含む。

Transforms はいわゆる三種の神器（Resize、ToTensor、Normalize）で、訓練セットには RandomCrop と RandomHorizontalFlip を追加。

class Resize:
    def __init__(self, size):
        self.size = size

    def __call__(self, sample):
        image = sample["image"]
        mask = sample["mask"]
        image = F.resize(image, self.size, InterpolationMode.BILINEAR)
        mask = F.resize(mask, self.size, InterpolationMode.NEAREST) if mask is not None else None
        return {"image": image, "mask": mask}

class ToTensor:
    def __call__(self, sample):
        image = sample["image"]
        mask = sample["mask"]
        image = F.to_tensor(image)
        mask = torch.as_tensor(np.array(mask), dtype=torch.long) if mask is not None else None
        return {"image": image, "mask": mask}

class Normalize:
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, sample):
        image = sample["image"]
        mask = sample["mask"]
        image = F.normalize(image, mean=self.mean, std=self.std)
        return {"image": image, "mask": mask}

損失関数と評価：Losses と Metrics

混同しないでね — losses はモデルがトレーニング中に評価するためのもの。Metrics は画面の前のあなたが検証・テスト時に確認するためのもの。

単一ステップ学習：train_step

1ステップの学習：フォワード、バックワード、オプティマイザ更新 — 古典的な学習の5ステップ。

def train_step(model, images, masks, optimizer, scaler, criterion) -> float:
    optimizer.zero_grad()
    device_type = images.device.type

    with torch.amp.autocast(device_type):
        logits = model(images)
        loss = criterion(logits, masks)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    return loss.item()

トレーニング：train

すべてのモジュールを統合してモデルをトレーニング。

結果：predict

最終的な結果を表示。

実行方法

簡単です：

uv パッケージをインストール：pip install uv（またはより高速な方法）
お使いのGPUに合ったパッケージをインストール（PS：これに丸一日ハマりました — 本当に絶望的だった）
ターミナルで uv run python train.py を実行
トレーニングが完了するのを静かに待つ（約1時間）
ターミナルで uv run python predict.py を実行 — output フォルダに画像が生成される

まとめ

ディープラーニングを始めたばかりの初心者として、チュートリアルを見るだけから実際にモデルを再現して構築するレベルに成長できました (⊙﹏⊙)（Ciallo～(∠・ω< )⌒★）