クラウドエンジニアブログ

PyTorch-Lightning を Azure Machine Learning Services で動かす

白石 卓也

はじめに

こんにちは。データサイエンスチームの白石です。
深層学習モデルを構築するときに、当社では主に Azure Machine Learning Services(以下、AMLS)を使っています。
この記事では、PyTorch で実装したモデルを AMLS 上で使う方法について説明します。特に、PyTorch-Lightning(以下、Lightning)という高機能な PyTorch のラッパーライブラリを使うことで、効果的な訓練を実現する方法を紹介します。
まず、Lightning と AMLS がそれぞれどのようなものなのかを簡単に紹介します。次にこれらを使って具体的にモデルを訓練する手順を紹介します。最後に、Lightning や AMLS を使う上で生じやすい様々なトラブルについて、解決のための手引を紹介します。

「機械学習導入支援サービス」に関する資料請求・お問い合わせはこちら


PyTorch-Lightning とは

Lightning は、PyTorch のモデルをラッピングし、効率的な Experiment を実現するためのラッパーライブラリです。このようなラッパーライブラリは、しばしば TensorFlow に対する Keras の立ち位置を表す "Keras-Like" なライブラリと呼ばれます。

PyTorch における Keras-like と呼べそうなライブラリは、以下のように複数乱立している状況です。


これらのライブラリの中で私が Lightning を気に入っている理由は、学習コストの低さと柔軟性にあります。公式リポジトリの冒頭には、以下のように Lightning の立ち位置が明言されています。以下、意訳です。

Lightning は、研究のためのコードをエンジニアリングから切り離すために PyTorch コードを整理する方法です。フレームワークというよりも、PyTorch のスタイルガイドのようなものです。
Lightning では、コードを3つの異なるカテゴリに整理します。

  1. 研究コード(LightningModule に実装する)。
  2. エンジニアリングコード (Trainer が担当する)。
  3. ロギングなどの研究にとって本質的ではないコード (Callback 類が担当する)。

Lightning を使うためには、PyTorch で実装したネットワークの定義や損失関数、訓練・評価処理などを LightningModule というひとつの Python クラスの各メソッドとして実装する必要があります。これらのコードは3つのカテゴリのうちの1番目の研究コードに該当します。この作業は、いままで何十行も必要だったコードがわずか数行に収めらるというニュアンスの Keras-like という単語から受ける印象とはだいぶ異なっています。結局 PyTorch の基本的な処理は記述する必要があるので、実装にかかる時間は PyTorch を生で実装するのとほとんど変わりません。これが、「(Lightning は)PyTorch のスタイルガイドのようなものです」ということの意味です。

Lightning が素晴らしいのは、研究コードのスタイルガイドを守って LightningModule を実装しさえすれば、様々な機能を低コストで導入できる、という点にあります。
それを担っているのが2、3番目のカテゴリのエンジニアリングコードと各種 Callback です。LightningModule に適切に各処理を実装すれば、様々な機能を内包している Trainer と Callback を自由に組み合わせてモデルを訓練できるのです。たとえば、Trainer の基本的な機能の中に、簡単に PyTorch モデルを分散環境で訓練可能にできる、というものがあります。同一の LightningModule を用いて、使用する分散環境や GPU 数を、Trainer に与える引数で指定して実行することができます。PyTorch の標準である NCCL 以外にも、Horovod を指定したり、(Azure ではなく GoogleCloudPlatform や Colaboratory では)TPU も指定することもできます。もちろん、単一ノードの GPU や CPU での訓練も可能です。生の PyTorch 実装で分散環境での訓練を行ったことがある方にとっては、DistributedDataParallel まわりの処理を自分で書く必要がなくなる、といえばわかりやすいでしょうか。

また、モデルの訓練時に使用できる様々なテクニックが Trainer には実装されています。例えば、限られた GPU 数で大規模なバッチサイズの訓練をしたいというときに、擬似的に巨大なバッチでの訓練を行うという accumlate gradient や、モデルのパラメータの更新に使う勾配の大きさをある範囲内に収める gradient clipping、また、バッチサイズを GPU のメモリサイズに自動的に適応させる Auto scaling of batch size などがオプションとして指定できます。

さらに、各種の Callback が用意されています。ロギングの Callback として、デフォルトでは TensorBoard が使用でき、LightningModule に定義した訓練処理や評価処理で計算した損失などが自動で TensorBoard のログとして記録されます。TensorBoard の他には、CometML や MLflow といったその他のサービスやライブラリへの対応が可能です。その他の Callback としては、たとえば、早期終了(Eary Stopping)などの機能が用意されています。

以上のように、Lightning は生の PyTorch の処理を LightningModule に配置するだけで高機能な Trainer と Callback を利用できるようになるライブラリです。研究コードは通常の PyTorch で書けるため、学習コストの低さや柔軟性をもたらしています。先程挙げたいくつかの Keras-like なライブラリの中には、Lightning に比べるとより抽象度の高いラッピングがなされているために、そのライブラリ固有の学習コストが高くなってしまったり、PyTorch そのものが持っている柔軟性を疎外してしまうために複雑なモデルの訓練が困難になったりすることもあります。もちろん、Lightning の抽象化の粒度も完璧ではありませんし、まだまだ発展途上の部分もありますが、筆者は一番気に入っています。

以上、PyTorch Lightning の紹介でした。興味のある方は、PyTorch から Lightning に移行するためのガイドが、メインコントリビュータである William Falcon 氏によって書かれているので、一読をおすすめします。


Azure Machine Learning Services とは

Azure Machine Learning Services(AMLS)についても、簡単にですが紹介しておきます。
AMLS はデータセット管理、訓練、デプロイを一気通貫で行えるように整備されたマネージドサービスです。ドラッグ&ドロップによって簡単に機械学習モデルを構築できる Visual Designer や、自動的に最適な機械学習モデルを作ってくれる AutoML、Jupyter Notebook を手軽に構築できる Notebook VM、手軽に訓練済みのモデルを公開できる WebAPI 機能なども用意されています。

今回は、独自の深層学習モデルを訓練するために最も汎用的な機能である Experiment を使用します。Experiment では、モデルの構築のために記述したスクリプトや依存ファイルを、SDK を通じて AMLS にジョブとして提出することで、Azure 上の計算資源を使ってモデルを訓練することができます。ジョブを実行する環境である ComputingTarget は自由に選ぶことができ、ちょっとした訓練であれば1つ GPU を積んだ1台の仮想マシンを指定することができますし、より大規模な環境として8GPU のノードを10台使うといった指定もできます。もちろん、リッチな ComputingTarget はそれなりにお値段がかかるので注意が必要です。以下のような設定の ComputingTarget を作ることでお値段を抑えて活用することができます。

  • 最小ノード数を0にすることで、ジョブが実行されていないタイミングではノード数が0になるようにできます。ノード数が0の間はコンピューティングの費用はかかりません。その代わり、ノード数が0のときに新しいジョブを送信すると、実行が開始されるまでに少々時間がかかります。

  • 低優先度にすることで、急ぎではない訓練を Azure 全体のリソースの隙間時間に実行することができます。低優先度の ComputingTarget は、価格が約30%になります。


Lightning on AMLS

さて、ここからが本題です。AMLS 上で Lightning の訓練を行ってみましょう。手順としては以下のとおりです。

  1. 訓練に使用するスクリプトを作成する
  2. AMLS 環境を準備する
    1. Workspace を作成する
    2. ComputingTarget を作成する
    3. Environment を作成する
  3. Estimator を提出して訓練を実行する
  4. ログを TensorBoard で可視化する

以降は、私が書いたソースコードや jupyter notebook も併せてご確認ください。


訓練スクリプトの作成

まずは訓練に使用するスクリプトを作成しましょう。とはいえ、あまり凝ったものを作るのではなく、Lightning の公式リポジトリのサンプルとして使われている、MNIST を対象にした GAN の訓練を使いましょう。

AMLS への対応のために、Trainer に対していくつかの変更を行います。
今回は、Lightning の0.7.4以降でサポートされている Horovod を分散訓練環境として使用することにしましょう。Horovod を使用するときの注意すべき点として、`gpus` 引数があります。通常、この引数には全ノードにまたがって存在する GPU の数を指定しますが、Horovod のときのみ「1」を指定する必要があります。これによって、Horovod のプロセス数分の GPU を自動的に使用できるようになっています。
次にロギング用の Callback として、カスタマイズした TensorBoardLogger を指定します。Lightning ではデフォルトで TensorBoardLogger が使用されますが、ログの出力先がデフォルトでは lightning_logs というディレクトリになってしまいます。AMLS がサポートしている TensorBoard のログ出力先は、logs というディレクトリで固定なので、そのように指定した TensorBoardLogger を指定します。
その他の変更点として、あまり長時間訓練を行っても仕方ないので、`max_epoch` 引数を100に指定しています。

以上の変更を反映すると、main 関数の中は以下のように修正されます。


def main(args: Namespace) -> None:
    # ------------------------
    # 1 INIT LIGHTNING MODEL
    # ------------------------
    model = GAN(**vars(args))

    # ------------------------
    # 2 INIT TRAINER
    # ------------------------
    # If use distubuted training  PyTorch recommends to use DistributedDataParallel.
    # See: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel
    tb_logger = TensorBoardLogger(save_dir="logs", name="mnist-gan")
    trainer = Trainer(
        distributed_backend="horovod",
        gpus=1,
        logger=tb_logger,
        max_epochs=10,
        prepare_data_per_node=True,
    )

    # ------------------------
    # 3 START TRAINING
    # ------------------------
    trainer.fit(model)


AMLS の環境準備

ここからは、ローカル環境や Notebook VM、Azure Notebook などで Jupyter Notebook を起動し、以下のドキュメントを参考に AMLS 環境を整えてきます。

https://docs.microsoft.com/ja-jp/azure/machine-learning/tutorial-1st-experiment-sdk-setup
https://docs.microsoft.com/ja-jp/azure/machine-learning/tutorial-train-models-with-aml


  • SDK のインストール
    まずは、azml-sdk がインストールされていない場合は、pip でインストールしておきましょう。Tensorboard も使うので、追加モジュールとして tensorboard を指定しています。
    
    pip install -u azml-sdk[tensorboard]
    

  • Workspace の作成
    次に、Azure Portal 上から Workspace を作成します。この Workspace の名前とサブスクリプション ID、リソースグループを指定した config.json という以下のようなファイルを作成しておきましょう。
    
    (json)
    {
        "subscription_id": "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX",
        "resource_group": "XXXXXXXXXX",
        "workspace_name": "XXXXXXXXXXXXXXXXXXX"
    }
    

    この json ファイルを読み込んで azml-sdk を通じた認証を通すことで、Workspace を Python から操作することができるようになります。
    
    ws = Workspace.from_config()
    

  • Experiment と ComputingTarget の作成
    Workspace を操作できるようになったら、Experiment と ComputingTarget を作成します。
    まずは Experiment を作成します。Experiment は、これから何回も実行されるジョブをまとめる概念です。以下のようなコマンドで作成されます。
    
    exp = Experiment(workspace=ws, name=EXPERIMENT_NAME)
    

    次に、Workspace に紐づく ComputingTarget を作成します。以下のコードは、Workspace に同一の名称の ComputingTarget が存在しない場合は新たに ComputingTarget を作成します。今回は、`min_nodes=0` とし、ジョブが走っていないときには費用がかからないようにします。また、`vm_priority="lowpriority"` にすることで低優先度での実行になる代わりにお安くなるプランをしていしています。
    
    if COMPUTE_NAME in ws.compute_targets:  # ComputingTargetがすでに存在していたらそれを読み込む
        compute_target = ws.compute_targets[COMPUTE_NAME]
        if compute_target and type(compute_target) is AmlCompute:
            print('found compute target. just use it. ' + COMPUTE_NAME)
    else:  # ComputingTargetが存在していなければ新規作成する
        print('creating a new compute target...')
        provisioning_config = AmlCompute.provisioning_configuration(
            vm_size=VM_SIZE,
            vm_priority="lowpriority",
            min_nodes=MIN_NODE,  # 0
            max_nodes=MAX_NODE
        )
    
        # create the cluster
        compute_target = ComputeTarget.create(
            ws, COMPUTE_NAME, provisioning_config)
    
        # can poll for a minimum number of nodes and for a specific timeout.
        # if no min node count is provided it will use the scale settings for the cluster
        compute_target.wait_for_completion(
            show_output=True, min_node_count=None, timeout_in_minutes=20)
    
        # For a more detailed view of current AmlCompute status, use get_status()
        print(compute_target.get_status().serialize())
    


Estimator の作成と提出

最後に、Estimator オブジェクトを作成します。Estimator オブジェクトは、どのような環境で、どのようなスクリプトを実行するのかを定義するオブジェクトです。今回は PyTorch に特化した Estimator クラスを使用し、以下のような項目を指定することで作成します。

  • 実行環境情報
    • ComputingTarget
    • ノード数(ComputingTarget で指定された max_nodes を上限として、いくつノードを使用するかを指定)
    • GPU を使用するか否か
    • Docker コンテナの共有メモリの容量
    • 分散学習の設定(下の例では、1ノードあたり2GPU を使用するための Horovod 用の設定)
    • conda や pip でインストールするライブラリ
    • PyTorch のバージョン

  • クリプトの情報
    • 実行対象のスクリプトが含まれるディレクトリ(conda や pip でインストールできないパッケージや独自パッケージ、軽いデータなどもここに含めます)
    • 実行対象のスクリプト名
    • スクリプトのコマンドライン引数

ソースコードは以下のようになります。


mpi_config = MpiConfiguration()
mpi_config.process_count_per_node = 2

est = PyTorch(
    source_directory=SCRIPT_DIR,
    entry_script=ENTRY_SCRIPT,
    script_params=SCRIPT_PARAMS,
    compute_target=compute_target,
    use_gpu=True,
    conda_packages=["pip"],
    pip_packages=["pytorch-lightning"],
    node_count=NUM_NODE,
    distributed_training=mpi_config,
    framework_version="1.5",
    shm_size="8G",
)

一つ注意していただきたいのは、`shm_size` 引数です。これは、Docker コンテナが使用できる共有メモリの大きさを表します。PyTorch ではプロセス間のデータのやり取りを共有メモリ上で行うため、ある程度大きな値を指定する必要があります。デフォルトでは512MB とやや貧弱なので、数 GB を指定しましょう。

以下のコマンドで Experiment に対して Estimator オブジェクトを提出できます。返り値として Run オブジェクトが返ってきます。

run = exp.submit(config=est)

このように、Estimator オブジェクトを Experiment に提出することでジョブの登録が完了し、キューに保管されます。他のジョブが走っていないタイミングを見計らって、AMLS が勝手にジョブを開始してくれます。Run オブジェクトはその実行ステータスや実行結果へのリンクなどを管理しています。

また、AMLS のポータルからは、環境作成を含む実行時のログやジョブのステータスを確認できます。訓練が完了したあとで各ジョブの「出力とログ」を確認すると、TensorBoard 用のログやその中にある Lightning のチェックポイントなどが確認できます。


Estimatorの作成と提出


TensorBoard による可視化

Run オブジェクトは、ログの保管場所なども管理しているので、この情報を元に、TensorBoard でログを確認することができます。


from azureml.tensorboard import Tensorboard
tb = Tensorboard([run], port=8008)
tb.start()

TensorBoard が起動し、8008番ポートからアクセスできるようになります。最初は何も表示されていないかもしれませんが、ジョブが実行されてから徐々にログが見えるようになるので、少し待ちましょう。training_step や validation_epoch_end などで出力している各種損失、epoch_end で設定している画像出力を確認することができます。以下の画像では version が2つできていますが、これは2つのノードそれぞれでログを出力しているために生じているものです。


TensorBoardによる可視化


TensorBoard は、以下のように停止することができます。


tb.stop()

訓練手順についての説明は以上です。


トラブルシューティング

Lightning を AMSL で使用するための方法を紹介してきました。とても簡単に分散環境でのモデルの構築ができるのですが、いくつか落とし穴があります。ここではそれらの解決するための手引を共有します。


AMLS

まずは AMLS についてです。AMLS はまず環境構築に落とし穴があります。一度やり方を確立するまでは、Docker イメージの構築や依存ライブラリのインストールなどに苦戦するかもしれません。

また、スクリプトレベルで解決しなければならない問題として、分散環境やマルチプロセス固有の問題が挙げられます。例えば、複数のプロセスから同一のファイルを触ってしまう可能性がある実装になっている場合に、この問題が生じやすいです。

いずれの場合も、まずは AMLS のログを確認する、ということを心がけましょう。ログは AMLS のポータルから確認することができます。各ジョブの「出力とログ」タブの azureml-logs ディレクトリの下に以下のような形で表示されます。


AMLS



基本的には、環境構築に関するログは55から始まるファイル、実際にスクリプトが実行されたあとの各プロセスのログは70から始まるファイルに吐き出されます。この区分は問題の切り分けに使えますが、実行時エラーの原因が環境構築にある、というタイプの問題に注意しましょう。例えば、Estimator オブジェクトを作成するときの `shm_size` が小さすぎると、環境構築はうまく行っても実行時にエラーが生じることがあります。ログに表示されている内容をよく読んで、内部でどのような問題が生じているのかを把握した上で各種設定を最適化する必要があります。


Lightning

続けて、Lightning で生じる問題の解決方法を確認しましょう。
まずは、Lightning の公式ドキュメントにあるデバッグのための機能をよく読みましょう。これらの機能は、訓練プロセスの問題を明らかにすることがあります。
特に、分散環境での訓練を行う前に、そのそも単一のシングル GPU や CPU 環境下で期待どおりの訓練が行えるのかを確かめましょう。`fast_dev_run` や `overfit_batches` といった引数は、限られたデータで期待通りの訓練が進むかを事前に検証できる機能です。

次に、LightningModule についてのドキュメントをよく読み、各メソッドの実行サイクルや分散環境のどのランクで実行されるのかを把握しましょう。これを理解せずに LightningModule を記述すると、同一のファイルを複数のプロセスで同時に触ってしまい、訓練に使用するデータが壊れてしまうといった問題が生じることがあります。

ドキュメントを読んでもよくわからない、ということであれば、公式リポジトリの issue を検索してみましょう。Lightning は日々改善されており、最新版を再インストールするだけで問題が解決する可能性もあります。もちろん、デグレードなどが生じることもあるので、安定した動作が確認できたら、AMLS の Estimator オブジェクト作成に使用する引数でバージョンを指定することも必要です。


まとめ

以上、AMLS 上で Lightning を使ってみるための手順について説明しました。
最後に少し大変そうな話をしましたが、大前提として、Lightning と AMLS の組み合わせは PyTorch モデル構築の非常に強力な手段となります。

深層学習のモデル構築は気にかけなければならないことが盛り沢山ですが、少しでも考え事を低減させるために、Lightning のように便利なラッパーライブラリや AMLS のようなマネージドサービスを活用することをおすすめします。



関連ページ

機械学習導入支援サービス

【総合】お問い合わせ

製品・サービスに関するお問い合わせはお気軽にご相談ください

ピックアップ

セミナー情報
クラウドエンジニアブログ
clouXion
メールマガジン登録