wandbのARIAで全自動で研究させてみた。

はじめに

こんにちは、私は機械知能研究室に所属し、CV分野の研究をしています

機械知能研究室ではロボットの視覚機能や自律的に行動するための知能システムについての研究に取り組んでいます (ホームページは下記のリンクから)

今回はWeights & Biases(wandb)の新機能AI Research and Iteration Agent(ARIA)を用いて、全自動で研究させてみました!!

wandbサイト内のARIAというAIチャットに、「~~の実験をして」と投げるだけで、スクリプトの作成、実行をしてくれます...それだけではありません!!実行結果をwandbに保存し、結果からARIAがフィードバックをして、スクリプトの変更をし、再度実行をするという最適化を繰り返してくれます!!

様々なAIツールがありますが、本記事を参考にARIAを用いた実験も検討してみてください。

注意事項(オプトアウト設定について)

現在、私は学生の無料版としてwandbを使用しています。
無料版ではARIAのオプトアウト設定ができず、会話履歴などが学習に使用されます。企業での使用や、未発表研究に使用する場合、本記事を読んだ後に、有料版を使用するなどを検討してみてください。

ARIAとlaunchについて

ARIAとは、wandbのサイト内で使用できる対話型AIです。
wandbのサイト内で、結果やログに対して考察やレポートの生成などができます。
今回はこのARIAをwandb launchという機能と連携して、コードの生成、実行、フィードバックをしてスクリプトを更新という流れを自動で繰り返して実験をしてもらうことを目的とします。
つまり、ARIAに「この~の実験をしてください」と指示したら、そのタスクを何回も実験をし、試行錯誤しながら最適化してくれるということです。

wandb launchとは、実行したいjobを設定すると、そのjobをキューで順番に並べて、エージェントがキューの順番に処理するというパイプラインのようなイメージです。
人間が作成したスクリプトをjobにするのではなく、ARIAにプロンプトを投げることで、ARIAがスクリプトの生成をし、jobに追加してくれるという仕組みで、自動化します。

実験概要

目的

• Fashion-MNISTのテスト精度95%以上にするハイパーパラメータを見つけること
  (かつ最もエポック数が少ないパラメータ)

条件

• モデル：畳み込み層2層、全結合層1層のCNN
• 最適化アルゴリズム：SGD
• データ拡張なし

最適化対象

• 学習率
• 学習率スケジューラ
• エポック数
• 重み減衰率
• momentum

グリッドサーチなどではなく、実行後に毎回結果を分析し、次のパラメータを決めてもらいます。

実験の流れ

今回は、ローカルのGPUを積んだPCでスクリプトを実行をします。
つまり、ARIAにプロンプトを投げ、wandb launchを使用して、ローカルのPCのGPUを自動的に動かすことをします。

1.dockerの準備

インストールスクリプトのダウンロード

curl -fsSL https://get.docker.com -o get-docker.sh

スクリプトの実行（Dockerのインストール）

sudo sh get-docker.sh

コンテナ内からGPUを認識させるためのツールキッド

APTキーリング用フォルダの作成

sudo mkdir -p /etc/apt/keyrings

NVIDIA公式リポジトリのGPG鍵とパッケージリストの登録

curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /etc/apt/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/etc/apt/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

パッケージリストの更新とツールキットのインストール

sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit

Dockerのランタイム設定をNVIDIA用に構成し、Dockerサービスを再起動

sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

ARIAが操作できるように、sudoを打たなくてもいいようにする

現在のユーザーをdockerグループに追加し、設定を即時反映

sudo usermod -aG docker $USER
newgrp docker

【テスト1】sudoなしでDockerコンテナが起動できるか確認

docker run --rm hello-world

【テスト2】コンテナ内からGPU（NVIDIA-SMI）が正常に認識されるか確認

＊私の環境の場合のコマンド
docker run --rm --gpus all nvidia/cuda:12.2.2-base-ubuntu22.04 nvidia-smi

wandb Launch Agent用のローカルGPUのデフォルト設定

設定ディレクトリの作成

mkdir -p ~/.config/wandb

デフォルトの起動引数（args）に --gpus=all を指定した設定ファイルの書き込み

cat <<EOF > ~/.config/wandb/launch-config.yaml
docker:
  args:
    - "--gpus=all"
EOF

2.ローカル環境との紐付け

webブラウザのwandbから、リンクにlaunchを追加(リンクのhomeは削除)して、https://wandb.ai/launchに移動します。
Create a queueからキューを作成します。

次に、ローカルマシンでWandb Launch Agentを管理・起動する専用のvenvを作成します。
＊venvの作成、アクティベート方法は割愛

作成した仮想環境にwandbのライブラリをインストールしてログイン

pip install wandb

wandb login

wandb launch Agentの起動

wandb launch-agent -e [ユーザー名] -q [キューの名前]

これで紐付け完了です。

3.ARIAに指示

いくつか実験を試したのですが、最終的に下記のプロンプトで実験することにしました。
(データセットによっては、読み込みが長すぎるのかエラーで読み込めない場合がありました)

プロンプトで特別意識したことはありませんが、わかりやすく端的に書くことと、起きてしまいそうなミス(グリッドサーチになってしまうこと)を事前に伝えています。

目的：Fashion-MNISTでテスト精度95%を達成するハイパーパラメータを探索してください。
条件：最も小さいエポック数を探索してください(効率性)。また、モデルは畳み込み層2層と全結合層1層のCNNで固定です。オプティマイザもSGDで固定です。
探索するパラメータは、学習率、重み減衰率、momentum, エポック数、スケジューラ(種類やハイパーパラメータも探索対象です。)です。

グリッドサーチではなく、毎回実行後に結果をあなたが確認、考察して、次のハイパーパラメータを決めてください。

一連の実験をレポートに記録してください。

上記をARIAのプロンプトに投げると、様々な作業が始まり、chat画面でThinkingやshellといった文字が表示されます。
このshellが、実際にpythonスクリプトの作成や実行をしている部分です。

ターミナルにはログが表示されます。

gpu_name': 'NVIDIA GeForce RTX 3080', 'conv_layers_actual': 6, 'channels_actual': '64,128,256', 'parameter_count': 2330442}
epoch=01/45 lr=0.004450 train_acc=0.8227 val_acc=NA val_loss=NA
epoch=02/45 lr=0.007186 train_acc=0.8893 val_acc=NA val_loss=NA
epoch=03/45 lr=0.011378 train_acc=0.9028 val_acc=NA val_loss=NA
epoch=04/45 lr=0.016520 train_acc=0.9119 val_acc=NA val_loss=NA
epoch=05/45 lr=0.021991 train_acc=0.9183 val_acc=NA val_loss=NA
epoch=06/45 lr=0.027132 train_acc=0.9227 val_acc=NA val_loss=NA
epoch=07/45 lr=0.031321 train_acc=0.9277 val_acc=NA val_loss=NA
epoch=08/45 lr=0.034054 train_acc=0.9325 val_acc=NA val_loss=NA
epoch=09/45 lr=0.035000 train_acc=0.9354 val_acc=NA val_loss=NA
epoch=10/45 lr=0.034933 train_acc=0.9379 val_acc=NA val_loss=NA
epoch=11/45 lr=0.034734 train_acc=0.9414 val_acc=NA val_loss=NA
epoch=12/45 lr=0.034403 train_acc=0.9440 val_acc=NA val_loss=NA
epoch=13/45 lr=0.033944 train_acc=0.9465 val_acc=NA val_loss=NA
epoch=14/45 lr=0.033360 train_acc=0.9477 val_acc=NA val_loss=NA
epoch=15/45 lr=0.032655 train_acc=0.9498 val_acc=NA val_loss=NA
s=NAh=16/45 lr=0.031835 train_acc=0.9520 val_acc=NA val_los
epoch=17/45 lr=0.030906 train_acc=0.9535 val_acc=NA val_loss=NA
epoch=18/45 lr=0.029875 train_acc=0.9552 val_acc=NA val_loss=NA
epoch=19/45 lr=0.028749 train_acc=0.9575 val_acc=NA val_loss=NA
epoch=20/45 lr=0.027539 train_acc=0.9584 val_acc=NA val_loss=NA
epoch=21/45 lr=0.026252 train_acc=0.9598 val_acc=NA val_loss=NA
epoch=22/45 lr=0.024898 train_acc=0.9617 val_acc=NA val_loss=NA
epoch=23/45 lr=0.023488 train_acc=0.9628 val_acc=NA val_loss=NA
epoch=24/45 lr=0.022033 train_acc=0.9654 val_acc=NA val_loss=NA
epoch=25/45 lr=0.020543 train_acc=0.9653 val_acc=NA val_loss=NA
epoch=26/45 lr=0.019030 train_acc=0.9676 val_acc=NA val_loss=NA
epoch=27/45 lr=0.017505 train_acc=0.9688 val_acc=NA val_loss=NA
epoch=28/45 lr=0.015981 train_acc=0.9701 val_acc=NA val_loss=NA
epoch=29/45 lr=0.014468 train_acc=0.9718 val_acc=NA val_loss=NA
epoch=30/45 lr=0.012979 train_acc=0.9728 val_acc=NA val_loss=NA
epoch=31/45 lr=0.011523 train_acc=0.9756 val_acc=NA val_loss=NA

もちろん、実行中の精度や損失などの記録はwandbのサイト上で確認できます!!
ここで、ARIAは、一回実行して終わりではなく、１回目の実験が終了したら、私の指示を待たずに結果からフィードバックして２回目の実験を始めます。
(私の実験では、10回以上繰り返した後にやっと一度終了しました。)

4.結果

実行が終わると、chat画面に実験結果や、レポートのリンクが記載されます。
レポートはしっかり記載されており、グラフもあります。

今回、合計で22回の試行をしてくれました。
途中7試行目や15試行目ほどで、「この設定では限界だよ」と言われたので、データ拡張やモデルのアーキテクチャの制限を緩和しました。
そして、最終的に目標の95%以上を達成しました。

上記のグラフは、各試行ごとのテストデータの精度を表しています。
このグラフの結果が最も大切と言っても過言ではないです!!
古い実行(下)から徐々にテスト精度が上がっていることがわかります。
ARIAが実行毎にフィードバックをして改善できている結果と言えるでしょう。

まとめ

今回は、wandbの新機能ARIAを用いて、全自動での研究を試みました。
結果としては、"想像より実用化できそうだった"というのが私の感想です。
ARIAに目的のプロンプトを投げるだけで、基本的に放置で動作します。
また、wandbに記録され、レポートを生成してくれるのが非常に良いと思いました。
しかし、データセットの読み込みがうまく行かないことを始めとした、いくつか苦戦する部分もありましたので、すべてがスムーズにできたわけではありません。
本記事を参考にARIAを活用していただけたら幸いです。

参考にしたサイト