Vision Transformerを一から実装してFoodVision Miniに適用する¶

# 必要なバージョンの確認とインストール
try:
    import torch
    import torchvision
    assert int(torch.__version__.split(".")[1]) >= 12 or int(torch.__version__.split(".")[0]) == 2, "torch version should be 1.12+"
    assert int(torchvision.__version__.split(".")[1]) >= 13, "torchvision version should be 0.13+"
    print(f"torch version: {torch.__version__}")
    print(f"torchvision version: {torchvision.__version__}")
except:
    print(f"[INFO] torch/torchvision versions not as required, installing nightly versions.")
    !pip3 install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    import torch
    import torchvision
    print(f"torch version: {torch.__version__}")
    print(f"torchvision version: {torchvision.__version__}")

# 必要なライブラリのインポート
import matplotlib.pyplot as plt
import torch
import torchvision

from torch import nn
from torchvision import transforms
from torchinfo import summary
from going_modular import data_setup, engine
from helper_functions import download_data, set_seeds, plot_loss_curves

# デバイス設定（GPU利用可能性の確認）
device = "mps" if torch.mps.is_available() else "cpu"
device

# GitHubからピザ、ステーキ、寿司画像をダウンロード
image_path = download_data(source="https://github.com/vinsmoke-three/deeplearning-with-pytorch/raw/main/data/pizza_steak_sushi.zip",
                           destination="pizza_steak_sushi")
image_path

# 訓練用とテスト用ディレクトリのパスを設定
train_dir = image_path / "train"
test_dir = image_path / "test"

# 画像サイズの設定（ViT論文のTable 3から）
IMG_SIZE = 224

# 手動で変換パイプラインを作成
manual_transforms = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),  # 224x224にリサイズ
    transforms.ToTensor(),                    # テンソルに変換
])
print(f"手動作成された変換: {manual_transforms}")

# バッチサイズの設定（ViT論文より小さく設定）
BATCH_SIZE = 32  # ViT論文より小さいが、ハードウェア制約のため

# データローダーの作成
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
    train_dir=train_dir,
    test_dir=test_dir,
    transform=manual_transforms,  # 手動作成した変換を使用
    batch_size=BATCH_SIZE
)

train_dataloader, test_dataloader, class_names

# 画像バッチを取得
image_batch, label_batch = next(iter(train_dataloader))

# バッチから単一の画像を取得
image, label = image_batch[0], label_batch[0]

# バッチの形状を確認
image.shape, label

# matplotlibで画像をプロット
plt.imshow(image.permute(1, 2, 0))  # 画像次元を調整 [color_channels, height, width] -> [height, width, color_channels]
plt.title(class_names[label])
plt.axis(False);

# サンプル値を作成
height = 224 # H ("Training resolution is 224.")
width = 224 # W
color_channels = 3 # C
patch_size = 16 # P

# N（パッチ数）を計算
number_of_patches = int((height * width) / patch_size**2)
print(f"パッチサイズ(P={patch_size})における画像高さ(H={height})、幅(W={width})でのパッチ数(N): {number_of_patches}")

# 入力形状（単一画像のサイズ）
embedding_layer_input_shape = (height, width, color_channels)

# 出力形状
embedding_layer_output_shape = (number_of_patches, patch_size**2 * color_channels)

print(f"入力形状（単一2D画像）: {embedding_layer_input_shape}")
print(f"出力形状（パッチに平坦化された単一2D画像）: {embedding_layer_output_shape}")

# 単一画像を表示
plt.imshow(image.permute(1, 2, 0)) # matplotlibに合わせて調整
plt.title(class_names[label])
plt.axis(False)

# matplotlib用に画像形状を変更 [color_channels, height, width] -> [height, width, color_channels]
image_permuted = image.permute(1, 2, 0)

# 上部行のパッチ化されたピクセルをプロットするためのインデックス
patch_size = 16
plt.figure(figsize=(patch_size, patch_size))
plt.imshow(image_permuted[:patch_size, :, :])

# ハイパーパラメータの設定とimg_sizeとpatch_sizeの互換性確認
img_size = 224
patch_size = 16
num_patches = img_size/patch_size
assert img_size % patch_size == 0, "画像サイズはパッチサイズで割り切れる必要があります"
print(f"行ごとのパッチ数: {num_patches}\nパッチサイズ: {patch_size} pixels x {patch_size} pixels")

# 一連のサブプロットを作成
fig, axs = plt.subplots(nrows=1,
                        ncols=img_size // patch_size, # 各パッチに1列
                        figsize=(num_patches, num_patches),
                        sharex=True,
                        sharey=True)

# 上部行のパッチ数を反復処理
for i, patch in enumerate(range(0, img_size, patch_size)):
    axs[i].imshow(image_permuted[:patch_size, patch:patch+patch_size, :]) # 高さインデックスを一定に保ち、幅インデックスを変更
    axs[i].set_xlabel(i+1) # ラベルを設定
    axs[i].set_xticks([])
    axs[i].set_yticks([])

# ハイパーパラメータの設定
img_size = 224
patch_size = 16
num_patches = img_size/patch_size
print(f"行ごとのパッチ数: {num_patches}\
        \n列ごとのパッチ数: {num_patches}\
        \n総パッチ数: {num_patches*num_patches}\
        \nパッチサイズ: {patch_size} pixels x {patch_size} pixels")

# サブプロットシリーズを作成
fig, axs = plt.subplots(nrows=img_size // patch_size,
                        ncols=img_size // patch_size,
                        figsize=(num_patches, num_patches),
                        sharex=True,
                        sharey=True)

# 画像の高さと幅をループ
for i, patch_height in enumerate(range(0, img_size, patch_size)): # 高さを反復
    for j, patch_width in enumerate(range(0, img_size, patch_size)): # 幅を反復
        # 置換された画像パッチをプロット
        axs[i, j].imshow(image_permuted[patch_height:patch_height+patch_size,
                                        patch_width:patch_width+patch_size,
                                        :])
        # ラベル情報を設定
        axs[i, j].set_ylabel(i+1,
                             rotation="horizontal",
                             horizontalalignment="right",
                             verticalalignment="center")
        axs[i, j].set_xlabel(j+1)
        axs[i, j].set_xticks([])
        axs[i, j].set_yticks([])
        axs[i, j].label_outer()

# スーパータイトルを設定
fig.suptitle(f"{class_names[label]} -> パッチ化済み", fontsize=16)
plt.show()

from torch import nn

# パッチサイズを設定
patch_size=16

# ViT論文のハイパーパラメータでConv2d層を作成
conv2d = nn.Conv2d(in_channels=3, # カラーチャンネル数
                   out_channels=768, # Table 1より: Hidden size D、これが埋め込みサイズ
                   kernel_size=patch_size,
                   stride=patch_size,
                   padding=0)

# 画像を畳み込み層に通す
image_out_of_conv = conv2d(image.unsqueeze(0)) # 単一バッチ次元を追加
print(image_out_of_conv.shape)

# ランダムな5つの畳み込み特徴マップをプロット
import random
random_indexes = random.sample(range(0, 758), k=5)
print(f"インデックスからのランダムな畳み込み特徴マップを表示: {random_indexes}")

# プロット作成
fig, axs = plt.subplots(nrows=1, ncols=5, figsize=(12, 12))

# ランダムな画像特徴マップをプロット
for i, idx in enumerate(random_indexes):
    image_conv_feature_map = image_out_of_conv[:, idx, :, :]
    axs[i].imshow(image_conv_feature_map.squeeze().detach().numpy())
    axs[i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[]);

# 現在のテンソル形状
print(f"現在のテンソル形状: {image_out_of_conv.shape} -> [batch, embedding_dim, feature_map_height, feature_map_width]")

# 平坦化層を作成
flatten = nn.Flatten(start_dim=2, # feature_map_height（次元2）を平坦化
                     end_dim=3) # feature_map_width（次元3）を平坦化

# 1. 単一画像を表示
plt.imshow(image.permute(1, 2, 0)) # matplotlib用に調整
plt.title(class_names[label])
plt.axis(False)
print(f"元画像の形状: {image.shape}")

# 2. 画像を特徴マップに変換
image_out_of_conv = conv2d(image.unsqueeze(0)) # 形状エラーを避けるためバッチ次元を追加
print(f"画像特徴マップの形状: {image_out_of_conv.shape}")

# 3. 特徴マップを平坦化
image_out_of_conv_flattened = flatten(image_out_of_conv)
print(f"平坦化された画像特徴マップの形状: {image_out_of_conv_flattened.shape}")

# 平坦化された画像パッチ埋め込みを正しい形状に取得
image_out_of_conv_flattened_reshaped = image_out_of_conv_flattened.permute(0, 2, 1) # [batch_size, P^2•C, N] -> [batch_size, N, P^2•C]
print(f"パッチ埋め込みシーケンス形状: {image_out_of_conv_flattened_reshaped.shape} -> [batch_size, num_patches, embedding_size]")

# 1. nn.Moduleをサブクラス化するクラスを作成
class PatchEmbedding(nn.Module):
    """2D入力画像を1D学習可能埋め込みベクトルシーケンスに変換します。

    Args:
        in_channels (int): 入力画像のカラーチャンネル数。デフォルトは3。
        patch_size (int): 入力画像を変換するパッチのサイズ。デフォルトは16。
        embedding_dim (int): 画像を変換する埋め込みのサイズ。デフォルトは768。
    """
    # 2. 適切な変数でクラスを初期化
    def __init__(self,
                 in_channels:int=3,
                 patch_size:int=16,
                 embedding_dim:int=768):
        super().__init__()

        # 3. 画像をパッチに変換する層を作成
        self.patcher = nn.Conv2d(in_channels=in_channels,
                                 out_channels=embedding_dim,
                                 kernel_size=patch_size,
                                 stride=patch_size,
                                 padding=0)

        # 4. パッチ特徴マップを単一次元に平坦化する層を作成
        self.flatten = nn.Flatten(start_dim=2, # 特徴マップ次元のみを単一ベクトルに平坦化
                                  end_dim=3)

    # 5. forwardメソッドを定義
    def forward(self, x):
        # 入力が正しい形状かチェックするアサーションを作成
        image_resolution = x.shape[-1]
        assert image_resolution % patch_size == 0, f"入力画像サイズはパッチサイズで割り切れる必要があります。画像形状: {image_resolution}、パッチサイズ: {patch_size}"

        # フォワードパスを実行
        x_patched = self.patcher(x)
        x_flattened = self.flatten(x_patched)
        # 6. 出力形状が正しい順序になるようにする
        return x_flattened.permute(0, 2, 1) # 埋め込みが最後の次元になるように調整

set_seeds()

# パッチ埋め込み層のインスタンスを作成
patchify = PatchEmbedding(in_channels=3,
                          patch_size=16,
                          embedding_dim=768)

# 単一画像を通す
print(f"入力画像形状: {image.unsqueeze(0).shape}")
patch_embedded_image = patchify(image.unsqueeze(0)) # エラーを避けるため0番目のインデックスに追加のバッチ次元を追加
print(f"出力パッチ埋め込み形状: {patch_embedded_image.shape}")

# パッチ埋め込みとパッチ埋め込み形状を表示
print(patch_embedded_image)
print(f"パッチ埋め込み形状: {patch_embedded_image.shape} -> [batch_size, number_of_patches, embedding_dimension]")

# バッチサイズと埋め込み次元を取得
batch_size = patch_embedded_image.shape[0]
embedding_dimension = patch_embedded_image.shape[-1]

# 埋め込み次元と同じサイズを共有する学習可能パラメータとしてクラストークン埋め込みを作成
class_token = nn.Parameter(torch.ones(batch_size, 1, embedding_dimension), # [batch_size, number_of_tokens, embedding_dimension]
                           requires_grad=True) # 埋め込みが学習可能であることを確認

# class_tokenの最初の10例を表示
print(class_token[:, :, :10])

# class_token形状を出力
print(f"クラストークン形状: {class_token.shape} -> [batch_size, number_of_tokens, embedding_dimension]")

# クラストークン埋め込みをパッチ埋め込みの前面に追加
patch_embedded_image_with_class_embedding = torch.cat((class_token, patch_embedded_image),
                                                      dim=1) # 最初の次元で連結

print(f"クラストークンが前に付けられたパッチ埋め込みのシーケンス形状: {patch_embedded_image_with_class_embedding.shape}")

# N（パッチ数）を計算
number_of_patches = int((height * width) / patch_size**2)

# 埋め込み次元を取得
embedding_dimension = patch_embedded_image_with_class_embedding.shape[2]

# 学習可能な1D位置埋め込みを作成
position_embedding = nn.Parameter(torch.ones(1,
                                             number_of_patches+1,
                                             embedding_dimension),
                                  requires_grad=True) # 学習可能であることを確認

print(f"位置埋め込み形状: {position_embedding.shape} -> [batch_size, number_of_patches, embedding_dimension]")

# 位置埋め込みをパッチおよびクラストークン埋め込みに追加
patch_and_position_embedding = patch_embedded_image_with_class_embedding + position_embedding
print(f"パッチ埋め込み、クラストークン前付け、位置埋め込み追加の形状: {patch_and_position_embedding.shape}")

set_seeds()

# 1. パッチサイズを設定
patch_size = 16

# 2. 元画像テンソル形状を出力し、画像次元を取得
print(f"画像テンソル形状: {image.shape}")
height, width = image.shape[1], image.shape[2]

# 3. 画像テンソルを取得してバッチ次元を追加
x = image.unsqueeze(0)
print(f"バッチ次元付き入力画像形状: {x.shape}")

# 4. パッチ埋め込み層を作成
patch_embedding_layer = PatchEmbedding(in_channels=3,
                                       patch_size=patch_size,
                                       embedding_dim=768)

# 5. パッチ埋め込み層を通して画像を渡す
patch_embedding = patch_embedding_layer(x)
print(f"パッチ化埋め込み形状: {patch_embedding.shape}")

# 6. クラストークン埋め込みを作成
batch_size = patch_embedding.shape[0]
embedding_dimension = patch_embedding.shape[-1]
class_token = nn.Parameter(torch.ones(batch_size, 1, embedding_dimension),
                           requires_grad=True) # 学習可能であることを確認
print(f"クラストークン埋め込み形状: {class_token.shape}")

# 7. クラストークン埋め込みをパッチ埋め込みの前に追加
patch_embedding_class_token = torch.cat((class_token, patch_embedding), dim=1)
print(f"クラストークン付きパッチ埋め込み形状: {patch_embedding_class_token.shape}")

# 8. 位置埋め込みを作成
number_of_patches = int((height * width) / patch_size**2)
position_embedding = nn.Parameter(torch.ones(1, number_of_patches+1, embedding_dimension),
                                  requires_grad=True) # 学習可能であることを確認

# 9. 位置埋め込みをクラストークン付きパッチ埋め込みに追加
patch_and_position_embedding = patch_embedding_class_token + position_embedding
print(f"パッチと位置埋め込み形状: {patch_and_position_embedding.shape}")

# 1. nn.Moduleを継承するクラスを作成
class MultiheadSelfAttentionBlock(nn.Module):
    """マルチヘッド自己注意ブロック（「MSAブロック」の略）を作成します。
    """
    # 2. Table 1のハイパーパラメータでクラスを初期化
    def __init__(self,
                 embedding_dim:int=768, # Table 1のViT-BaseのHidden size D
                 num_heads:int=12, # Table 1のViT-BaseのHeads
                 attn_dropout:float=0): # 論文ではMSABlocksでドロップアウトを使用していないようです
        super().__init__()

        # 3. Norm層（LN）を作成
        self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim)

        # 4. Multi-Head Attention（MSA）層を作成
        self.multihead_attn = nn.MultiheadAttention(embed_dim=embedding_dim,
                                                    num_heads=num_heads,
                                                    dropout=attn_dropout,
                                                    batch_first=True) # バッチ次元が最初に来るか？

    # 5. データを層に通すforward()メソッドを作成
    def forward(self, x):
        x = self.layer_norm(x)
        attn_output, _ = self.multihead_attn(query=x, # クエリ埋め込み
                                             key=x, # キー埋め込み
                                             value=x, # バリュー埋め込み
                                             need_weights=False) # 重みが必要か、層の出力だけか？
        return attn_output

# MSABlockのインスタンスを作成
multihead_self_attention_block = MultiheadSelfAttentionBlock(embedding_dim=768, # Table 1より
                                                             num_heads=12) # Table 1より

# パッチと位置画像埋め込みをMSABlockに通す
patched_image_through_msa_block = multihead_self_attention_block(patch_and_position_embedding)
print(f"MSAブロックの入力形状: {patch_and_position_embedding.shape}")
print(f"MSAブロックの出力形状: {patched_image_through_msa_block.shape}")

# 1. nn.Moduleを継承するクラスを作成
class MLPBlock(nn.Module):
    """層正規化多層パーセプトロンブロック（「MLPブロック」の略）を作成します。"""
    # 2. Table 1とTable 3のハイパーパラメータでクラスを初期化
    def __init__(self,
                 embedding_dim:int=768, # Table 1のViT-BaseのHidden Size D
                 mlp_size:int=3072, # Table 1のViT-BaseのMLP size
                 dropout:float=0.1): # Table 3のViT-BaseのDropout
        super().__init__()

        # 3. Norm層（LN）を作成
        self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim)

        # 4. Multilayer perceptron（MLP）層を作成
        self.mlp = nn.Sequential(
            nn.Linear(in_features=embedding_dim,
                      out_features=mlp_size),
            nn.GELU(), # "MLPには、GELU非線形性を持つ2つの層が含まれています（3.1節）。"
            nn.Dropout(p=dropout),
            nn.Linear(in_features=mlp_size, # 上の層のout_featuresと同じin_featuresを取る必要があります
                      out_features=embedding_dim), # embedding_dimに戻す
            nn.Dropout(p=dropout) # "ドロップアウトが使用される場合、すべての密結合層の後に適用されます.."
        )

    # 5. データを層に通すforward()メソッドを作成
    def forward(self, x):
        x = self.layer_norm(x)
        x = self.mlp(x)
        return x

# MLPBlockのインスタンスを作成
mlp_block = MLPBlock(embedding_dim=768, # Table 1より
                     mlp_size=3072, # Table 1より
                     dropout=0.1) # Table 3より

# MSABlockの出力をMLPBlockに通す
patched_image_through_mlp_block = mlp_block(patched_image_through_msa_block)
print(f"MLPブロックの入力形状: {patched_image_through_msa_block.shape}")
print(f"MLPブロックの出力形状: {patched_image_through_mlp_block.shape}")

# 1. nn.Moduleを継承するクラスを作成
class TransformerEncoderBlock(nn.Module):
    """Transformer Encoderブロックを作成します。"""
    # 2. Table 1とTable 3のハイパーパラメータでクラスを初期化
    def __init__(self,
                 embedding_dim:int=768, # Table 1のViT-BaseのHidden size D
                 num_heads:int=12, # Table 1のViT-BaseのHeads
                 mlp_size:int=3072, # Table 1のViT-BaseのMLP size
                 mlp_dropout:float=0.1, # Table 3のViT-Baseの密結合層のドロップアウト量
                 attn_dropout:float=0): # アテンション層のドロップアウト量
        super().__init__()

        # 3. MSAブロック（数式2）を作成
        self.msa_block = MultiheadSelfAttentionBlock(embedding_dim=embedding_dim,
                                                     num_heads=num_heads,
                                                     attn_dropout=attn_dropout)

        # 4. MLPブロック（数式3）を作成
        self.mlp_block =  MLPBlock(embedding_dim=embedding_dim,
                                   mlp_size=mlp_size,
                                   dropout=mlp_dropout)

    # 5. forward()メソッドを作成
    def forward(self, x):
        # 6. MSAブロック用の残差接続（入力を出力に追加）
        x =  self.msa_block(x) + x

        # 7. MLPブロック用の残差接続（入力を出力に追加）
        x = self.mlp_block(x) + x

        return x

# TransformerEncoderBlockのインスタンスを作成
transformer_encoder_block = TransformerEncoderBlock()

# 入力形状(1, 197, 768) -> (batch_size, num_patches, embedding_dimension)でのサマリーを表示
summary(model=transformer_encoder_block,
        input_size=(1, 197, 768),
        col_names=["input_size", "output_size", "num_params", "trainable"],
        col_width=20,
        row_settings=["var_names"])

# torch.nn.TransformerEncoderLayer()で上記と同じものを作成
torch_transformer_encoder_layer = nn.TransformerEncoderLayer(d_model=768, # Table 1のViT-BaseのHidden size D
                                                             nhead=12, # Table 1のViT-BaseのHeads
                                                             dim_feedforward=3072, # Table 1のViT-BaseのMLP size
                                                             dropout=0.1, # Table 3のViT-Baseの密結合層のドロップアウト量
                                                             activation="gelu", # GELU非線形活性化
                                                             batch_first=True, # バッチが最初に来るか？
                                                             norm_first=True) # MSA/MLP層の前または後に正規化するか？

torch_transformer_encoder_layer

# 1. nn.Moduleを継承するViTクラスを作成
class ViT(nn.Module):
    """デフォルトでViT-Baseハイパーパラメータを持つVision Transformerアーキテクチャを作成します。"""
    # 2. Table 1とTable 3のハイパーパラメータでクラスを初期化
    def __init__(self,
                 img_size:int=224, # ViT論文のTable 3のTraining resolution
                 in_channels:int=3, # 入力画像のチャンネル数
                 patch_size:int=16, # パッチサイズ
                 num_transformer_layers:int=12, # Table 1のViT-BaseのLayers
                 embedding_dim:int=768, # Table 1のViT-BaseのHidden size D
                 mlp_size:int=3072, # Table 1のViT-BaseのMLP size
                 num_heads:int=12, # Table 1のViT-BaseのHeads
                 attn_dropout:float=0, # アテンション射影のドロップアウト
                 mlp_dropout:float=0.1, # 密結合/MLP層のドロップアウト
                 embedding_dropout:float=0.1, # パッチと位置埋め込みのドロップアウト
                 num_classes:int=1000): # ImageNetのデフォルトですがカスタマイズ可能
        super().__init__() # super().__init__()を忘れずに！

        # 3. 画像サイズがパッチサイズで割り切れることを確認
        assert img_size % patch_size == 0, f"画像サイズはパッチサイズで割り切れる必要があります。画像サイズ: {img_size}、パッチサイズ: {patch_size}。"

        # 4. パッチ数を計算（高さ * 幅/パッチ^2）
        self.num_patches = (img_size * img_size) // patch_size**2

        # 5. 学習可能なクラス埋め込みを作成（パッチ埋め込みのシーケンスの前に配置する必要があります）
        self.class_embedding = nn.Parameter(data=torch.randn(1, 1, embedding_dim),
                                            requires_grad=True)

        # 6. 学習可能な位置埋め込みを作成
        self.position_embedding = nn.Parameter(data=torch.randn(1, self.num_patches+1, embedding_dim),
                                               requires_grad=True)

        # 7. 埋め込みドロップアウト値を作成
        self.embedding_dropout = nn.Dropout(p=embedding_dropout)

        # 8. パッチ埋め込み層を作成
        self.patch_embedding = PatchEmbedding(in_channels=in_channels,
                                              patch_size=patch_size,
                                              embedding_dim=embedding_dim)

        # 9. Transformer Encoderブロックを作成（nn.Sequential()を使用してTransformer Encoderブロックをスタック）
        # 注："*"は"すべて"を意味します
        self.transformer_encoder = nn.Sequential(*[TransformerEncoderBlock(embedding_dim=embedding_dim,
                                                                            num_heads=num_heads,
                                                                            mlp_size=mlp_size,
                                                                            mlp_dropout=mlp_dropout) for _ in range(num_transformer_layers)])

        # 10. 分類器ヘッドを作成
        self.classifier = nn.Sequential(
            nn.LayerNorm(normalized_shape=embedding_dim),
            nn.Linear(in_features=embedding_dim,
                      out_features=num_classes)
        )

    # 11. forward()メソッドを作成
    def forward(self, x):
        # 12. バッチサイズを取得
        batch_size = x.shape[0]

        # 13. クラストークン埋め込みを作成してバッチサイズに合わせて拡張（数式1）
        class_token = self.class_embedding.expand(batch_size, -1, -1) # "-1"は次元を推論することを意味

        # 14. パッチ埋め込みを作成（数式1）
        x = self.patch_embedding(x)

        # 15. クラス埋め込みとパッチ埋め込みを連結（数式1）
        x = torch.cat((class_token, x), dim=1)

        # 16. 位置埋め込みをパッチ埋め込みに追加（数式1）
        x = self.position_embedding + x

        # 17. 埋め込みドロップアウトを実行（付録B.1）
        x = self.embedding_dropout(x)

        # 18. パッチ、位置、クラス埋め込みをtransformer encoderレイヤーに通す（数式2&3）
        x = self.transformer_encoder(x)

        # 19. 0インデックスのロジットを分類器に通す（数式4）
        x = self.classifier(x[:, 0]) # バッチの各サンプルで0インデックスで実行

        return x

set_seeds()

# 単一画像と同じ形状のランダムテンソルを作成
random_image_tensor = torch.randn(1, 3, 224, 224) # (batch_size, color_channels, height, width)

# 作業中のクラス数（pizza、steak、sushi）でViTのインスタンスを作成
vit = ViT(num_classes=len(class_names))

# ランダム画像テンソルをViTインスタンスに渡す
vit(random_image_tensor)

from going_modular import engine

# ViT論文のハイパーパラメータを使用してViTモデルパラメータを最適化するオプティマイザーを設定
optimizer = torch.optim.Adam(params=vit.parameters(),
                             lr=3e-3, # Table 3のViT-* ImageNet-1kのBase LR
                             betas=(0.9, 0.999), # デフォルト値、ViT論文4.1節にも記載
                             weight_decay=0.3) # ViT論文4.1節とTable 3のViT-* ImageNet-1kより

# 多クラス分類用の損失関数を設定
loss_fn = torch.nn.CrossEntropyLoss()

# シードを設定
set_seeds()

# モデルを訓練し、訓練結果を辞書に保存
results = engine.train(model=vit.to(device),
                       train_dataloader=train_dataloader,
                       test_dataloader=test_dataloader,
                       optimizer=optimizer,
                       loss_fn=loss_fn,
                       epochs=10,
                       device=device,
                       writer=None)

# 1. ViT-Baseの事前訓練済み重みを取得
pretrained_vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT # torchvision >= 0.13が必要

# 2. 事前訓練済み重みでViTモデルインスタンスを設定
pretrained_vit = torchvision.models.vit_b_16(weights=pretrained_vit_weights).to(device)

# 3. ベースパラメータをフリーズ
for parameter in pretrained_vit.parameters():
    parameter.requires_grad = False

# 4. 分類器ヘッドを変更（線形ヘッドで同じ初期化を保証するためシードを設定）
set_seeds()
pretrained_vit.heads = nn.Linear(in_features=768, out_features=len(class_names)).to(device)

# 事前訓練済みViT重みから自動変換を取得
pretrained_vit_transforms = pretrained_vit_weights.transforms()
print(pretrained_vit_transforms)

# データローダーを設定
train_dataloader_pretrained, test_dataloader_pretrained, class_names = data_setup.create_dataloaders(
    train_dir=train_dir,
    test_dir=test_dir,
    transform=pretrained_vit_transforms,
    batch_size=32
)

from going_modular import engine

# オプティマイザーと損失関数を作成
optimizer = torch.optim.Adam(params=pretrained_vit.parameters(),
                             lr=1e-3)
loss_fn = torch.nn.CrossEntropyLoss()

# 事前訓練済みViT特徴抽出器モデルの分類器ヘッドを訓練
set_seeds()
pretrained_vit_results = engine.train(model=pretrained_vit,
                                      train_dataloader=train_dataloader_pretrained,
                                      test_dataloader=test_dataloader_pretrained,
                                      optimizer=optimizer,
                                      loss_fn=loss_fn,
                                      epochs=10,
                                      device=device,
                                      writer=None)

# 損失曲線をプロット
from helper_functions import plot_loss_curves

plot_loss_curves(pretrained_vit_results)

# モデルを保存
from going_modular import utils

utils.save_model(model=pretrained_vit,
                 target_dir="models",
                 model_name="09_pretrained_vit_feature_extractor_pizza_steak_sushi.pth")

from pathlib import Path

# モデルサイズをバイトで取得してメガバイトに変換
pretrained_vit_model_size = Path("models/09_pretrained_vit_feature_extractor_pizza_steak_sushi.pth").stat().st_size // (1024*1024)
print(f"事前訓練済みViT特徴抽出器モデルサイズ: {pretrained_vit_model_size} MB")

# 画像で予測を行い、プロットする関数をインポート
from going_modular.predictions import pred_and_plot_image

# カスタム画像パスを設定
custom_image_path = Path("data/sushi.jpg")

# カスタム画像で予測
pred_and_plot_image(model=pretrained_vit,
                    image_path=custom_image_path,
                    class_names=class_names)