Python で快適（主観）な機械学習実験

この記事は夏のブログリレー 4日目の記事です。

こんにちは。ikura-hamu です。traP では Go で Web アプリなどを作っています。一方で、研究室では Python で機械学習っぽいことをやっています。この記事では、Python で機械学習の実験のプログラムを書くにあたって、これをやったらうまくいいったな、というのを紹介します。機械学習の実験は人や分野によって書くプログラムの内容が変わってくるのと思うので、全員に言えることではないと思うのですが、誰かの参考になればいいなと思って書きます。

始めに、自分が行う実験や書くプログラムの特徴を説明しておきます。これと共通する部分が多い人は、この記事で参考になる部分も多いと思います。

PyTorch を使う。ResNet-18 とかの既存のモデルをそのまま使うことが多い。
データセットも、CIFAR-10 とかの整った形で世に出回ってるものを使うことが多い。
回帰タスクより分類タスクを行うことが多い。
1 つの実験にかかる時間は数時間から 1 日程度。
同じコードベースを数か月から 1 年の長い間扱う。
同時並行でちょっと設定を変えたたくさんの実験を回したい。
環境は研究室の共用 GPU サーバーか、スーパーコンピューター（TSUBAME 4）を使う。

コンテナを使う

コンテナを使うと、環境が分離出来ていい感じになります。研究室サーバーでは Podman という rootless コンテナが整備されているので、それを使っています。Docker など他のコンテナでもいいと思います。

共用サーバーだと、欲しいツールをインストールするのに管理者権限が必要だったりして面倒なことがありますが、コンテナに閉じ込めてしまえば好きに環境を作ることができます。

おすすめなのは Python のプログラムを bind mount するのではなく、そのままコンテナイメージに載せてしまうことです。bind mount するとプログラムを書き換えた結果がコンテナ内にそのまま反映されてしまい、例えばコンテナ内部で繰り返してプログラムを実行するようにしていた場合は意図せずプログラムが変更されてしまう場合があります。コンテナイメージにプログラムを固定してあげることで、そのような事故を防ぐことができます。

静的な型チェッカーを活用する

Python は動的型付き言語ですが、様々な静的な型を扱うツールが揃っています。これらを活用することで、長い実験の終盤にプログラムのミスで実験が失敗する、というような悲しい事象が起こる確率を下げることができます。

ty や mypy といった型をチェックするツールを使用することで、書いたプログラムで型がちゃんと扱われているかを調べることができます。ty は、uv を作っている Astral が作った Rust 製の型チェッカーで、とにかく高速なのが特長です。プログラムを書いてから型チェックの結果が表示されるまでが短いと、非常に快適にプログラムを書くことができるのでおすすめです。

pydantic + tyro で実験設定をバリデーションする

pydantic と tyro は、どちらも Python のライブラリです。

pydantic は、Web アプリ開発などでもよく使われるライブラリで、クラスを定義してデータをバリデーションすることができます。クラスのコンストラクタに渡されたデータが決まったデータと違うものの場合はバリデーションエラーになります。

a.pyfrom pydantic import BaseModel, PositiveInt

class Config(BaseModel):
  n_epochs: PositiveInt
  dataset: Literal["cifar10", "mnist"]

Config(n_epochs=3, dataset="cifar10") # ok
Config(n_epochs=-1, dataset="cifar100") # バリデーションエラー

tyro は、pydantic や dataclass から CLI のインターフェースを作ることができるライブラリです。pydantic などで定義したデータ構造をもとに、コマンドラインオプションのパースやヘルプの表示を行ってくれます。バリデーションに失敗したときのエラーメッセージも非常に丁寧です。

a.pyconfig = tyro.cli(Config)

$ uv run a.py
╭─ Required options ────────────╮
│ Missing from a.py:            │
│     --n-epochs INT            │
│         (required)            │
│     --dataset {cifar10,mnist} │
│         (required)            │
│ ───────────────────────────── │
│ For full helptext, run:       │
│     a.py --help               │
╰───────────────────────────────╯

これを組み合わせることで、設定を簡単に変えながら実験を行うことができ、設定のバリデーション処理をライブラリに移譲して実験本体のプログラムに注力することができます。ユーザーからの入力は静的型チェッカーが対応できない部分なのでバリデーションが重要になりますが、pydantic + tyro で容易に実現できます。

また、たくさんの実験を同時に行う際は、コマンドライン引数ではなく JSON などのファイルに設定を書き出して実行したくなります。その場合は、ファイルを読んで辞書を pydantic のクラスに渡してあげるとバリデーションを行ってくれます。

with open("config.json", "r", encoding="utf-8") as f:
    config_dict = json.load(f)
config = Config(**config_dict)

`typing.NewType` を使ってプリミティブ型に別名を付ける

typing.NewType を使うと、型チェッカーが元の型と別の型として認識してくれます。これを使って int 型に名前を付けて、普通の数とは別の意味を持つ数として扱うことができます。自分のプログラムでは、データセットのラベルを表す LabelID などを定義して使っています。

from typing import NewType
LabelID = NewType("LabelID", int)

これを使うと、変数などの型の意味が分かりやすくなります。例えば、データセットのラベルごとのサンプル数を表す型として、通常の int を使うと、インデックスを用いてlist[int]としたり、dict[int, int]としたりしますが、ラベルの番号とは違う意味の int 型をキーに使ってしまう可能性があります。これを独自のLabelIDとしてdict[LabelID, int]のように表すことで、何を表す型なのかが明確になり、ミスが減ります。

Git でコミットしてから実験する

Git でコミットしてから実験すると、あとからそこに戻って実験をやり直したいときに簡単に戻ることができます。コミットハッシュは実験結果や設定と一緒に記録するようにしています。自分はタスクランナーに Task を使っており、コンテナイメージをビルドするタスクに Git の状態を確認するコマンドを含め、未コミットの変更がある場合は失敗するようにしています。

notebook は marimo を使う

実験のプログラムそのものは通常の Python のプログラムを書いて使っていますが、データを分析して可視化するときは画像などが表示できる notebook が欲しくなります。notebook というと Jupyter Notebook が有名ですが、自分は marimo を使っています。marimo は、Jupyter Notebook に比べていくつか優れた点があると考えています。

Jupyter Notebook の実体は JSON ファイルであり、直接読むのは難しく、コーディングエージェントなどに書かせるのも大変です。一方、marimo は実体が Python のスクリプトであるため、読みやすく、コーディングエージェントにも優しい仕様になっています。
また、Jupyter Notebook は実行結果が notebook に保存されるのに対し、marimo は実行結果は保存されません。これは一見良くないように見えますが、Git 管理を考えると優れた仕組みであることが分かります。notebook を実行するたびに差分が生じてしまい、それを Git 管理に含めると、何を意図した変更なのかが分かりにくくなってしまいます。

marimo は他にも、セルが変更された際にそれに依存したセルが自動で実行されたり、UI を簡単に表示できたりと、便利な機能がたくさんあり、とてもおすすめです。

おわり

なんとかして Python を快適に書けるように頑張っています。継続していろいろ試していきたいです。

明日の担当は @fken_57 さんです。