[堅牢.py #1] テストを書かない研究者に送る、最初にテストを書く実験コード入門 / Let's start your ML project by writing tests

pytest と「最初にテストを書く」という考え方を活用することで、機械学習の実験コードに再現性と設計の明示性を確保する手法を、主にテストを書くのが億劫だなと感じている研究者・機械学習エンジニアに向けて紹介します

■ 堅牢.py #1
https://kenro.connpass.com/event/371009/

■ 登壇概要
タイトル：[堅牢.py #1] テストを書かない研究者に送る、最初にテストを書く実験コード入門

■ サンプルコード
🐙：https://github.com/shunk031/pytest-ml-tdd-example

Transcript

1. © LY Corporation テストを書かない研究者に送る “最初にテストを書く” 実験コード入門 – オレオレ最強 main.py から抜け出すために – サンプルレポジトリ: github.com/shunk031/pytest-ml-tdd-example AI Corporate Business Unit / Visual Generation Div. Shunsuke Kitada, Ph.D. HP: shunk031.me / 𝕏: @shunk031 堅牢.py #1 招待講演 @ 株式会社ユーザベース #kenro_py
2. © LY Corporation 経歴 • ‘23/04 LINE ➜ ‘23/10 LINEヤフー Research Scientist • ‘23/03 法政大学大学院 彌冨研 博士 (工学) / 学振 DC2 研究分野 • 自然言語処理 (NLP) / 画像処理 (CV) ◦ 摂動に頑健で解釈可能な深層学習 [Kitada+ IEEE Access’21, Appl. Intell.’22] • 計算機広告 (Multi-modal / Vision & Language) ◦ 効果の高いデジタル広告の作成支援 [Kitada+ KDD’19] ◦ 効果の低いデジタル広告の停止支援 [Kitada+ Appl. Sci.’22] • デザイン生成 AI ◦ 離散拡散モデルで生成されたレイアウトの再調整 [Iwai+ ECCV’24] ◦ LLMによるレイアウトの生成に対する自己修正 [Zhang+ arXiv’24] 自己紹介: 北田俊輔 Shunsuke KITADA 2
 🏠: shunk031.me / 𝕏: @shunk031 画像生成AIにおける 拡散モデルの理論と実践 リサーチサイエンティスト 北田俊輔 www.youtube.com/watch?v =-IPEUOcPTas Pythonで学ぶ画像生成 北田俊輔 インプレス社 https://book.impress.co.j p/books/1123101104
3. © LY Corporation 経歴 • ‘23/04 LINE ➜ ‘23/10 LINEヤフー Research Scientist • ‘23/03 法政大学大学院 彌冨研 博士 (工学) / 学振 DC2 研究分野 • 自然言語処理 (NLP) / 画像処理 (CV) ◦ 摂動に頑健で解釈可能な深層学習 [Kitada+ IEEE Access’21, Appl. Intell.’22] • 計算機広告 (Multi-modal / Vision & Language) ◦ 効果の高いデジタル広告の作成支援 [Kitada+ KDD’19] ◦ 効果の低いデジタル広告の停止支援 [Kitada+ Appl. Sci.’22] • デザイン生成 AI ◦ 離散拡散モデルで生成されたレイアウトの再調整 [Iwai+ ECCV’24] ◦ LLMによるレイアウトの生成に対する自己修正 [Zhang+ arXiv’24] 自己紹介: 北田俊輔 Shunsuke KITADA 3
 🏠: shunk031.me / 𝕏: @shunk031 画像生成AIにおける 拡散モデルの理論と実践 リサーチサイエンティスト 北田俊輔 www.youtube.com/watch?v =-IPEUOcPTas Pythonで学ぶ画像生成 北田俊輔 インプレス社 https://book.impress.co.j p/books/1123101104
4. © LY Corporation 経歴 • ‘23/04 LINE ➜ ‘23/10 LINEヤフー Research Scientist • ‘23/03 法政大学大学院 彌冨研 博士 (工学) / 学振 DC2 研究分野 • 自然言語処理 (NLP) / 画像処理 (CV) ◦ 摂動に頑健で解釈可能な深層学習 [Kitada+ IEEE Access’21, Appl. Intell.’22] • 計算機広告 (Multi-modal / Vision & Language) ◦ 効果の高いデジタル広告の作成支援 [Kitada+ KDD’19] ◦ 効果の低いデジタル広告の停止支援 [Kitada+ Appl. Sci.’22] • デザイン生成 AI ◦ 離散拡散モデルで生成されたレイアウトの再調整 [Iwai+ ECCV’24] ◦ LLMによるレイアウトの生成に対する自己修正 [Zhang+ arXiv’24] 自己紹介: 北田俊輔 Shunsuke KITADA 4
 🏠: shunk031.me / 𝕏: @shunk031 画像生成AIにおける 拡散モデルの理論と実践 リサーチサイエンティスト 北田俊輔 www.youtube.com/watch?v =-IPEUOcPTas Pythonで学ぶ画像生成 北田俊輔 インプレス社 https://book.impress.co.j p/books/1123101104 “Python で学ぶ” 画像生成なので • 筋の良い Python コードを 読者の人に書いてほしい • 研究始めたての大学生や我流で コードを書いてきた研究者や 機械学習エンジニアに送りたい という気持ちがあった
5. © LY Corporation “Pythonで学ぶ画像生成” のコラムにねじ込もう！ 5
 無料で公開しています！詳しくは： 🔍「Pythonで学ぶ画像生成 コラム」 もしくは🔗 「https://bit.ly/py-img-gen-note-column」 本文には1ページ 程度しか書けない 😭 ➜ せや！note で 補足公開したろ！
6. © LY Corporation # --- 俺の考えた最強の main.py --- # torch.manual_seed(0) # FIXME: よくわかんないけどコメントアウト # TODO: config.json にしたい（いつか） EPOCHS = 20; BATCH_SIZE = 128; LR = 0.0003; USE_CUDA = True # モデル定義（ファイル分けるのが面倒だった） class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(28*28, 256) self.fc2 = nn.Linear(256, 256) # self.fc3 = nn.Linear(256, 10) # わからん def forward(self, x): x = x.view(-1, 28*28) ... model = Net().device("cuda") # 決め打ち！ train_loader = # ここに最強のデータローダーが爆誕 ... for epoch in range(EPOCHS): # \突然現れる train loop/ total_loss = 0 for i, (x, y) in enumerate(train_loader): x, y = x.device("cuda"), y.device("cuda") optimizer.zero_grad() out = model(x) loss = F.cross_entropy(out, y) loss.backward() optimizer.step() 6
 こんな main.py 書いてませんか？ そもそも “Jupyter Notebook 書き散らしですごめんなさい！” の人は心を改めてくださいね
7. © LY Corporation # --- 俺の考えた最強の main.py --- # torch.manual_seed(0) # FIXME: よくわかんないけどコメントアウト # TODO: config.json にしたい（いつか） EPOCHS = 20; BATCH_SIZE = 128; LR = 0.0003; USE_CUDA = True # モデル定義（ファイル分けるのが面倒だった） class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(28*28, 256) self.fc2 = nn.Linear(256, 256) # self.fc3 = nn.Linear(256, 10) # わからん def forward(self, x): x = x.view(-1, 28*28) ... model = Net().device("cuda") # 決め打ち！ train_loader = # ここに最強のデータローダーが爆誕 ... for epoch in range(EPOCHS): # \突然現れる train loop/ total_loss = 0 for i, (x, y) in enumerate(train_loader): x, y = x.device("cuda"), y.device("cuda") optimizer.zero_grad() out = model(x) loss = F.cross_entropy(out, y) loss.backward() optimizer.step() 7
 こんな main.py 書いてませんか？ とりあえず深夜テンションで ここまで書いた！ これは SoTA だな（確信）
8. © LY Corporation # --- 俺の考えた最強の main.py --- # torch.manual_seed(0) # FIXME: よくわかんないけどコメントアウト # TODO: config.json にしたい（いつか） EPOCHS = 20; BATCH_SIZE = 128; LR = 0.0003; USE_CUDA = True # モデル定義（ファイル分けるのが面倒だった） class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(28*28, 256) self.fc2 = nn.Linear(256, 256) # self.fc3 = nn.Linear(256, 10) # わからん def forward(self, x): x = x.view(-1, 28*28) ... model = Net().device("cuda") # 決め打ち！ train_loader = # ここに最強のデータローダーが爆誕 ... for epoch in range(EPOCHS): # \突然現れる train loop/ total_loss = 0 for i, (x, y) in enumerate(train_loader): x, y = x.device("cuda"), y.device("cuda") optimizer.zero_grad() out = model(x) loss = F.cross_entropy(out, y) loss.backward() optimizer.step() 8
 こんな main.py 書いてませんか？ とりあえず深夜テンションで ここまで書いた！ これは SoTA だな（確信） ちょっとだけ実験するつもりが 1ヶ月後には誰も触れられない コードに大変身
9. © LY Corporation # --- 俺の考えた最強の main.py --- # torch.manual_seed(0) # FIXME: よくわかんないけどコメントアウト # TODO: config.json にしたい（いつか） EPOCHS = 20; BATCH_SIZE = 128; LR = 0.0003; USE_CUDA = True # モデル定義（ファイル分けるのが面倒だった） class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(28*28, 256) self.fc2 = nn.Linear(256, 256) # self.fc3 = nn.Linear(256, 10) # わからん def forward(self, x): x = x.view(-1, 28*28) ... model = Net().device("cuda") # 決め打ち！ train_loader = # ここに最強のデータローダーが爆誕 ... for epoch in range(EPOCHS): # \突然現れる train loop/ total_loss = 0 for i, (x, y) in enumerate(train_loader): x, y = x.device("cuda"), y.device("cuda") optimizer.zero_grad() out = model(x) loss = F.cross_entropy(out, y) loss.backward() optimizer.step() 9
 こんな main.py 書いてませんか？ とりあえず深夜テンションで ここまで書いた！ これは SoTA だな（確信） ちょっとだけ実験するつもりが 1ヶ月後には誰も触れられない コードに大変身 適当にハードコード したネ申パラメータ
10. © LY Corporation # --- 俺の考えた最強の main.py --- # torch.manual_seed(0) # FIXME: よくわかんないけどコメントアウト # TODO: config.json にしたい（いつか） EPOCHS = 20; BATCH_SIZE = 128; LR = 0.0003; USE_CUDA = True # モデル定義（ファイル分けるのが面倒だった） class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(28*28, 256) self.fc2 = nn.Linear(256, 256) # self.fc3 = nn.Linear(256, 10) # わからん def forward(self, x): x = x.view(-1, 28*28) ... model = Net().device("cuda") # 決め打ち！ train_loader = # ここに最強のデータローダーが爆誕 ... for epoch in range(EPOCHS): # \突然現れる train loop/ total_loss = 0 for i, (x, y) in enumerate(train_loader): x, y = x.device("cuda"), y.device("cuda") optimizer.zero_grad() out = model(x) loss = F.cross_entropy(out, y) loss.backward() optimizer.step() 10
 こんな main.py 書いてませんか？ とりあえず深夜テンションで ここまで書いた！ これは SoTA だな（確信） ちょっとだけ実験するつもりが 1ヶ月後には誰も触れられない コードに大変身 適当にハードコード したネ申パラメータ 適切に処理が分離 されておらず入出力 がよくわからん
11. © LY Corporation # FIXME: とりあえず動かすために seed はコメントアウト（なぜ） # torch.manual_seed(0) # TODO: config.json にしたい（いつか） EPOCHS = 20; BATCH_SIZE = 128; LR = 0.0003; USE_CUDA = True class Net(nn.Module): # モデル定義（ファイル分けるのが面倒だった） def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(28*28, 256) self.fc2 = nn.Linear(256, 256) # self.fc3 = nn.Linear(256, 10) # わからん def forward(self, x): x = x.view(-1, 28*28) ... model = Net().device("cuda") # 決め打ち！ train_loader = # ここに最強のデータローダーが爆誕 ... for epoch in range(EPOCHS): # \突然現れる train loop/ total_loss = 0 for i, (x, y) in enumerate(train_loader): x, y = x.device("cuda"), y.device("cuda") optimizer.zero_grad() out = model(x) loss = F.cross_entropy(out, y) loss.backward() optimizer.step() “再現性” の罠 • 環境・乱数・設定がバラバラ ◦ どのように実行したか忘れる ◦ 環境は uv ✨ で幸せになれるが… “拡張性” の罠 • モデルを差し替えるたびにコード修正 ◦ モデルを比較するだけで一苦労 “テスト” の罠 • いつのまにか壊れている ◦ どこが壊れたかわからない これら、pytest で解決できるかも？ 11
 main.py が “最強” になる 3 つの罠
12. © LY Corporation # train 時の依存は model, data_loader, epochs, … etc. def train( model: nn.Module, data_loader: DataLoader, epochs: int, lr: float ) -> nn.Module: optimizer = optim.Adam(model.parameters(), lr=lr) for _ in range(epochs): for x, y in data_loader: optimizer.zero_grad() loss = F.cross_entropy(model(x), y) loss.backward() optimizer.step() return loss def test_train_smoke(): # pytest 実行時はこの関数が動作対象 model = Net() dataset = make_fake_dataset(size=128) loss = train(model, dataset, epochs=1, lr=1e-3) assert loss <= 0.01 # loss が十分小さくなっているかテスト テストを書くことで “何が入力で” “何を期待するか” を明示できる • 実験設計と同じ行為 • コードの設計 = 研究の再現設計 train 関数への切り出しとテスト実行 • 機械学習モデルの実験で 重要な train のフェーズを明示 ◦ 関数の入出力が決まることで 実験の設計が明らかになる 12
 テストを書く = 実験設計を明示する
13. © LY Corporation “依存性注入” と呼ばれているが 要は “オブジェクト注入” • オブジェクトの 「作成」と「利用」を分ける “外から渡す” ことで コードを柔軟に保つ • train 関数に注目 • model, data_loader 等に依存 pytest の fixture 機構が DI 機能を提供 • オブジェクトの作成を fixture 化 • 作成・利用と分離で見通しup 13
 依存性注入 (dependency injection; DI) の発想 def train( model : nn.Module, data_loader : DataLoader, epochs : int, lr: float ) -> nn.Module: optimizer = optim.Adam(model.parameters(), lr=lr) for _ in range(epochs): for x, y in data_loader: optimizer.zero_grad() loss = F.cross_entropy(model(x), y) loss.backward() optimizer.step() return loss def test_train_smoke(): model = Net() dataset = make_fake_dataset(size=128) loss = train(model, dataset, epochs=1, lr=1e-3) assert loss <= 0.01 # loss が十分小さくなっているかテスト model data_loader epochs lr
14. © LY Corporation 実験条件を fixture に明示 • seed, dataset, model ... pytest により以下が容易に • 実験環境準備 ◦ fix_seed によるシード値固定 • 依存注入 ◦ 実験したいパラメータの比較 • 実験再実行・並列実行 ◦ test_ 関数を起点とした実行 ◦ parametrize + pytest-xdist で テストの並列実行が可能 14
 テストを書くことが 研究を再現可能にする @pytest.fixture def seed() -> int: return 19950815 # シード値は小倉唯さんの誕生日 @pytest.fixture(autouse=True) def fix_seed(seed: int): # シード値を受け取って固定 random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) @pytest.fixture def dataset(transform: transforms.Compose) -> Dataset: return datasets.MNIST(train=True, transform=transform) @pytest.fixture def model(device: torch.device) -> nn.Module: return nn.Sequential(nn.Flatten(), nn.Linear(28*28, 10)) @pytest.mark.parametrize("lr", [1e-2, 1e-3, 1e-4]) @pytest.mark.parametrize("batch_size", [32, 64]) def test_train_smoke(dataset, model, lr, bs, device): # 適宜型付けして ね data_loader = DataLoader(dataset, batch_size=bs, shuffle=True) optim = torch.optim.Adam(model.parameters(), lr=lr) for x, y in loader: x, y = x.to(device), y.to(device) opt.zero_grad() … # 以降 train loop
15. © LY Corporation pytest をテストだけでなく 実験のタスクランナーとして使う 特徴 • 依存を fixture / 引数で明示可能 • テスト単位で切り替えが簡単 ➜ @pytest.mark.parametrize • pytest 1 コマンドだけ 覚えていれば全実験を再現可能 ◦ python main.py --hoge fuga … のようなコマンドを忘れていてもOK 15
 pytest は 軽量な実験実行くん @pytest.fixture def seed() -> int: return 19950815 # シード値は小倉唯さんの誕生日 @pytest.fixture(autouse=True) def fix_seed(seed: int): # シード値を受け取って固定 random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) @pytest.fixture def dataset(transform: transforms.Compose) -> Dataset: return datasets.MNIST(train=True, transform=transform) @pytest.fixture def model(device: torch.device) -> nn.Module: return nn.Sequential(nn.Flatten(), nn.Linear(28*28, 10)) @pytest.mark.parametrize("lr", [1e-2, 1e-3, 1e-4]) @pytest.mark.parametrize("batch_size", [32, 64]) def test_train_smoke(dataset, model, lr, bs, device): # 適宜型付けして ね data_loader = DataLoader(dataset, batch_size=bs, shuffle=True) optim = torch.optim.Adam(model.parameters(), lr=lr) for x, y in loader: x, y = x.to(device), y.to(device) opt.zero_grad() … # 以降 train loop $ pytest # もしくは uv run pytest
16. © LY Corporation MNIST 訓練コードのスモークテスト • とりあえず動くか確かめる 最小のテストのこと 元はハードウェアの用語： 「電源を入れて煙が出なければOK」 動作説明 • test_ とつけるだけで、pytest が実験環境を自動で構築 • tmp_path に実験用のディレクトリが自動生成 • テストは実行の最小単位= 研究の最小構成単位 その他 pytest のチュートリアルは：Get Started - pytest documentation https://docs.pytest.org/en/stable/getting-started.html 16
 main.py の代わりに 書く最小実験 def test_train_smoke( tmp_path: pathlib.Path # 自動で実験用ディレクトリが作成されて注入され る ): cfg = make_cfg() ds = make_dataloaders(cfg, tmp_path) model = make_model(cfg) trainer = make_trainer(cfg) metrics = trainer.train(model, ds) assert "loss" in metrics
17. © LY Corporation • @pytest.mark.parametrize(“lr, epochs”, [(1e-3, 1), (1e-4, 3)]) 説明 • 設定の組み合わせを自動で総当り • 実験探索とテストが同じ構文で書ける 各パラメータの組み合わせを実行可能 17
 pytest の parametrize で実験探索 ... @pytest.mark.parametrize("lr", [1e-2, 1e-3, 1e-4]) @pytest.mark.parametrize("batch_size", [32, 64]) def test_train_smoke(dataset, model, lr, bs, device): # 適宜型付けして ね data_loader = DataLoader(dataset, batch_size=bs, shuffle=True) optim = torch.optim.Adam(model.parameters(), lr=lr) for x, y in loader: x, y = x.to(device), y.to(device) opt.zero_grad() … # 以降 train loop 学習率の候補を列挙 バッチサイズの 候補を列挙
18. © LY Corporation • 以下のように段階を踏んで実験コードを育てていくのをおすすめします 18
 実験コードを育てる 5 フェーズ 段階 内容 目的 1 pytest でスモークテストを回す 実験の再現性を確保 2 パラメータを注入 実験探索を自動化 3 関数化・共通化 再利用性を担保 4 モジュール抽出 設計を安定化 5 ライブラリ化・CLI 化 論文公開と同時に 使ってもらえるように main.py tests/main_test.py src/module.py
19. © LY Corporation pytest を使うと • fixture を通じて依存を分離 • 関数が純粋関数化 • 実験が小さく再現可能に 19
 pytest で始めると自然に設計が良くなる # 依存が中に埋まっている model = nn.Sequential(nn.Flatten(), nn.Linear(784, 10)) optimizer = optim.Adam(model.parameters(), lr=1e-3) dataset = MNIST(train=True, transform=transforms.ToTensor()) train_loader = DataLoader(dataset, batch_size=64, shuffle=True) for epoch in range(5): for x, y in train_loader: optimizer.zero_grad() loss = nn.functional.cross_entropy(model(x), y) loss.backward() optimizer.step() @pytest.fixture def dataset() -> Dataset: return MNIST(train=True, transform=transforms.ToTensor()) @pytest.fixture def model() -> nn.Module: return nn.Sequential(nn.Flatten(), nn.Linear(784, 10)) def train(model: nn.Module, loader: DataLoader) -> float: opt = torch.optim.Adam(model.parameters(), lr=1e-3) for x, y in loader: opt.zero_grad() loss = F.cross_entropy(model(x), y) loss.backward(); opt.step() return loss.item() def test_train_smoke(model: nn.Module, dataset: Dataset): loader = DataLoader(dataset, batch_size=64, shuffle=True) loss = train(model, loader) assert loss < 10 ✅ After：pytest スタイル （依存が明示され、構造が見える） 🧨 Before：main.py スタイル （手続き的・密結合）
20. © LY Corporation 実行は pytest コマンドだけ！ • 1ヶ月後、半年後でも覚えていられる • 実行パラメータは fixture として残されている • 誰でも簡単に動かせる・コードも読める 20
 何も覚えていなくても動く実行環境 $ pytest # もしくは uv run pytest Prompt: 「何も覚えていなくても動く Python 実行環境」に関し て、研究者が困っていそうな感じの画像を生成してください。文字 は書かなくていいです。もやもやを出してください。
21. © LY Corporation Before After 依存が隠蔽されている fixture / 引数で明示 コードが壊れやすい テストで守られている 再現が困難 pytest で実験構成を固定 実行手順が不透明 pytest コマンドだけで動く • テストを書くことは 研究を設計することである • オレオレ最強 main.py から抜け出して、 pytest を最初の実験ランナーにしよう おすすめ pytest plugin • pytest-sugar: pytest の結果を見やすくしてくれる君 • pytest-xdist: pytest のテストケースを並列で実行してくれる君 • pytest-lazy-fixture: pytest.fixture を pytest.mark.paramtrize に渡せるようにする君 関連研究 • Pinjected: 研究開発向けPythonライブラリ(Dependency Injection等) https://zenn.dev/proboscis/articles/4a10d26b13a940 21
 🍵 まとめ