PyTorch Project Template

PyTorchのProjectテンプレートです．

Features

• Docker + uvで環境構築
• PyTorchのDistributed Data Parallel(DDP), Data Parallel, Fully Shared Distributed Parallel(FSDP)によるマルチGPU Training
• MLflowとwandbを使った実験管理
• OmegaConfを使ったコンフィグ管理

Environments

• Python 3.11
• CUDA 12.1
• PyTorch 2.4.0

Install

環境構築はDockerで行います． Dockerコンテナに./datasetをマウントするため，この中に各データセットのディレクトリを入れてください。 Docker Imageのビルドは以下のコマンドで行えます。

./docker/build.sh

MLflow

本テンプレートでは，MLflowによる実験管理が行えます。 MLflowのデータをローカルに保存する場合と，外部のサーバに送信する場合の両方に対応しています． デフォルトはローカル保存となり，result/mlrunsに保存されます．

外部サーバを利用する場合は，dotenvに設定を書き込む必要があります． template.envをコピーして利用してください．

$ cp template.env .env
$ vim .env

SLACK_TOKEN="HOGE"
MLFLOW_TRACKING_URI=""

# Basic Auth
# MLFLOW_TRACKING_USERNAME=""
# MLFLOW_TRACKING_PASSWORD=""

ローカルで保存している場合，mlflow server(ui)のコマンドは以下です．

./docker/run.sh --mlflow-ui ./script/run_mlflow.sh

Optional: Wandb

Wandbによる実験管理も行えます。 .envのWANDB_API_KEYを設定し、configのwandb.useをtrueにすれば結果が送信されます。

./docker/run.sh python train.py config/dummy.yaml wandb.use=true wandb.project_name="hoge"

Optional: Slackによる通知

学習の終了や評価の終了時にSlackに通知を行うことができます． 通知を行うには.envにSlackのトークン(Webhookではない)を書き込む必要があります． デフォルトでは，通知はchannel="#通知", username="通知"で行われます． .envのSLACK_TOKENにAPI tokenを入れて下さい。

Usage

Dockerコンテナ内でコマンドを実行します． そのため，実行するコマンドの先頭に./docker/run.shをつけてください．

./docker/run.sh python train.py config/model/model.yaml gpu.use=1
./docker/run.sh python test.py result/20220911/config.yaml gpu.use=1

yaml内のConfigの値をCLIから変更することもできます．以下のように，.で連結して，=で値を指定してください．

gpu:
  use: 1

./docker/run.sh python train.py config/model/ResNet.yaml gpu.use=2

Train

学習を行う場合は，train.pyを使います．このスクリプトでは，学習が終わった後にテストも一緒に行われます． train.pyではconfig以下のyamlファイルを指定します．

./docker/run.sh python train.py config/model/ResNet.yaml

複数GPUを用いた学習を行う場合は，実行するコマンドのpythonを消して， 前に./torchrun.sh [GPU数]を入れ，gpu.use="0,1"のように，Configの値を変更します． この時，GPUのIDの順番はnvidia-smiコマンドで並ぶPCIeの順番になっています．

./docker/run.sh ./torchrun.sh 4 train.py config/model/ResNet.yaml gpu.use="0,1,2,3"

学習結果はresult/[train_dataset.name]/[日付]_[model.name]_[dataset.name]_[tag]以下のディレクトリに保存されます．

Test

評価を行うスクリプトはtest.pyです． 学習で動かすtrain.pyでも評価は実行されますが，手動で行う場合はこちらを使用してください． test.pyでは，学習結果の保存されているディレクトリにあるconfig.yamlを第1引数に指定します．

./docker/run.sh python test.py result/ImageNet/hoge_hoge/config.yaml gpu.use=7

上記のコマンドは，学習時のログでも表示されています．学習時のログは， result/[train_dataset.name]/[日付]_[model.name]_[dataset.name]_[tag]/train.logに保存されています．

テスト時のログやデータは result/[train_dataset/name]/[日付]_[model.name]_[dataset.name]_[tag]/runs/以下に， ディレクトリが作成され，その中に保存されます．

Scripts

Config一括編集

./docker/run.sh python script/edit_configs.py [config_path or recursive directory] "params.hoge=aa,params.fuga=bb"

MLflow UI

./docker/run.sh --mlflow-ui ./script/run_mlflow.sh

JupyterLab

./script/run_notebook.sh

Test

./docker/run.sh ./script/run_test.sh

Clean result

MLflowで削除された実験結果をローカルの./result/から削除する

./docker/run.sh
# In container
python script/clean_result.py | xargs -I{} -P 2 rm -rf {}

集計

MLflowの結果を集計し、./doc/result_csv以下に保存する。

./docker/run.sh python script/aggregate_mlflow.py [dataset_name or all]