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を使ったコンフィグ管理
• データセットの一部をRAMにキャッシュする機能
• 学習の再開機能

Environments

• Python 3.11
• CUDA 12.4
• PyTorch 2.5.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]以下のディレクトリに保存されます．

Train Option: RAM Cache

データセットの一部をRAMにキャッシュする機能があります。キャッシュはtorch.Tensorのみ対応しています。

./docker/run.sh python train.py config/model/ResNet.yaml gpu.use=1 use_ram_cache=true ram_cache_size_gb=16

この機能を使うには、datasetの実装を以下のように工夫する必要があります。

if self.cache is not None and idx in self.cache:
    image = self.cache.get(idx)
else:
    image = read_image(str(image_path), mode=ImageReadMode.RGB)
    if self.cache is not None:
        self.cache.set(idx, image)

Train Option: Resume Training

学習の終了後、もしくは学習を中断した場合に、 結果のディレクトリに保存されているconfig.yamlである、 result/[train_dataset.name]/[日付]_[model.name]_[dataset.name]_[tag]/config.yamlを指定して実行すると，学習を再開できます。 この時、元のepochが100だった場合でも、epoch=150をコマンドラインで指定して実行すると、configが上書きされて150epochまで学習が継続されます。

./docker/run.sh python train.py config/config.yaml # 100epochまで完了

# 上記の結果を利用して学習を再開 or 継続
./docker/run.sh python train.py result/ImageNet/hoge_hoge/config.yaml epoch=150 gpu.use=7 # 150epochまで学習を継続

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

Regression 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]