PyTorchLightning
https://scrapbox.io/files/642730b2f4f5ed001b190317.png
Lightning(稲妻)と言っているだけあり、高速に多量の実験をしやすくするためのライブラリ。
学習自体が早くなるのではなくて、学習までのコード変更にかかる時間が短縮できる。
ただ使用前の注意として、PyTorchLightningを使用する上で必要になる知識は多い。
1度しかモデルを訓練しないのであればPyTorchLightningは非効率。
だけどそんなことはなく、何度も実験するので導入する価値がある。
逆に言えば、アドホックな実装が多くなってPyTorchで作成するのに比べて時間がかかりすぎるのであればPyTorchに戻るのも一手
PyTorchLightningを使う上で,普遍性が失われて多量実験ができなくなるような拡張をしてはいけない
適度にCallback等を使って切り分ける。
PyTorchのイメージ
100の力を使ってコード作成
実験
10の力を使ってコード変更
実験
10の力を使ってコード変更
実験
PytorchLightningのイメージ
100の力を使ってPyTorchLightningについて学ぶ
100の力を使ってコード作成
実験
5の力を使ってコード変更
実験
5の力を使ってコード変更
実験
→ 20回くらい実験すれば効率がPyTorchLightningの勉強を上回る
この便利さから他の深層学習プロジェクトでもPyTorchLightningを使用することが増えている。
どうせ利用する日が来るのだから、今やっておいて損はないはず。
PyTorchLightningでは、モデルの訓練に必要なものを大きく4要素に分割し整理している。
これらはすべてクラス(つまりオブジェクト)として分割する。
1. 研究コード: Lightning Module
学習の流れを分割して整理したもの、最も大きくなるオブジェクト
モデルと各step(epochごとやバッチごと)の挙動をまとめたクラス
今までのモデルクラスに訓練の流れや訓練に必要なもの(optimizer/shedularオブジェクト)を全部のせしちゃったイメージ
クラス内での関数名がすでに指定されているのでこれを埋めていく形
モデルを定義するだけでなくバッチに対するlossの計算、optimizerまで定義するクラスになっている
作成したモデルをTrainerクラスに渡してfitメソッドで学習
2. データ: Pytorch Dataloader or LightningDataModule
Dataset周りを定義するクラス。これも関数名が事前に決まっている
Dataloaderを返却する
Datasetのダウンロードや交差検証への対応もここで行う
データの用意や前処理もここにまとめることで分かりやすくすることができる。
従来のDataloaderに前処理やダウンロード方法も乗せちゃったイメージ。
LightningDataModuleはDatasetを読みとる形で構成するとPyTorchでも使えて嬉しいかも
3. エンジニアリングコード: Trainer
fitするだけで実行できる、学習や推論ループを回す本体
使うときは引数を埋めるだけ。
4. 非必須の研究コード: Callback
訓練に直接関係のない追加のコード。Loggerやearlystopping等。
early stopping(モデルが改善しなさそうなら学習を止める)
Wandb logger(ログを取る)
ModelCheckpoint(モデルを保存する)
自作のCallbackも作成できる
訓練に直接関係のない可視化のためのログ集め+可視化
LightningModuleまたはTrainerに渡すことができる
モデル構造によって挙動が変わるならLightningModuleに渡す
どんな訓練でも使えそうならTrainerに渡す
私の思う,PyTorchLightningの利点を以下に挙げておく。
.to(device) や loss.backward()、optimizer.step(),with torch.no_grad()を書く必要がない
コードを使い回しやすい
Callbackを利用して学習に普遍的に必要なものを後から色々な実験に付与することができる
FP16への減少などが比較的簡単にできる
学習済のモデルオブジェクトを再度Python上に展開するのが楽
DDPを簡単に書ける
公式サイトの説明が潤沢
他の人が書いたPyTorchLightningのコードは読みやすく参考にしやすい(構造化されてるから)
逆に問題点は
PyTorchよりも実装に時間がかかってしまうことがわりと多い
細かい実装をしたいときにモンキーハックしなきゃいけない時がある