Jupyter Notebookファイルのバッチ化を考える

図は機械学習システムの運用の泥臭さに打ちのめされるエンジニアのイラストです。

はじめに

記事概要

機械学習(Machine Learning; 以下ML)プロジェクトにおいて、PoCフェーズから製品開発フェーズに移るステップの1つとしてJupyter Notebookのファイル(以下ipynbファイル)をpython scriptのバッチファイルにする作業があります。
この記事ではバッチ化の意義を製品の運用効率化の観点から論じた後、MNIST分類器をケーススタディとしてipynbファイルをバッチ化をしてみます。

想定読者

• Jupyter Notebookでいいモデルが学習できたので、このモデルを組み込んだサービスを公開する予定の方
• 今はまだJupyter Notebookで実験しているけど、いずれは自分のMLモデルを載せたサービスをリリースしようと考えている方
• 誰かに「このipynbファイル、バッチ化しといて」と頼まれた方¹

本記事におけるバッチ化の定義

この記事では、ipynbファイルのコードを学習や評価をする上で使う機能ごとに分割して、コマンドライン実行で学習できる状態にしたpythonスクリプトを作ることとします。
間違っても、train(...)の中にipynbファイルのセルの中身を詰め込んだtrain.pyを作ることではありません。

バッチ化の意義編

手頃に使えるJupyter Notebookですが、なぜバッチ化をするかをJupyter Notebookの得失とモデル再学習の観点から見ていきます。

Jupyter Notebookの得失

Jupyter Notebookの利点

Jupyter Notebookやjupyter-labのよい点としては、

• それなりにコーディングしやすい開発環境がすぐに整えられる
  • コードのハイライト、コード補完があるので、一応開発できなくはない。
  • ワーキングディレクトリが開いているノートブックなので、resource not foundやimport errorが起きにくい
• 小回りを利かせてコードを書けるので試行錯誤がしやすい
  • コードの即時実行により、EDAやデバッグがしやすい
    • データがどんな傾向を持っているのかの調査
    • どんな変数が入っているか、処理によってどんな画像ができたのかの簡易的な表示できる
  • それ故、実行結果をすぐにコードにフィードバックしてモデルの性能改善がしやすい
• 実験を再現しやすい
  • セルを順次実行すれば誰でも同じ評価結果になる

したがって、個人で遂行する実験用途やPoCに適していると言えます。

Jupyter Notebookの欠点

一方で、欠点としては、次のようなものが挙げられます。上2つはJupyter Notebookそのものに起因する欠点では無い気がしますが、自戒も込めて記しました。

• コードが暫定対応を繰り返されていて、可読性・保守性が悪い
  • 柔軟にコードを変えて実験をするがゆえに、当事者しか読み方がわからないようなコードになっている
    • セルをまたいだ変数が関数内にグローバルに参照されている
    • セルの実行順が上から下でない
    • 簡易的な変数名(例: a1,bb=a*calc(x, net) # コメントはセル内をスクロールしないと読めない文字数の説明が〜)
• スケーラビリティが悪い
  • 特徴抽出のアルゴリズムやMLモデルの定義やモデル学習/評価ロジックが一緒くたに入っている。
    • 巨大なブラックボックス状態
    • 共通コードがpython scriptファイルに別途まとめられていることもあるが、実際にはロジックの単なる抜き出しで、他の関数や他のプロジェクトで再利用できないことも。。。
  • したがって、別のモデルで試したいとか、別のデータも追加したいとか、前処理を追加したいといった要件の変更に対するスケーラビリティが少ない
• 学習の自動実行ができないため、モデル再学習などの運用業務と相性が悪い(詳しくは後述の「なぜ再学習は必要か」にて)
  • 新規データの追加、パラメータ探索となった時、手作業でipynbのコードを書き直す羽目になる
    • ただでさえ暫定対応のコードにさらにforループなどが付け加わり、運用を行えば行うほど大変なことになっていきます。
  • バッチサーバ、パイプラインツール、MLプラットフォームへの登録ができない
    • データバージョニングなどのトレーサビリティ確保や、計算リソースの割当てなどの恩恵に与ることができない。

したがって、ipynbでの学習は、PoC以降の製品開発や保守/改修/運用フェーズで問題になります²。

なぜ再学習は必要か

ここまで読んでみて、「なぜ再学習という作業をするのか」という疑問をお持ちの方もいらっしゃるかもしれません。精度が出たらもう学習は不要なのでは、ということですね。答えとしては、データの質の変化に伴うモデルの性能劣化に対応するためです。

MLは使用データに依存した技術です。根本となるデータの性質がモデルの学習時と変わるとモデルの予測性能が衰えます。例えば、コロナ禍のせいで消費者マインドが変わったことでレコメンドがうまくいかなかったり、法改正によって有権者の年齢が下がることで若年層の投票傾向が変わるとかですね。
恐ろしいことに、たとえコードを一切変えていなくても、モデルの性能が変わる(大抵は悪くなる)のです。データの性質は自然に、あるいは人為的にどちらの要因でも変わっていくので、これを考慮しないでサービスをリリースすると、例えML以外のロジックにバグが一切無いとしても、やがてサービス品質は劣化するでしょう。

本当は、機械学習モデルそのものに対する品質保証手法、すなわちMLモデルに問題があるかを事前に検査できる方法があればよいのですが、現時点では研究段階という印象です。したがって、データの質的変化に伴うモデルの劣化を、モデルを再学習して事後対応する形でサービスを保守するのが現実的です。

繰り返しのモデル学習において、人手で動かすJupyter Notebookを使うのは運用的に辛いものがあるので、バッチ化をして学習を半自動化できるようにしましょうというわけです。

バッチ化実践編

それでは実際にipynbのバッチ化を行っていきます。最終的にできあがるものはこちらのリポジトリに配置しています。

題材

今回のバッチ化の題材として使うipynbファイルは「PyTorchニューラルネットワーク実装ハンドブック」の4.1節 AlexNetで画像分類するタスクです。

フレームワークはPyTorchで話を進めていますが、実装の流れはどのフレームワークを使っても同じ話になると思います。また、想定するタスクは画像分類ですが、その他のタスクでも考え方は通用すると思います。
もとのipynbファイルで使用されているデータはCIFAR-10ですが、本記事の続編として手書き文字認識Webアプリを作る記事を執筆予定なので、MNISTにします。

バッチの呼び出し方

コマンドラインやバッチサーバで実行するとき、train(config_yml, working_root)を呼ぶ想定です。
この関数の引数は2つで、コンフィグファイルとワーキングルート(どこを起点としてデータを参照するか)です。

バッチ化の方針

バッチ化の方針としては、ふぁむたろうさんが執筆されたPyTorch 三国志（Ignite・Catalyst・Lightning）の「2.2.1 共通部分」、「2.2.1 ベースコード(素の学習用コード)」、「3. Catalyst・Ignite・Lightning を自由に行き来するために」をパクリ非常に参考にしています。こちらの記事はPyTorchのモデル学習ラッパー3種類を比較する旨のものですが、本題とは裏腹に、機械学習コードの抽象化論も提案されています。

下記はその項目の抜粋です(列挙順は若干変えています)。

・ Config を一つにまとめる
・ optimizer や model を呼び出す関数を作る
・ ループの中身はなるべく取り出しておく
・ 関数の引数を増やしすぎない

本記事もこの方針に倣って、バッチ化をしていきます。

具体的には

• コンフィグを辞書に統一する
• 学習に必要なオブジェクトをコンフィグを元に生成する関数を実装する
• 学習フェーズの最適化ループの中身は関数として抜き出す
• 関数の引数はconfigのdictを渡す

を行っていきます。

バッチ化

先に私がバッチ化した学習ロジック部分を掲載します。少々長いので折りたたみました。軽く眺めるだけで大丈夫です。
ipynbファイルのコード量とさほど変わらないと思われるかもしれませんが、スケーラビリティの点で違います。

↓の☆の隣にある三角をクリックでコードを展開します

☆クリックでコードを展開します

import sys
import os

import yaml
import torch
import torchvision
import torch.nn as nn
from tqdm import tqdm

import sys
from pathlib import Path
this_file_dir = Path(__file__).parent.resolve()

sys.path += [
    str(this_file_dir),
]

from models.model_factory import get_model, load_weight
from datasets.dataset_factory import get_dataloader_through_dataset
from losses.loss_factory import get_loss
from optimizers.optimizer_factory import get_optimizer
from output_data.data_saver import DataSaver


def train(config_yml, 
          working_root=str(this_file_dir / '..')):
    """画像分類モデルを学習するエンドポイント
    Args:
        config_yml (str): コンフィグ用のyamlファイル
        working_root (str, optional): どこを起点としてデータを参照するか. Defaults to str(this_file_dir / '..').
    """

    with open(config_yml, 'r') as f:
        config = yaml.safe_load(f)

    # ====
    # データ用意
    # ====
    train_loader = get_dataloader_through_dataset(
        config['data']['train'],
        working_root
    )
    test_loader  = get_dataloader_through_dataset(
        config['data']['eval'],
        working_root,
    )

    # ネットワーク用意
    net = get_model(config['model'])
    if config['model'].get('model_state_dict'):
        model_state_dict_path = Path(working_root) / config['model']['model_state_dict']
        load_weight(net, str(model_state_dict_path))
    
    # optimizer定義
    optimizer = get_optimizer(net, config['optimizer'])

    # 損失関数定義
    criterion = get_loss(config['loss'])

    # データを保存する機能を持つオブジェクト
    datasaver = DataSaver(config['output_data'])
    datasaver.save_config(config_yml)

    # ======
    # メインループ
    # ======
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if 'cuda' in config:
        device = 'cuda' if config['cuda'] == True else 'cpu'

    net.to(device)

    num_epochs = config['num_epochs']
    for epoch in range(num_epochs):
        print(epoch)
        metrics_dict = {}

        #train
        print('train phase')
        metrics = run_train(net, train_loader, criterion, optimizer, device)
        metrics_dict.update(metrics)

        # eval
        print('eval phase')
        metrics, result_detail = run_eval(net, test_loader, criterion, device)
        metrics_dict.update(metrics)

        # 評価指標の記録
        datasaver.save_metrics(metrics_dict, epoch)
        datasaver.save_model(net, epoch)
        datasaver.save_result_detail(result_detail, epoch)


def run_train(model, data_loader, criterion, optimizer, device, grad_acc=1):
    model.train()

    # zero the parameter gradients
    optimizer.zero_grad()

    total_loss = 0.
    train_acc = 0
    for i, (inputs, gt_labels, _, _) in tqdm(enumerate(data_loader), total=len(data_loader)):
        
        # 画像が合っているかを確認
        if os.environ.get('DEBUG', None) is not None:
            import cv2
            im = (inputs.numpy()*255)[0].transpose(1,2,0)
            cv2.imwrite('./sample.png', im)
        
        inputs = inputs.to(device)
        gt_labels = gt_labels.to(device)

        outputs = model(inputs)

        loss = criterion(outputs, gt_labels)
        loss.backward()

        # Gradient accumulation
        if (i % grad_acc) == 0:
            optimizer.step()
            optimizer.zero_grad()

        total_loss += loss.item()
        train_acc += (outputs.max(1)[1] == gt_labels).sum().item()

    total_loss /= len(data_loader)
    avg_train_acc = train_acc / len(data_loader.dataset)
    
    metrics = {'train_loss': total_loss,
               'train_acc': avg_train_acc}

    return metrics


def run_eval(model, data_loader, criterion, device):
    model.eval()
    val_acc = 0.
    
    result_detail = []

    with torch.no_grad():
        total_loss = 0.
        for inputs, gt_labels, filenames, dataset_name in tqdm(data_loader, total=len(data_loader)):
            inputs = inputs.to(device)
            gt_labels = gt_labels.to(device)

            outputs = model(inputs)
            est = outputs.max(1)[1]

            loss = criterion(outputs, gt_labels)

            total_loss += loss.item()
            correct = (est == gt_labels)
            val_acc += correct.sum().item()
    
            l = []
            for i in range(len(filenames)):
                gt_label = model.class_labels[gt_labels[i].data.item()]
                est_label = model.class_labels[est.data[i].item()]
                result_detail.append([dataset_name[i],
                                      filenames[i], 
                                      gt_label, 
                                      est_label])
            
        total_loss /= len(data_loader)  # iterした回数で割る
        val_acc /= len(data_loader.dataset)  # 画像の総枚数で割る
        metrics = {'val_loss': total_loss, 
                   'val_acc': val_acc}

    return metrics, result_detail


if __name__ == '__main__':
    argvs = sys.argv
    train(argvs[1], argvs[2])

それでは、先述の方針に沿って、どこをどう変えたのかを説明します。

コンフィグを辞書に統一する

MLにおけるコンフィグ(configuration)は学習をどう行うかを決める情報です。例えば、データはどういうデータセットを使うか、モデルはどういうハイパーパラメータかなど、一連のモデル学習を構成しているパラメータの情報です。

元のipynbファイルでは、データセットもMNIST、モデルもAlexNet、optimizerもAdam、といった風に決め打ちですので、下記のようにコード中に設定が書かれています。

device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = AlexNet(num_classes).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)

num_epochs = 20

しかし、要素を変えて学習したい(例えばoptimizerにSGDを使いたい)という場合は、コンフィグを一つにまとめて管理した方が、学習条件の視認性がよくなります。すなわち、この精度はどういう条件で学習した結果出たものなのかを、わざわざコードを見なくてもコンフィグファイルを見れば把握できます。

したがって、これを以下のような形でyamlファイルでまとめておきます。

# 学習エポック数
num_epochs: 20
# GPUを使用するか
cuda: true
# 成果物を出力するディレクトリ。working rootからの相対パス。nullの場合は .myoutput に作られる
output_data:

data:
  # 学習用
  train:
    # データセット
    dataset:
      # データセット
      # MNISTのdatasetクラスを取得するためのパラメータ
      mnist_train:
        # データの置き場
        dir: &dataset data/images
        train_mode: true
        # 画像の前処理。上から下へ順番にかける。
        transform: &transform
          - Resize:
            width: 32
            height: 32
          - Gray2BGR:
          - Invert:
          - ToTensor:
    dataloader: &dataloader
      batch_size: 64
      shuffle: true
  # 評価
  eval:
    dataset:
      mnist_eval:
        dir: *dataset
        train_mode: false
        transform: *transform
    dataloader: 
      batch_size: 64
      shuffle: false

# ネットワークアーキテクチャ
model:
  model_name: alexnet
  # クラス数
  num_classes: 10
  # クラスラベル
  class_labels: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

loss:
  loss_name: cross_entropy_loss

optimizer:
  optimizer_name: sgd
  lr: 0.01
  momentum: 0.9
  weight_decay: 0.0005

もちろん、このコンフィグをもとにオブジェクトを生成する関数を別途用意する必要がありますが、それは後述します。
このように役割ごと(データセット、モデル、optimizerなど)にパラメータをまとめておくことで、オブジェクト生成関数に情報を渡しやすくなります。

どのようなパラメータをコンフィグに含めるのか、という議論はあると思います。それに対する私のスタンスはコンフィグを構成するパラメータは後から変えればいいという立場をとります。やや楽観的かもしれません。

学習に必要なオブジェクトをコンフィグを元に生成する関数を実装する

別の前処理やモデルやoptimizerを試したい時、ipynbファイルだけで完結させようとすると、下記のようにコードに書く形でオブジェクトの生成方法を手動で切り替えるかもしれません。

#optimizer = torch.optim.SGD(...)
#optimizer = torch.optim.Adagrad(...)
optimizer = torch.optim.Adam(...)

これはコード中にコンフィグを書いている状態といってよいでしょう。このままだと設定の切り替えミス、戻し忘れミスなどが発生するかもしれないので、コンフィグ辞書を元に自動的に生成したいです。
そのための関数を用意します。ちなみにオブジェクトを生成する関数は、factoryと名のつくファイル内に、含められることが多いです。

下記は、optimizerのオブジェクトを生成する関数の例です³。

import torch.optim as optim

def get_optimizer(net, config):
    
    optimizer_name = config['optimizer_name']
    
    if optimizer_name == 'sgd':
        return optim.SGD(
            net.parameters(), 
            lr=config['lr'], 
            momentum=config['momentum'], 
            weight_decay=config['weight_decay'])
    elif optimizer_name == 'adam':
        # ...
    else:
        raise ValueError(f'invalid optimizer name {optimizer_name}')

このように、オブジェクト生成関数を実装することで、コードとコンフィグを切り離すことが可能になります。

この生成関数によって、後からでも柔軟にクラスを追加することができるようになります。たとえば、オフィシャルには採用されていない有志が実装したoptimizerのクラスを試したいときは、get_optimizer()にそのクラスのオブジェクトを生成する旨を追記します。あとはコンフィグの'optimizer_name'に対応する情報を変えれば、新しいoptimizerで学習を進めることができるようになります。
このように前処理、モデル、optimizer、といった単位でオブジェクト生成関数をつくっておくことで、学習ロジックが抽象化できます。
このオブジェクト生成関数の存在は地味ですが、学習のスケーラビリティを上げてくれることに非常に貢献してくれるパーツです。

学習フェーズの最適化ループの中身は関数として抜き出す

学習のforループが長くなって処理の流れがどうなっているのか追いきれなくなることがあるので、関数として抜き出しておきます。
また、関数化しておくと、関数単体で再利用できる可能性が高まります。

オリジナルのipynbファイルでは

train_loss_list = []
train_acc_list = []
val_loss_list = []
val_acc_list = []

for epoch in range(num_epochs):
    train_loss = 0
    train_acc = 0
    val_loss = 0
    val_acc = 0
    
    #train
    net.train()
    for i, (images, labels) in enumerate(train_loader):
      #view()での変換をしない
      images, labels = images.to(device), labels.to(device)
      
      optimizer.zero_grad()
      outputs = net(images)
      loss = criterion(outputs, labels)
      train_loss += loss.item()
      train_acc += (outputs.max(1)[1] == labels).sum().item()
      loss.backward()
      optimizer.step()
    
    avg_train_loss = train_loss / len(train_loader.dataset)
    avg_train_acc = train_acc / len(train_loader.dataset)
    
    #val
    net.eval()
    with torch.no_grad():
      for images, labels in test_loader:
        #view()での変換をしない
        images = images.to(device)
        labels = labels.to(device)
        outputs = net(images)
        loss = criterion(outputs, labels)
        val_loss += loss.item()
        val_acc += (outputs.max(1)[1] == labels).sum().item()
    avg_val_loss = val_loss / len(test_loader.dataset)
    avg_val_acc = val_acc / len(test_loader.dataset)
    
    print ('Epoch [{}/{}], Loss: {loss:.4f}, val_loss: {val_loss:.4f}, val_acc: {val_acc:.4f}' 
                   .format(epoch+1, num_epochs, i+1, loss=avg_train_loss, val_loss=avg_val_loss, val_acc=avg_val_acc))
    train_loss_list.append(avg_train_loss)
    train_acc_list.append(avg_train_acc)
    val_loss_list.append(avg_val_loss)
    val_acc_list.append(avg_val_acc)

この部分をモデル学習関数run_train()とモデル評価関数run_eval()で外部に出しておくと、可読性が上がりforループの中を追いやすくなります。
↓

for epoch in range(num_epochs):
    metrics_dict = {}

    #train
    metrics = run_train(net, train_loader, criterion, optimizer, device)
    metrics_dict.update(metrics)

    #val
    metrics = run_eval(net, test_loader, criterion, device)
    metrics_dict.update(metrics)

関数の引数はconfigのdictを渡す

関数にまとめるのはいいのですが、関数がパラメータを引数に取りすぎると、ごちゃごちゃするので、パラメータをまとめたもので渡すとすっきりします。
今一度、optimizerのオブジェクトを生成する関数を再掲します。

def get_optimizer(net, config):
    
    optimizer_name = config['optimizer_name']
    
    if optimizer_name == 'sgd':
        return optim.SGD(
            net.parameters(), 
            lr=config['lr'], 
            momentum=config['momentum'], 
            weight_decay=config['weight_decay'])
    elif optimizer_name == 'adam':
        # ...

必要なパラメータを辞書から適宜ひっぱってきています。
関数インタフェースがシンプルになるとはいえ、第三者がその関数を見た時に、必須なコンフィグの要素を把握できるようにしておく必要はあります。

[備考] その他バッチ化に伴って追加した機能

• トレーサビリティの確保
  • 実験管理のような機能もつけています。DataSaverというクラスを実装しました。
  • 実行すると、デフォルトで.myoutput/yyyy-mm-dd_hhMMssに、学習したコード、コンフィグ、アーティファクト(重み)、メトリクス、評価の詳細ファイルを吐くようにしています。

バッチ化における、今後の課題

• 関数に必要なdictの要素をどう情報提示するか
  • 関数のインタフェースをを辞書にまとめたが故に、パラメータの与え方が自明ではなくなってしまいます。そのため、何が必須パラメータなのか分からないということに陥ります。
  • readmeやdocstringに書くこともできますが、関数の機能変更の時に都度書き直す手間が発生します。その手間の削減をどうすればいいか分かっていません。
• コンフィグ保持を辞書からクラスに変える
  • 今は辞書型でコンフィグを持っていますが、書籍「自走プログラマー」のプラクティスNo.12によると、データの保持には辞書ではなく(dataclassの)クラスで持つべきだと書いてあります。クラスだとデータの検証ができるのが強みとのことです。
• アーティファクトの保存や、メトリクスのロギングに外部パッケージを使用
  • 今回はDataSaverというクラスを自前で実装しました。このクラスのオブジェクトが、学習コードと、アーティファクト、メトリクスの保存をおこなってくれます。
    • しかし、MLflowを使ったり、tensorbordXなどのOSSのロガーを使ってメトリクスを管理したほうが、自分の開発範囲を抑えることが期待できます。
• 認識対象のスケーラビリティ
  • 今回は手書き数字読み取りであるとタスクを決めたので、クラス数、ラベルをコンフィグの中で与えています。
  • しかし、例えば追加要件で、手書きの英字も読めるようにしてほしいとなった時は、コンフィグの構成要素のうち、影響のある部分(モデル定義)を変える必要があります。
• コンフィグファイルはyamlでいいのか
  • yamlは変数が使えないからyamlが長くなり、コンフィグが読みにくくなることがあります。
  • 以前、pythonのコンフィグ辞書だけを書いたスクリプトをコンフィグに入力する方法も考えましたが、未だにそれに代わるアイディアが浮かんでいません。
• テストケースの用意
  • 構成要素に分けるので、テストケースはぜひとも用意したいです。しかし、学習は再現が難しいですし、テスト実行一回あたりの時間も長くなってしまいます。
  • 十分ではないですが案としては、学習データを1件だけ学習し、さらにそれを評価データにしてちゃんと読めるかを調べるのが妥協ラインかと思います。

おわりに

バッチ化が必要な背景を、ipynbの得失とモデルの再学習の観点から説明しました。
また、ipynbファイルのバッチ化の一例を示しました。コンフィグとコードを分離して、コンフィグに沿ったオブジェクトを生成し、学習ループを回すようにしました。
効率よくモデル学習する方法論はまだまだ模索中なので、今後も事例を収集していきたいです。

参考文献

• "継続的改善をし続けるための機械学習基盤の課題", http://ibisml.org/ibis2019/files/2019/12/slide_ariga.pdf
• "AI/MLシステム開発の難しさ", https://speakerdeck.com/azaazato/mlsisutemukai-fa-falsenan-sisa
• "機械学習プロジェクトを頑健にする施策", https://speakerdeck.com/takahiko03/ji-jie-xue-xi-puroziekutowowan-jian-nisurushi-ce-ml-ops-study-number-2
• "Avoiding technical debt with ML pipelines", https://blog.maiot.io/technical_debt/

1. 他人にバッチ化をやらせるのは、プロジェクトマネジメントの観点から見てアンチパターンなので分業体制を見直したほうがいいと思いますが... ↩

2. 逆に言えば製品開発に移行しないのであれば、ipynbファイルのままでもよいでしょう。 ↩

3. もっとも、pytorchに標準実装されているoptimizerだけを生成したければ、ふぁむたろうさんの実装のように、getattrでクラスを参照した方がスッキリすると思います。 ↩