Graphcore IPUのためのPyTorch Lightning初回リリースのご紹介

PyTorch LightningがGraphcore IPUをサポートしました。PyTorch Lightningチームは、ここ数ヶ月、IPUとの連携の構築に精力的に取り組んできましたが、このたび、リリース1.4をコミュニティに提供することになりました。IPUを開発者にとって使いやすいものにするという当社の使命を果たすために、Graphcoreチームとの緊密な協力関係を築いてくれたことにとても感謝しています。

Graphcoreは、柔軟で高速なAI計算ソリューションを求めるAIリサーチコミュニティの要望に応えるというPyTorch Lightningのミッションに広く協力しています。

PyTorch Lightningは、データサイエンティストやディープラーニングの専門家を、大変なエンジニアリング業務（データ配信、ループ管理、ロギング処理など）から解放し、モデリングやデータ理解に集中することを可能にします。言い換えれば研究により多くの時間を費やすことを可能にします。

この新しい連携により、PyTorch Lightningユーザーは、当社のPopTorch APIを使用しているIPU用のPyTorchモデルを、最小限のコード変更で実行し、同様の高いパフォーマンスを得られるようになりました。

PopTorchは、PyTorchのモデルをIPUハードウェア上で直接実行できるようにするPyTorchの拡張機能で、IPU上で実行するためにできるだけコードを変更しないように設計されています。

PyTorch Lightningが選ばれる理由

当社がPyTorch Lightningを選ぶ理由は、AI研究者と同じです。PyTorch Lightningは、PyTorchモデルのトレーニングループの複雑さを解消するシンプルなラッパーで、基本的なプラットフォームを抽象化するため、IPUのユーザーエクスペリエンスは他のプラットフォームに近いものとなります。

多くの定型コードが削除され、よりクリーンで使いやすい実装になっています。

PyTorch Lightningでは、PyTorchモデルのPopTorch構成と同じように、ユーザーがモデル関数を指定する必要があるため、親しみやすいユーザーエクスペリエンスを実現しています。

IPU固有の最適化や分解に干渉しないので、IPU上のPyTorch Lightningモデルは、IPU用の標準的なPyTorchモデルと同様の高いパフォーマンスを得ることができます。

使用を開始する方法

PyTorch Lightningをインストールします。https://github.com/PyTorchLightning/pytorch-lightning#step-0-install

デフォルトでは、PyTorch Lightningは最新版のPyTorchをインストールします。PopTorchでサポートされているバージョンのPyTorchを確実にインストールするには、PyTorch Lightningのインストール時に依存性のないpip3 installを使用するか、サポートされているバージョンのPyTorchを後からインストールする必要があります。

pip3 install pytorch-lightning

pip3 uninstall torch

pip3 install torch==1.7.1+cpu -f

https://download.pytorch.org/whl/torch_stable.html

PopTorchをインストールする

Graphcoreドキュメントポータルのお使いのIPUシステムに関連する「Getting Started」ガイドに記載されている通りにPoplar SDKをインストールします。

PopTorchは、インストール可能なホイールファイルとしてPoplar SDKにパッケージされています。正しくインストールされているかどうかを確認するための詳しい手順は、PopTorchユーザーガイドに記載されています。

PopTorchは現在PyTorchのバージョン1.7.1を使用していますが、今後のリリースではバージョンを上げていく予定です。パッケージによっては、より新しいバージョンのPyTorchがインストールされている場合があり、サポートされているバージョンをpipでインストールする必要があります。ユーザーガイドのインストール手順の「Version Compatibility」をご覧ください。

重要

PopTorch の依存パッケージを正しくインストールするには、pip >= 18.1 が必要です。

基本的な例を実行する

以下のコード例は、シンプルなMNISTサンプルを使ってトレーニングモデルを実行する方法を示しています。

	# Copyright The PyTorch Lightning team.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.

	import torch
	from torch.nn import functional as F

	from pytorch_lightning import LightningModule, Trainer
	from pytorch_lightning.demos.mnist_datamodule import MNISTDataModule


	class LitClassifier(LightningModule):
	def __init__(self, hidden_dim: int = 128, learning_rate: float = 0.0001):
	super().__init__()
	self.save_hyperparameters()

	self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
	self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)

	def forward(self, x):
	x = x.view(x.size(0), -1)
	x = torch.relu(self.l1(x))
	x = torch.relu(self.l2(x))
	return x

	def training_step(self, batch, batch_idx):
	x, y = batch
	y_hat = self(x)
	loss = F.cross_entropy(y_hat, y)
	return loss

	def validation_step(self, batch, batch_idx):
	x, y = batch
	probs = self(x)
	# we currently return the accuracy as the validation_step/test_step is run on the IPU devices.
	# Outputs from the step functions are sent to the host device, where we calculate the metrics in
	# validation_epoch_end and test_epoch_end for the test_step.
	acc = self.accuracy(probs, y)
	return acc

	def test_step(self, batch, batch_idx):
	x, y = batch
	logits = self(x)
	acc = self.accuracy(logits, y)
	return acc

	def accuracy(self, logits, y):
	# currently IPU poptorch doesn't implicit convert bools to tensor
	# hence we use an explicit calculation for accuracy here. Once fixed in poptorch
	# we can use the accuracy metric.
	acc = torch.sum(torch.eq(torch.argmax(logits, -1), y).to(torch.float32)) / len(y)
	return acc

	def validation_epoch_end(self, outputs) -> None:
	# since the training step/validation step and test step are run on the IPU device
	# we must log the average loss outside the step functions.
	self.log("val_acc", torch.stack(outputs).mean(), prog_bar=True)

	def test_epoch_end(self, outputs) -> None:
	self.log("test_acc", torch.stack(outputs).mean())

	def configure_optimizers(self):
	return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)


	if __name__ == "__main__":
	dm = MNISTDataModule(batch_size=32)

	model = LitClassifier()

	trainer = Trainer(max_epochs=2, accelerator="ipu", devices=8)

	trainer.fit(model, datamodule=dm)
	trainer.test(model, datamodule=dm)

view raw PyTorch Lightning mnist hosted with ❤ by GitHub

IPU上で開始する方法

大学の研究者は、GraphcoreのAcademic Programmeに応募することで、IPUにアクセスする機会が得られます。このプログラムは、IPUを使用して研究を行ったり、コースワークや教育の中で研究を発表する研究者を支援するために設計されています。選ばれた研究者は、GraphcoreのIPU計算プラットフォームをクラウドで無償利用できるほか、ソフトウェアツールやサポートを受けることができます。

IPU上でPyTorch Lightningモデルを1行のコードで実行する方法については、最新のデベロッパーチュートリアルウォークスルーブログをご覧ください。