PyTorch Project Template

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

Features

Docker + uvで環境構築
PyTorchのDistributed Data Parallel(DDP), Data Parallel, Fully Shared Distributed Parallel(FSDP)によるマルチGPU Training
Multi Node Multi GPU Trainingのサポート
MLflowとwandbを使った実験管理
OmegaConfを使ったコンフィグ管理
データセットの一部をRAMにキャッシュする機能
学習の再開機能

Environments

Python 3.11
CUDA 12.4
PyTorch 2.5.1

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 file path]
./docker/run.sh python train.py config/model/ResNet.yaml

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

Single Node Multi GPU Training

シングルノード・複数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"

Multi Node Multi GPU Training

マルチノード・複数GPUを用いた学習を行う場合は，実行するコマンドのpythonを消して，前に./multinode.sh [ノード数] [GPU数] [ジョブID] [ノードランク] [マスターノードのホスト名:マスターノードのポート]を入れ， gpu.use="0,1"のように，Configの値を変更します．

# Master Node
./docker/run.sh ./multinode.sh 2 4 12345 0 localhost:12345 train.py config/model/ResNet.yaml gpu.use=0,1,2,3

# Worker Node
./docker/run.sh ./multinode.sh 2 4 12345 1 192.168.1.10:12345 train.py config/model/ResNet.yaml gpu.use=4,5,6,7

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

Test Source

./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]

実装済みモジュールの確認

python script/show_options.py

DATASET_REGISTRY
Registry of DATASET:
╒══════════════╤════════════════════════════════════════════════╕
│ Names        │ Objects                                        │
╞══════════════╪════════════════════════════════════════════════╡
│ DummyDataset │ <class 'src.dataloaders.dataset.DummyDataset'> │
╘══════════════╧════════════════════════════════════════════════╛
EVALUATOR_REGISTRY
Registry of EVALUATOR:
╒════════════════╤══════════════════════════════════════════════════╕
│ Names          │ Objects                                          │
╞════════════════╪══════════════════════════════════════════════════╡
│ DummyEvaluator │ <class 'src.evaluator.evaluator.DummyEvaluator'> │
╘════════════════╧══════════════════════════════════════════════════╛
MODEL_REGISTRY
Registry of MODEL:
╒════════════╤═══════════════════════════════════════╕
│ Names      │ Objects                               │
╞════════════╪═══════════════════════════════════════╡
│ DummyModel │ <class 'src.models.model.DummyModel'> │
╘════════════╧═══════════════════════════════════════╛
BACKBONE_REGISTRY
['bat_resnext26ts',
 'beit_base_patch16_224',
 'beit_base_patch16_384',
 'beit_large_patch16_224',
...

Structure

.//
├── config/            # 実験とモデルの設定ファイル
│   └── __base__/      # 基本設定ファイル
├── dataset/           # データセットを保存するディレクトリ（Dockerコンテナにマウント）
├── doc/               # ドキュメントファイルと実験結果
├── docker/            # Dockerセットアップとユーティリティスクリプト
├── etc/               # その他のファイル
├── notebook/          # 探索と分析用のJupyterノートブック
├── result/            # 学習結果、チェックポイント、ログ
├── script/            # 様々なタスク用のユーティリティスクリプト
│   ├── aggregate_mlflow.py  # MLflow結果を集計するスクリプト
│   ├── clean_result.py      # resultディレクトリをクリーンアップ
│   ├── edit_configs.py      # 設定ファイルの一括編集
│   ├── show_*.py            # 様々な情報を表示するスクリプト
│   ├── run_mlflow.sh        # MLflow UIを起動
│   ├── run_notebook.sh      # JupyterLabを起動
│   └── run_test.sh          # テストを実行
├── src/               # ソースコード
│   ├── config/        # 設定の処理
│   ├── dataloaders/   # データ読み込みユーティリティ
│   ├── evaluator/     # モデル評価コード
│   ├── models/        # モデル定義
│   ├── optimizer/     # 最適化アルゴリズム
│   ├── scheduler/     # 学習率スケジューラ
│   ├── transform/     # データ変換
│   ├── utils/         # ユーティリティ関数
│   ├── sampler.py     # データサンプリングユーティリティ
│   ├── tester.py      # テストループの実装
│   ├── trainer.py     # 学習ループの実装
│   └── types.py       # 型定義
├── tests/             # テストコード
├── README.md
├── template.env       # 環境変数のテンプレート
├── train.py
├── test.py
├── torchrun.sh        # 単一ノードでの分散学習用スクリプト
├── multinode.sh       # マルチノード分散学習用スクリプト
├── pyproject.toml
└── uv.lock

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Repository files navigation

PyTorch Project Template

Features

Environments

Install

MLflow

Optional: Wandb

Optional: Slackによる通知

Usage

Train

Single Node Multi GPU Training

Multi Node Multi GPU Training

Train Option: RAM Cache

Train Option: Resume Training

Test

Scripts

Config一括編集

MLflow UI

JupyterLab

Test Source

Clean result

集計

実装済みモジュールの確認

Structure

About

Uh oh!

Releases

Packages

Uh oh!

Languages

Name		Name	Last commit message	Last commit date
Latest commit History 601 Commits
.devcontainer		.devcontainer
.vscode		.vscode
config		config
dataset		dataset
doc		doc
docker		docker
etc		etc
notebook		notebook
result		result
script		script
src		src
tests		tests
.dockerignore		.dockerignore
.gitignore		.gitignore
README.md		README.md
multinode.sh		multinode.sh
pyproject.toml		pyproject.toml
template.env		template.env
test.py		test.py
torchrun.sh		torchrun.sh
train.py		train.py
uv.lock		uv.lock

mjun0812/PyTorch-Project-Template

Folders and files

Latest commit

History

Repository files navigation

PyTorch Project Template

Features

Environments

Install

MLflow

Optional: Wandb

Optional: Slackによる通知

Usage

Train

Single Node Multi GPU Training

Multi Node Multi GPU Training

Train Option: RAM Cache

Train Option: Resume Training

Test

Scripts

Config一括編集

MLflow UI

JupyterLab

Test Source

Clean result

集計

実装済みモジュールの確認

Structure

About

Resources

Uh oh!

Stars

Watchers

Forks

Releases

Packages 0

Uh oh!

Languages

Packages