DatabricksでPyTorch Lightningを用いたディープラーニングを加速させる

[Accelerating Your Deep Learning with PyTorch Lightning on Databricks \- The Databricks Blog](https://www.databricks.com/blog/2022/09/07/accelerating-your-deep-learning-pytorch-lightning-databricks.html)の翻訳です。

:::note warn
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

PyTorch Lightningは、お手元のPyTorchコードをシンプルにし、ディぷラーニングのワークロードをブートスラップするための優れた手段となります。しかし、皆様のレイクハウスにあるすべてのデータを用いてタイムリーな結果を生み出すためにワークロードをスケールさせようとすると、PyTorch Lightningが持っている課題に直面することになります。本書では、Horovodを用いて、どのようにこれを達成するのか、コードをどのように効率的にスケールさせるのかを説明します。

# イントロダクション

企業は自分達の高度な機械学習アプリケーションを加速させるために、徐々にディープラーニングに移行しています。例えば、現在ではコンピュータービジョン技術が、[製造業における不良品検知](https://medium.com/@infopulseglobal_9037/intelligent-defect-inspection-powered-by-computer-vision-and-deep-learning-4c75fdf8673)を改善するために活用されています。[チャットボットを用いてビジネルプロセスを拡張](https://venturebeat.com/2021/05/20/despite-challenges-salesforce-says-chatbot-adoption-is-accelerating/)するために自然言語処理が用いられ、顧客に対する成果を改善するために[ニューラルネットワークベースのレコメンデーションシステム](https://venturebeat.com/2021/07/19/ai-powered-deep-neural-nets-increase-accuracy-for-credit-score-predictions/)が用いられています。

しかし、適切に最適化されたコードであっても、ディープラーニングモデルのトレーニングは遅いプロセスとなり、データサイエンスチームがクイックな実験の繰り返しを通じて、成果を提供する能力を制限してしまっています。このため、これをスケールアップさせるために、どのように計算能力を組み合わせるのがベストであるのかを知ることが重要となります。

本書では、最初にコードの再利用を最大化するためにお使いのコードベースをどのように構造化するのかを説明し、小規模な単一ノードインスタンスから、どのようにして完全なGPUクラスターにスケールさせるのかを説明します。また、完全なる実験トラッキングとモデルロギングを提供するために、すべてをMLflowとインテグレーションします。

# Part 1 - データのロードおよびPyTorch Lightningの導入

ターゲットとなるアーキテクチャからスタートしましょう。

## クラスターのセットアップ

ディープラーニングをスケールさせる際、高価なGPUリソースを効率的に活用できるように、実験を小規模なものからスタートし、徐々にスケールアップすることが重要です。コードの複雑性を削減するために、複数のノードにスケールさせる前に、単一ノードの複数のGPUで実行されるようにコードをスケールアップします。

Databricksでは、この有用なパターンをサポートするためのシングルノードクラスターをサポートしています。[Azure Single Node Clusters](https://docs.microsoft.com/en-us/azure/databricks/clusters/single-node)、[AWS Single Node Clusters](https://docs.databricks.com/clusters/single-node.html)、[GCP Single Node Clusters](https://docs.gcp.databricks.com/clusters/single-node.html)をご覧ください。インスタンスの選択に関しては、Nvidia T4 GPUがスタートするのに最適なコスト効率の高いインスタンスタイプを提供しています。AWSでは[G4インスタンス](https://aws.amazon.com/blogs/aws/now-available-ec2-instances-g4-with-nvidia-t4-tensor-core-gpus/)を利用することもできます。Azureでは[NCasT4\_v3インスタンス](https://docs.microsoft.com/en-us/azure/virtual-machines/nct4-v3-series)、GCPでは[A2インスタンス](https://cloud.google.com/compute/docs/accelerator-optimized-machines)を利用することができます。

ノートブックで説明している手順を進めるには、少なくとも64GBのRAMが非有用となります。モデリングのプロセスではメモリーを大量に必要とし、小規模なインスタンスでは以下のようなRAM不足のエラーになる場合があります。

```
Fatal error: The Python kernel is unresponsive.
```

ここで使用するコードは、Databricks機械学習ランタイム10.4LTSと11.1MLで構築され、テストされました。DBR 10.4 ML LTSではpytorch-lightning 1.6.5までがサポートされています。DBR 11.1 MLではpytorch-lightning 1.7.2でテストしました。我々はライブラリを[ワークスペースレベルライブラリ](https://docs.databricks.com/libraries/cluster-libraries.html)としてインストールしました。ドライバーノード上のアクティブなノートブックにのみライブラリをインストールする`%pip`と異なり、ワークスペースライブラリは後で行う分散トレーニングで必要となるすべてのノードにライブラリをインストールします。

DBR 10.4 LTS MLの設定
![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-1.png)

DBR 11.1 MLの設定
![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-2.png)

*図1. ライブラリの設定*

## ターゲットアーキテクチャ

![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-3.png)
*図2. キーとなるコンポーネント*

本書のゴールは、上のように構造化されたコードベースを構築することです。オープンソースのLinux Foundationプロジェクトである[Delta Lake](https://delta.io/)を用いてデータを格納します。Delta Lakeは内部で生のデータをParquetフォーマットで格納します。Petastormがデータロードの役割を担い、レイクハウスとディープラーニングモデルの間のインタフェースとなります。[MLflow](https://mlflow.org/)は、エクスペリメントトラッキングツールを提供し、モデルレジストリにモデルを保存できるようになります。

このセットアップによって、不要なデータ複製によるコストや、トレーニングするモデルの管理のコストを回避することができます。

# Part 2 - サンプルユースケースおよびライブラリの概要

## サンプルユースケース

このサンプルユースケースでは、[tensorflow flowers](https://www.tensorflow.org/datasets/catalog/tf_flowers)データセットを使用します。このデータセットは、花はどのクラスなのかを識別しようとする分類タイプの問題で使用されます。
![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-4.png)
*図3. Flowersデータセット*

## Petastormを用いてデータレイクをディープラーニングに活用する

歴史的に、[レイクハウス](https://www.databricks.com/research/lakehouse-a-new-generation-of-open-platforms-that-unify-data-warehousing-and-advanced-analytics)やデータウェアハウスのようなデータ管理システムは、機械学習フレームワークとインテグレーションするというよりは、並列に開発されてきました。このため、PyTorchデータローダーモジュールでは、簡単にParquetをサポートすることができません。また、これらはhiveメタストアのようなレイクハウスのメタデータ構造とインテグレーションされていません。

[Petastorm](https://petastorm.readthedocs.io/en/latest/)プロジェクトは、レイクハウスのテーブルとPyTorchの間のインタフェースを提供します。トレーニングノード間のデータのシャーディングをハンドリングし、キャッシングレイヤーを提供します。PetastormはDatabricks機械学習ランタイムにプレパッケージされています。

最初にデータセットやどのように操作するのかに慣れ親しんでみましょう。

```py:Python
peta_conv_df = make_spark_converter(preprocessed_df)
```

`spark_converter`オブジェクトを作成することで、以下のようにPyTorch Dataloaderに変換することができます。

```py:Python
with peta_conv_df.make_torch_dataloader(transform_spec=transform_func) as converted_dataset
```

これによって、通常pytorchコードで使うのと同じように使用できる`converted_dataset`の`DataLoader`を提供します。

[Exploring the flowers dataset](https://www.databricks.com/wp-content/uploads/notebooks/db-277-ptorch-dl/exploring-the-flowers-dataset.dbc)というノートブックを開いて試してみましょう。標準的なMLランタイムのクラスターで十分であり、GPUクラスターで実行する必要はありません。

## モデルをシンプルにし構造化する - PyTorch Lightningにようこそ

デフォルトでは、PyTorchコードは膨大なものとなる場合があります。モデルの定義、トレーニングループ、データローダーのセットアップが含まれます。デフォルトでは、このコードはすべて結び付けられており、迅速な実験においてはキーとなるデータセットやモデルを入れ替えることが困難となります。

PyTorch Lightningは、実験モデルやメインのトレーニングループのセットアップに必要な定型コードを大幅に削減することでシンプルにします。可読性とメンテナンス性を高めるための、PyTorchコードを構造化するためのアプローチとなります。

我々のプロジェクトにおいては、コードを以下のモジュールにブレークダウンします。

- PyTorchモデル
- データローダーとトランスフォーメーション
- メインのトレーニングループ

これによって、コードがよりポータブルになり、整理することができるようになります。これらのすべてのクラスとファンクションは、トレーニングのハイパーパラメーターが定義され、実際にコードを実行するメイン処理実行用のノートブックから`%run`経由で呼び出されることになります。

![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-5.png)
*図4. コードのレイアウト*

## モデル定義

このモジュールには、モデルクラスである[LightningModule](https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html)の中にモデルのアーキテクチャのコードが含まれています。ここにモデルのアーキテクチャが存在します。リファレンスとして、これが[timm](https://timm.fast.ai/)や[HuggingFace](https://huggingface.co/docs/transformers/index)のような人気のあるモデルフレームワークを活用するためにアップデートを必要とするモジュールとなります。また、このモジュールにはオプティマイザの定義も含まれます。この場合、SGDを使用しますが、他のタイプのオプティマイザをテストするためにパラメーター化を行うことができます。

## DataLoaderクラス

データローダーのコードとモデルのコードが結びつき合っているネイティブのPyTorchと異なり、PyTorch Lightningを用いることで、別の[LightningDataModule](https://pytorch-lightning.readthedocs.io/en/latest/data/datamodule.html?highlight=DataModule)クラスに分割することができます。これによって、データセットの管理を容易にし、データセットに対して異なる操作をクイックにテストできるようになります。

Petastormデータローダーを用いた`LightningDataModule`を構築する際には、生のSparkデータフレームではなく`spark_converter`オブジェクトにデータを投入します。Sparkデータフレームは内部のSparkクラスターで管理され最初から分散されておりますが、PyTorchデータローダーにおいては別の手段を用いて後ほど分散します。

## メインのトレーニングループ

これがメインのトレーニング関数となります。`Trainer`クラスにデータを投入する前に、モデルを定義する`LightningDataModule`と`LightningModule`を受け取ります。ここで、PyTorch Lightningのトレーナーのインスタンスを作成し、必要なすべてのコールバックを定義します。

後ほどこのトレーニングプロセスをスケールアップさせるので、すべての処理ノードでMLflowのロギングのような幾つかのプロセスを実行する必要はありません。このため、ここでは、最初のGPUでのみ実行するように処理を制限します。

```py:Python
if device_id == 0:
    
    # we only need this on node 0
    mlflow.pytorch.autolog()
```

トレーニングの過程でモデルのチェックポイントを作成することは進捗を保存する観点で重要ですが、PyTorch Lightningでは[デフォルトで](https://pytorch-lightning.readthedocs.io/en/stable/common/checkpointing.html#automatic-saving)チェックポイントの作成を行うのでコードに追加する必要はありません。

ノートブック[Building the PyTorch Lightning Modules](https://www.databricks.com/wp-content/uploads/notebooks/db-277-ptorch-dl/building-the-pytorch-lightning-modules.dbc)を試してください。

# Part 3 - トレーニングジョブをスケールさせる

単体のGPUによるトレーニングはCPUによるトレーニングよりもはるかに高速ですが、多くの場合十分ではありません。適切なプロダクションのモデルは大規模なものとなり得ますし、これらを適切にトレーニングするために必要となるデータセットも大規模なものとなり得ます。このため、複数のGPUでトレーニングを行えるようにスケールアップする方法を探す必要があります。

ディープラーニングモデルを分散させる主なアプローチは、それぞれのGPUにモデルのコピーを送信し、それぞれの異なるデータのシャードを与えるというデータ並列性を通じて行うというものです。これによって、[こちらの記事](https://qiita.com/taka_yayoi/items/322a09f88f6eaf421cb5)で議論されているように、トレーニング時間を改善するために、バッチサイズを増やし、高い学習率を活用できるようになります。

GPUに対するトレーニングジョブの分散を支援するために、Horovodを活用することができます。[Horovod](https://horovod.ai/)は別のLinux Foundationプロジェクトであり、複数ノードに対してpytorchのプロセスをマニュアルで起動することに対する代替策を提供します。Databricks MLランタイムにはデフォルトで[HorovodRunner](https://docs.databricks.com/applications/machine-learning/train-model/distributed-training/horovod-runner.html)クラスが含まれており、シングルノード、マルチノードのトレーニング両方でスケールさせることができるようになります。

Horovodを活用するためには、新たな「スーパー」トレーニングループを作成する必要があります。

```py:Python
def train_hvd():
  hvd.init()
  
  # MLflow setup for the worker processes
  mlflow.set_tracking_uri("databricks")
  os.environ['DATABRICKS_HOST'] = db_host
  os.environ['DATABRICKS_TOKEN'] = db_token

  
  hvd_model = LitClassificationModel(class_count=5, learning_rate=1e-5*hvd.size(), device_id=hvd.rank(), device_count=hvd.size())
  hvd_datamodule = FlowersDataModule(train_converter, val_converter, device_id=hvd.rank(), device_count=hvd.size())
  
  # `gpus` parameter here should be 1 because the parallelism is controlled by Horovod
  return train(hvd_model, hvd_datamodule, gpus=1, strategy="horovod", device_id=hvd.rank(), device_count=hvd.size())
```

この関数は、horovodの`hvd.init()`を起動させ、我々のDataModuleと`train`関数が、適切なノード番号である`hvd.rank()`、デバイスの総数である`hvd.size()`を用いて起動されることを保証します。[horovod の記事](https://horovod.readthedocs.io/en/stable/pytorch.html?highlight=scale%20learning%20rate#horovod-with-pytorch)で議論されているように、GPUの数を用いて学習率をスケールアップします。

```py:Python
hvd_model = LitClassificationModel(class_count=5, learning_rate=1e-5*hvd.size(), device_id=hvd.rank(), device_count=hvd.size())
```

そして、Horovodが並列性をハンドリングするので、GPUのカウントが1に設定された通常のトレーニングループを返却します。

[Main Execution notebook](https://www.databricks.com/wp-content/uploads/notebooks/db-277-ptorch-dl/main-execution-notebook.dbc)を試していく中で、単一のGPUから複数のGPUに移行する方法を見ていきます。

## Step 1 - 1ノードでスケールさせる

![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-6.png)
*図5. シングルノードにおけるスケーリング*

1ノードをスケールさせることが一番簡単な方法となります。また、複数ノードのトレーニングでは必要となるネットワークトラフィックを避けることができるので、非常の高速です。SparkネイティブのMLライブラリと異なり、多くのディープラーニングのトレーニングプロセスは、ノード障害から自動で復旧しません。しかし、PyTorch Lightningはトレーニングのエポックを復旧できるように自動でチェックポイントを作成します。

我々のコードでは、train関数の中でパラメーター`default_dir`をdbfsに指定しています。ここにPyTorch Lightningがチェックポイントを保存します。`ckpt_restore`のパスをチェックポイントを指すようにすると、train関数はそこにあるチェックポイントからトレーニングを再開します。

```py:Python
def train(model, dataloader, gpus:int=0, 
          strategy:str=None, device_id:int=0, 
          device_count:int=1, logging_level=logging.INFO,
          default_dir:str='/dbfs/tmp/trainer_logs',
          ckpt_restore:str=None,
          mlflow_experiment_id:str=None):
```

我々のtrain関数を1ノード上の複数のGPUにスケールさせるには、`HorovodRunner`を使用します。

```py:Python
from sparkdl import HorovodRunner

hr = HorovodRunner(np=-4, driver_log_verbosity='all')
hvd_model = hr.run(train_hvd)
```

`np`を負の値に設定するとシングルノードで動作します。この例では、ドライバーノードの4CPUで実行されます。`np`を正の値にすると、ワーカーノードにまたがって実行されます。

## Step 2 - ノードにまたがってスケールさせる

![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-7.png)
*図5. 複数ノードでのスケーリング*

我々は既にトレーニング関数をhorovodラッパーでラップしており、シングルノードにおける複数GPUでHorovodRunnerを活用しました。最後のステップは、マルチノード、マルチGPU環境に移行することとなります。ここまででシングルノードクラスターを使っていたのであれば、ここでマルチノードクラスターに切り替えることになります。以降のコードでは、以下の設定のクラスターを使用します。
![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-8.png)
*図6. マルチノードクラスターのセットアップ*

Databricksで分散トレーニングを実行する際、現時点ではオートスケーリングがサポートされていないので、事前に固定数のワーカーを設定しておきます。

```py:Python
hr = HorovodRunner(np=8, driver_log_verbosity='all')
hvd_model = hr.run(train_hvd)
```

分散ディープラーニングのジョブをスケールアップする際に生じる一般的な問題は、すべてのGPUがバッチのスプリットを取得できるようにPetastormのテーブルが適切にパーティショニングされていないということです。複数のGPUを持っているので、少なくとも多くのデータパーティションが作成されるようにする必要があります。

我々のコードでは、`num_devices`変数と`prepare_data`関数を用いてGPUの数を設定することでこの問題に対応しています。

```py:Python
flowers_df, train_converter, val_converter = prepare_data(data_dir=Data_Directory, num_devices=NUM_DEVICES)

datamodule = FlowersDataModule(train_converter=train_converter, 
                               val_converter=val_converter)
```

これは、シンプルに標準的なSparkの再パーティションコマンドを呼び出します。我々はパーティションの数をGPUの数である`num_devices`の倍数になるように設定しており、データセットがトレーニングプロセスに割り当てたすべてのGPUに対して十分なパーティションを持つようにしています。不十分なパーティションはアイドル状態のGPUの一般的な原因となります。

```py:Python
flowers_dataset = flowers_dataset.repartition(num_devices*2)
```

# 分析

ディープニューラルネットワークをトレーニングする際、ネットワークを下学習させないことが重要となります。これに対処する標準的な方法は、Early Stoppingを用いるというものです。このプロセスは、それぞれのエポックにおいてモニタリングするように設定したメトリックに依然として改善が認められるのかをチェックします。この場合は`val_loss`となります。

我々の実験においては、`min_delta`を`0.01`に設定しているので、それぞれのエポックで少なくとも`val_loss`に`0.01`の改善が認められることを期待します。`patience`を`10`に設定しているので、トレーニングを停止する前に、改善が認められないエポックが10発生するまでトレーニングを継続します。我々は最後のパフォーマンスの改善まで持たせるまでこの設定を行なっています。実験を短時間で終わらせられるように、`stopping_threshold`を`0.55`に設定しており、`val_loss`がこのレベルを下回った場合には、すぐにトレーニングプロセスを停止します。

これらのパラメーターを念頭において、以下でスケーリングの実験結果を見ていきましょう。

![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-9.png)
![](https://cms.databricks.com/sites/default/files/inline-images/db-277-blog-img-10.png)

Running Time vs Cluster Setupのグラフから分かるように、システムのリソースを増加させることで、トレーニング時間をほぼ半分にしています。異なるGPU間でのトレーニングプロセスの調整によるオーバーヘッドのため、線形にはスケーリングしていません。ディープラーニングをスケールさせる際、結果が減少しているかどうかを見ることは一般的であり、GPUを追加する前にトレーニングループが効率的になっていることを確認することが重要です。

これは全体像ではありませんが、以前のブログ記事[ディープラーニングを容易にスケールさせる\(させない\)6つのステップ](https://qiita.com/taka_yayoi/items/322a09f88f6eaf421cb5)でアドバイスしているベストプラクティスに基づいて、さまざまなトレーニングの実行によって達成される最終的な`validation loss`をチェックすることが重要なので`EarlyStopping`を活用しました。この場合、`stopping_threshold`を`0.55`に設定しました。面白いことに、単一のGPU環境では、複数GPUの環境よりも悪い`validation loss`で停止しています。単一のGPUでのトレーニングでは、`val_loss`に改善が認められなくなるまで処理が実行されました。

# 使い始める

ここまでで、DatabricksでどのようにPyTorch Lightningを活用できるのか、複数ノードにスケールさせるためにHorovodRunnerでどのようにラッピングするのかを説明し、`EarlyStopping`をどのように活用するのかに関してのガイドを説明しました。次はあなたがトライしてみる番です。

**ノートブック**

- [Exploring the flowers dataset](https://www.databricks.com/wp-content/uploads/notebooks/db-277-ptorch-dl/building-the-pytorch-lightning-modules.dbc)
- [Building the PyTorch Lightning Modules](https://www.databricks.com/wp-content/uploads/notebooks/db-277-ptorch-dl/exploring-the-flowers-dataset.dbc)
- [Main Execution Notebook](https://www.databricks.com/wp-content/uploads/notebooks/db-277-ptorch-dl/main-execution-notebook.dbc)

**参考文献**

- [HorovodRunner](https://docs.databricks.com/applications/machine-learning/train-model/distributed-training/horovod-runner.html)
- [Petastormを用いたデータロード](https://qiita.com/taka_yayoi/items/33086561e44c4eee027c)
- [Databricksにおけるディープラーニングのベストプラクティス](https://qiita.com/taka_yayoi/items/991917125c6f38357195)
- [ディープラーニングを容易にスケールさせる\(させない\)6つのステップ](https://qiita.com/taka_yayoi/items/322a09f88f6eaf421cb5)
- [Leveling the Playing Field: HorovodRunner for Distributed Deep Learning Training](https://www.databricks.com/blog/2021/01/14/leveling-the-playing-field-horovodrunner-for-distributed-deep-learning-training.html)

### Databricks 無料トライアル

[Databricks 無料トライアル](https://databricks.com/jp/try-databricks)