PyTorchモデルのデプロイメント - FoodVision Bigの構築とHugging Face Spacesへの公開¶

# 必要なライブラリのインポート
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

device = "mps" if torch.mps.is_available() else "cpu"
device

def create_effnetb2_model(num_classes:int=3, 
                          seed:int=42):
    """EfficientNetB2特徴抽出モデルと変換処理を作成します。

    Args:
        num_classes (int, optional): 分類器ヘッドのクラス数。デフォルトは3。
        seed (int, optional): ランダムシード値。デフォルトは42。

    Returns:
        model (torch.nn.Module): EfficientNetB2特徴抽出モデル
        transforms (torchvision.transforms): EfficientNetB2用の画像変換処理
    """
    # 1, 2, 3. EfficientNetB2の事前学習済み重み、変換処理、モデルを作成
    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
    transforms = weights.transforms()
    model = torchvision.models.efficientnet_b2(weights=weights)

    # 4. ベースモデルの全レイヤーを凍結
    for param in model.parameters():
        param.requires_grad = False

    # 5. 再現性のためランダムシードを設定して分類器ヘッドを変更
    torch.manual_seed(seed)
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.3, inplace=True),
        nn.Linear(in_features=1408, out_features=num_classes),
    )

    return model, transforms

# Food101の101クラスに対応したEfficientNetB2モデルを作成
effnetb2_food101, effnetb2_transforms = create_effnetb2_model(num_classes=101)

# Food101学習データ用の変換処理（学習画像にのみデータ拡張を適用）
food101_train_transforms = torchvision.transforms.Compose([
    torchvision.transforms.TrivialAugmentWide(),
    effnetb2_transforms,
])

print(f"学習用変換処理:\n{food101_train_transforms}\n") 
print(f"テスト用変換処理:\n{effnetb2_transforms}")

from torchvision import datasets
from pathlib import Path

# データディレクトリの設定
data_dir = Path("data")

# 学習データの取得（~750画像 × 101食品クラス）
train_data = datasets.Food101(root=data_dir,
                              split="train",
                              transform=food101_train_transforms,
                              download=True)

# テストデータの取得（~250画像 × 101食品クラス）
test_data = datasets.Food101(root=data_dir,
                             split="test",
                             transform=effnetb2_transforms,
                             download=True)

# Food101のクラス名を取得
food101_class_names = train_data.classes

# 最初の10個を表示
food101_class_names[:10]

['apple_pie',
 'baby_back_ribs',
 'baklava',
 'beef_carpaccio',
 'beef_tartare',
 'beet_salad',
 'beignets',
 'bibimbap',
 'bread_pudding',
 'breakfast_burrito']

def split_dataset(dataset:torchvision.datasets, split_size:float=0.2, seed:int=42):
    """データセットを指定された比率でランダムに分割します。

    Args:
        dataset (torchvision.datasets): PyTorchデータセット
        split_size (float, optional): 分割サイズ。デフォルトは0.2
        seed (int, optional): ランダムシード。デフォルトは42

    Returns:
        tuple: (分割1, 分割2)のタプル
    """
    # 元のデータセット長に基づいて分割長を作成
    length_1 = int(len(dataset) * split_size) # 希望する長さ
    length_2 = len(dataset) - length_1 # 残りの長さ

    # 情報を表示
    print(f"[INFO] 長さ{len(dataset)}のデータセットを以下のサイズに分割: {length_1} ({int(split_size*100)}%), {length_2} ({int((1-split_size)*100)}%)")

    # 指定されたランダムシードで分割を作成
    random_split_1, random_split_2 = torch.utils.data.random_split(dataset, 
                                                                   lengths=[length_1, length_2],
                                                                   generator=torch.manual_seed(seed))
    return random_split_1, random_split_2

# Food101の学習データの20%分割を作成
train_data_food101_20_percent, _ = split_dataset(dataset=train_data,
                                                 split_size=0.2)

# Food101のテストデータの20%分割を作成
test_data_food101_20_percent, _ = split_dataset(dataset=test_data,
                                                split_size=0.2)

len(train_data_food101_20_percent), len(test_data_food101_20_percent)

import os

BATCH_SIZE = 32
NUM_WORKERS = 2 if os.cpu_count() <= 4 else 4

# Food101 20%学習用DataLoaderを作成
train_dataloader_food101_20_percent = torch.utils.data.DataLoader(train_data_food101_20_percent,
                                                                  batch_size=BATCH_SIZE,
                                                                  shuffle=True,
                                                                  num_workers=NUM_WORKERS)
# Food101 20%テスト用DataLoaderを作成
test_dataloader_food101_20_percent = torch.utils.data.DataLoader(test_data_food101_20_percent,
                                                                 batch_size=BATCH_SIZE,
                                                                 shuffle=False,
                                                                 num_workers=NUM_WORKERS)

from going_modular import engine

# オプティマイザの設定
optimizer = torch.optim.Adam(params=effnetb2_food101.parameters(),
                             lr=1e-3)

# 損失関数の設定（クラス数が多いためラベルスムージングを使用）
loss_fn = torch.nn.CrossEntropyLoss(label_smoothing=0.1)

# 元のFood101論文の56.4%の精度を20%のデータで超えることが目標
set_seeds()    
effnetb2_food101_results = engine.train(model=effnetb2_food101.to(device),
                                        train_dataloader=train_dataloader_food101_20_percent,
                                        test_dataloader=test_dataloader_food101_20_percent,
                                        optimizer=optimizer,
                                        loss_fn=loss_fn,
                                        epochs=5,
                                        device=device,
                                        writer=None)

from helper_functions import plot_loss_curves

# FoodVision Bigの損失曲線を確認
plot_loss_curves(effnetb2_food101_results)

from going_modular import utils

# モデルパスの作成
effnetb2_food101_model_path = "10_pretrained_effnetb2_feature_extractor_food101_20_percent.pth" 

# FoodVision Bigモデルを保存
utils.save_model(model=effnetb2_food101,
                 target_dir="models",
                 model_name=effnetb2_food101_model_path)

# Food101対応のEfficientNetB2インスタンスを作成
loaded_effnetb2_food101, effnetb2_transforms = create_effnetb2_model(num_classes=101)

# 保存されたモデルのstate_dict()を読み込み
loaded_effnetb2_food101.load_state_dict(torch.load("models/10_pretrained_effnetb2_feature_extractor_food101_20_percent.pth"))

from pathlib import Path

# モデルサイズをバイトで取得し、メガバイトに変換
pretrained_effnetb2_food101_model_size = Path("models", effnetb2_food101_model_path).stat().st_size // (1024*1024)
print(f"事前学習済みEfficientNetB2特徴抽出器Food101モデルサイズ: {pretrained_effnetb2_food101_model_size} MB")

from pathlib import Path

# FoodVision Bigデモパスの作成
foodvision_big_demo_path = Path("demos/foodvision_big/")

# FoodVision Bigデモディレクトリの作成
foodvision_big_demo_path.mkdir(parents=True, exist_ok=True)

# FoodVision Bigデモ例ディレクトリの作成
(foodvision_big_demo_path / "examples").mkdir(parents=True, exist_ok=True)

# サンプル画像をコピー
!cp data/sushi.jpg demos/foodvision_big/examples/sushi.jpg

# 学習済みモデルをFoodVision Bigデモフォルダに移動
!mv models/10_pretrained_effnetb2_feature_extractor_food101_20_percent.pth demos/foodvision_big

# Food101クラス名のパスを作成
foodvision_big_class_names_path = foodvision_big_demo_path / "class_names.txt"

# Food101クラス名リストをファイルに書き込み
with open(foodvision_big_class_names_path, "w") as f:
    print(f"[INFO] Food101クラス名を{foodvision_big_class_names_path}に保存")
    f.write("\n".join(food101_class_names)) # 各クラス間に改行を挿入

# Food101クラス名ファイルを開いて各行をリストに読み込み
with open(foodvision_big_class_names_path, "r") as f:
    food101_class_names_loaded = [food.strip() for food in  f.readlines()]

# 読み込んだ最初の5つのクラス名を表示
food101_class_names_loaded[:5]

['apple_pie', 'baby_back_ribs', 'baklava', 'beef_carpaccio', 'beef_tartare']

%%writefile demos/foodvision_big/model.py
import torch
import torchvision

from torch import nn


def create_effnetb2_model(num_classes:int=3, 
                          seed:int=42):
    """EfficientNetB2特徴抽出モデルと変換処理を作成します。

    Args:
        num_classes (int, optional): 分類器ヘッドのクラス数。
            デフォルトは3。
        seed (int, optional): ランダムシード値。デフォルトは42。

    Returns:
        model (torch.nn.Module): EfficientNetB2特徴抽出モデル
        transforms (torchvision.transforms): EfficientNetB2画像変換処理
    """
    # EfficientNetB2の事前学習済み重み、変換処理、モデルを作成
    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
    transforms = weights.transforms()
    model = torchvision.models.efficientnet_b2(weights=weights)

    # ベースモデルの全レイヤーを凍結
    for param in model.parameters():
        param.requires_grad = False

    # 再現性のためランダムシードを設定して分類器ヘッドを変更
    torch.manual_seed(seed)
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.3, inplace=True),
        nn.Linear(in_features=1408, out_features=num_classes),
    )

    return model, transforms

%%writefile demos/foodvision_big/app.py
### 1. インポートとクラス名の設定 ### 
import gradio as gr
import os
import torch

from model import create_effnetb2_model
from timeit import default_timer as timer
from typing import Tuple, Dict

# クラス名を設定
with open("class_names.txt", "r") as f: # class_names.txtから読み込み
    class_names = [food_name.strip() for food_name in  f.readlines()]

### 2. モデルと変換処理の準備 ###    

# モデルを作成
effnetb2, effnetb2_transforms = create_effnetb2_model(
    num_classes=101, # len(class_names)も使用可能
)

# 保存された重みを読み込み
effnetb2.load_state_dict(
    torch.load(
        f="10_pretrained_effnetb2_feature_extractor_food101_20_percent.pth",
        map_location=torch.device("cpu"),  # CPUに読み込み
    )
)

### 3. 予測関数 ###

# 予測関数を作成
def predict(img) -> Tuple[Dict, float]:
    """画像を変換して予測を実行し、予測結果と所要時間を返します。
    """
    # タイマーを開始
    start_time = timer()

    # 対象画像を変換してバッチ次元を追加
    img = effnetb2_transforms(img).unsqueeze(0)

    # モデルを評価モードにして推論モードをオンにする
    effnetb2.eval()
    with torch.inference_mode():
        # 変換された画像をモデルに通し、予測ロジットを予測確率に変換
        pred_probs = torch.softmax(effnetb2(img), dim=1)

    # 各予測クラスの予測ラベルと予測確率の辞書を作成（Gradioの出力パラメータに必要な形式）
    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}

    # 予測時間を計算
    pred_time = round(timer() - start_time, 5)

    # 予測辞書と予測時間を返す 
    return pred_labels_and_probs, pred_time

### 4. Gradioアプリ ###

# タイトル、説明、記事の文字列を作成
title = "FoodVision Big"
description = "101種類の食品画像を分類するEfficientNetB2特徴抽出コンピュータービジョンモデル。"
article = "[10. PyTorchモデルデプロイメント](https://vinsmoke-three.com/PyTorch/10_pytorch_model_deployment/)で作成。"

# "examples/"ディレクトリから例のリストを作成
example_list = [["examples/" + example] for example in os.listdir("examples")]

# Gradioインターフェースを作成 
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(num_top_classes=5, label="予測結果"),
        gr.Number(label="予測時間 (秒)"),
    ],
    examples=example_list,
    title=title,
    description=description,
    article=article,
)

# アプリを起動！
demo.launch()

%%writefile demos/foodvision_big/requirements.txt
torch>=2.3.1,<2.8.0
torchvision>=0.18.1,<0.23.0
gradio>=5.35.0,<6.0.0