介绍适用于Graphcore IPU的PyTorch Lightning初始版本

2021/07/27

介绍适用于Graphcore IPU的PyTorch Lightning初始版本

作者:

我们非常兴奋地宣布PyTorch Lightning现支持Graphcore IPU。过去几个月里，PyTorch Lightning团队一直在为构建IPU集成而辛勤工作，并在现在将其通过1.4版本提供给社区。我们非常感谢他们与Graphcore团队的密切合作，帮助我们实现让IPU更易于开发人员使用的使命。

Graphcore大力支持PyTorch Lightning的使命，即满足AI研究社区对灵活、快速的AI计算解决方案不断增长的需求。

PyTorch Lightning将数据科学家和深度学习从业者从繁重的工程任务（数据分发、循环管理、日志处理等）中解放出来，使他们能够专注于建模和数据理解，换言之，他们可以将更多时间专注于研究。

这种新的集成意味着PyTorch Lightning用户现在可以针对IPU使用任何PyTorch模型，这些模型使用我们的PopTorch API，以最少的代码更改运行它们，并获得相同的高性能。

PopTorch是PyTorch的一组扩展，它使PyTorch模型能够直接在IPU硬件上运行，并且旨在在IPU上运行时尽可能少地更改代码。

我们为什么喜欢PyTorch Lightning

我们喜欢PyTorch Lightning的原因与AI研究人员相同：它是一个简单的包装器，消除了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无依赖项安装，或者在之后安装受支持的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系统的相关“入门”指南中的说明安装Poplar SDK。

PopTorch与Poplar SDK打包在一起，作为可安装的轮文件。可以在PopTorch用户指南中找到验证它是否已正确安装的完整说明。

PopTorch当前使用PyTorch版本1.7.1，我们将在未来版本中扩充它。某些软件包可能会安装较新版本的PyTorch，您可能需要pip安装受支持的版本。请参阅用户指南安装步骤的“版本兼容性”部分。

重要

要正确安装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
 
import pytorch_lightning as pl
from pl_examples.basic_examples.mnist_datamodule import MNISTDataModule
 
 
class LitClassifier(pl.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 = pl.Trainer(max_epochs=2, ipus=8)
 
    trainer.fit(model, datamodule=dm)
    trainer.test(model, datamodule=dm)

开始使用IPU

大学研究人员可以申请Graphcore学术计划，以获得使用IPU的机会。该项目旨在支持学者使用IPU进行研究并发表相关成果，或在他们的课程作业或教学中使用IPU。被选中参与的研究人员将受益于免费访问云上的Graphcore IPU计算平台以及软件工具和支持。

了解有关如何使用一行代码在IPU上运行PyTorch Lightning模型的更多信息，请阅读我们最新的开发人员教程博客。

参考资料和链接

Graphcore PyTorch Lightning教程和示例

PyTorch Lightning IPU支持文档

IPU的PyTorch：用户指南

PopTorch编译说明

PopTorch源码库