はじめまして、AI技術開発部の廣安です。
私たちのチームでは、機械学習を使用したシステムをRustで開発しています。本記事では、機械学習システムのテストに役立つスナップショットテストという手法をRustで実践する方法を紹介します。
ソフトウェアのテストは一般に実行結果と期待値と比較することで行います。しかし、私たちが開発している機械学習システムでは、システムの実行結果が行列データや座標情報であったり、またモデルを差し替えることで結果が変わったりするため、期待値を手動で定義してテストコードを運用するのは手間がかかります。
それを解決するための一手法としてスナップショットテストを使うことができます。この手法は、あるテストを実行した時点の結果をファイルに保存し、次の実行からはその結果を期待値としてテストをおこなうものです。これにより期待値の管理の手間が軽減できます。
スナップショットテストは、React等のフレームワークで生成されたWebUIのテストとしても使われています。この場合は、ある時点のUIのキャプチャをスナップショットテストに使い、意図せずUIの見た目が変わる事を防ぐという使い方になります。
私たちは機械学習システムのリグレッションテスト(CIで自動実行)の中にスナップショットテストを導入しており、コードの変更によって意図せずシステムの結果が変わってしまった場合はすぐに検出できる状態にしています。
Rustのスナップショットテストのライブラリとして、insta を採用しました。他のライブラリと比較してシンプルで使いやすそうと感じたためです。 instaはスナップショットテストを行うためのマクロを提供しています。テストのコード中にそのマクロを適宜書くだけでスナップショットとの比較とアサーションまでやってくれます。
下記は簡単な使用例です。ここでは、 Vec<i32> を返す関数をテストするケースを想定しています。 assert_debug_snapshot に引数として指定した変数がsnapshotテストの対象となり、過去の結果と比較されます。
use insta::assert_debug_snapshot;
fn func() -> Vec<i32> {
vec![1, 2, 3]
}
#[test]
fn test_func() {
let val = func();
// valが過去のsnapshotと一致しない場合はテスト失敗となる
assert_debug_snapshot!(&val);
}
比較元のスナップショットファイルは最初無いので、テストを実行し作成する必要があります。 cargo insta test --review コマンドでテストを実行するとスナップショットファイルが自動的に作成され、対話的にレビューが始まります。上のテストだと下のような画面が表示されます。
acceptする事でそのスナップショットが確定され、以降のテスト実行時に使われるようになります。
cargo test コマンドでテストが実行されると、マクロの中でスナップショットとの比較が行われます。もし一致しなければ、テストは失敗します。
機械学習モデルを更新した、関数の戻り値が変わった、等の理由でスナップショットを更新したい場合、同じようにcargo insta test --reviewを実行する事で対応できます。
私たちはスナップショットもgitリポジトリで管理し、以下のような運用で利用しています。
PRのレビュー時、スナップショットの更新が正しいかを判断しやすくするため、テスト結果を可視化して問題が無い事を確認しています。Rustでのテスト時の結果の可視化については弊社の記事も参考にしてください。
実際にinstaがどのような仕組みで動いているか簡単に説明します。
上で紹介した assert_debug_snapshot!ですが、このマクロに渡す引数の型はDebug traitを実装している必要があります。レビューが実行/承認されると、引数の内容がDebug traitに対して実装された文字列表現としてスナップショットファイルに書き出されます。つまり、このマクロを使用した場合に出力されるスナップショットファイルの内容は dbg!(&val) のようなDebug traitを使った出力と同じになります。
instaの内部実装では、このDebug traitで表現される文字列同士を比較して一致確認をする仕組みになっています。同じようなマクロとしてassert_display_snapshot!がありますが、こちらではDebug traitではなくDisplay traitで表現される文字列が使用されます。
この他にもJSONやYAMLとして保存する事ができるassert_json_snapshot!やassert_yaml_snapshot!といったマクロもあります。こちらは指定できる型の制約としてserdeのSerializeを実装する必要があります。
これらのマクロは対象となる型が実装しているtraitや、スナップショットの出力に使いたいフォーマット形式に応じて使い分けると良いと思います。
本記事では、スナップショットテストとRustのライブラリinstaの簡単な紹介を行いました。
instaは手軽にスナップショットテストを記述して利用する事ができるので、機械学習システムなど、期待値の定義が面倒なテストのリグレッションテスト等に導入するのは有用だと思います。便利と感じた方はぜひ使ってみてください。
Mobility Technologiesでは私たちのチームメンバを募集しています。興味の在る方はぜひご連絡ください。
エッジAIエンジニア | 株式会社Mobility Technologies
最後まで読んでいただきありがとうございました。