PyTorch を使用して高速な開発を実現する方法

※この投稿は米国時間 2021 年 2 月 4 日に、Google Cloud blog に投稿されたものの抄訳です。

長年、私は機械学習モデルを構築するために多くのフレームワークを使用してきましたが、PyTorch を試したのはつい最近です。導入チュートリアルの Deep Learning with PyTorch: A 60 Minute Blitz（PyTorch を使用したディープ ラーニング: 60 分）を見た結果、使い方を理解し始めました。ノートブックや事前構成された VM イメージなどの PyTorch サポートが Google Cloud に組み込まれているため、簡単に始めることができました。

しかし、1 つ難点がありました。その素晴らしい柔軟性により、さまざまな作業に対して採りうるアプローチが非常に多くなってしまうことです。どのようにトレーニングとテストのデータを読み込むべきか？損失の計算とログへの記録を行いながら、どのようにモデルをトレーニングするべきか？すべての作業を正常にこなすことができましたが、アプローチを改善できるのではないかと常に考えていました。作業の方法を整理し、問題の解決に集中できるようにする、より高いレベルの抽象化ができないか模索していたのです。

嬉しいことに PyTorch Lightning を見つけました。Lightning は、コードを整理し、よく使う機能のユーティリティを提供するのに役立つ、軽量の PyTorch ラッパーです。Lightning を使用すると、CPU、GPU、TPU で動作する標準の PyTorch モデルを簡単に作成できます。その仕組みと使用開始方法を詳しく見てみましょう。

本ブログ投稿では、PyTorch Lightning をご紹介するため、私のノートブック Training and Prediction with PyTorch Lightning（PyTorch Lightning を使用したトレーニングと予測）にあるサンプルコードを見ていきます。使用する UCI Machine Learning Repository のデータセットは、水中ソナーから金属製シリンダーや岩石に発射して返された測定値で構成されています。このモデルは、返された信号に基づいて、見つかったアイテムを分類します。音響データには、医用画像や地震探査などのさまざまな用途があり、機械学習はこのデータのパターンを検出するのに役立ちます。

PyTorch Lightning を使用してノートブック コードを整理する

Lightning をインストールした後、まず SonarDataset を標準の PyTorch Dataset から継承して作成しました。このクラスは、データの読み込み、反復処理、変換のためのロジックをカプセル化します。たとえば、「R」で岩石を、「M」で鉱物を表す元データを 0 と 1 にマッピングします。これにより、「これは鉱物か？」という二項分類の問題としてデータを解釈できるようになります。そのクラスのコード スニペットは次のとおりです。

次に、Lightning の LightningDataModule を継承して SonarDataModule を作成しました。このクラスは、データをトレーニング、テスト、検証の 3 セットに分割し、各セットを PyTorch DataLoader に読み込むための標準的な方法を提供します。SonarDataModule の setup() メソッドのコード スニペットは次のとおりです。

最後に、SonarModel を LightningModule から継承して作成しました。このクラスにはモデルのほかに、予測のための forward()、トレーニング損失計算のための training_step()、精度計算のための test_step() など、プロセスの各ステップのメソッドが含まれています。

モデルを使用したトレーニングと予測

Lightning の Trainer クラスを使用すれば、トレーニングが簡単になります。このクラスは、PyTorch DataLoaders とのやり取り、必要に応じたグラデーションの有効化または無効化、コールバック関数の呼び出し、適切なデバイスへのデータと計算のディスパッチなどの詳細を管理します。

チュートリアル ノートブックのメソッドをいくつか見てみましょう。まず、新しいトレーナーをインスタンス化し、使用する GPU の数やトレーニング期間などのオプションを指定します。fit() でモデルをトレーニングし、test() でテストデータの最終評価を実施できます。ハイパーパラメータを調整するための tune() メソッドも用意されています。

トレーニング プロセスの後、標準の PyTorch 関数を使用して、モデルの保存やモデルを使った予測ができます。たとえば、次のようになります。

Lightning の使用を開始する

Google Cloud の PyTorch のサポートにより、Lightning を使用してモデルを簡単に構築できます。それでは手順を見ていきましょう。まず、Cloud AI Platform Notebooks を使用してノートブック インスタンスを作成します。PyTorch DLVM イメージがプリロードされる PyTorch 対応インスタンスを選択します。必要に応じて GPU サポートを有効にすることもできます。ノートブック インスタンスがプロビジョニングされたら、[JupyterLab を開く] を選択して開始します。

PyTorch の依存関係はすでに構成されているため、ノートブックに「!pip install pytorch-lightning」の 1 行を含めるだけで Lightning の使用を開始できます。

このチュートリアルのサンプルにアクセスする場合は、新しいターミナルを開いて（[ファイル] > [新規] > [ターミナル]）、git clone https://github.com/GoogleCloudPlatform/ai-platform-samples を実行します。サンプルは、[ai-platform samples] > [notebooks] > [samples] > [pytorch] > [lightning] にあります。

Lightning を使用すれば、PyTorch の使用がこれまでになく簡単になります。また、フレームワークに組み込まれたベスト プラクティスと便利なユーティリティにより、ML の問題の解決に集中できるようになります。Lightning は標準の PyTorch コードを生成するため、モデルの開発、トレーニング、提供に Google Cloud の PyTorch サポートを利用できます。

-デベロッパー アドボカシー担当マネージャー Karl Weinmeister