作者：Erfandi Maula Yusnu, Lalu
编译：ronghuaiyang
首发：AI公园公众号

导读

对使用PyTorch Lightning的训练代码和原始的PyTorch代码进行了对比，展示了其简单，干净，灵活的优点，相信你会喜欢的。

PyTorch Lightning是为ML研究人员设计的轻型PyTorch封装。它帮助你扩展模型并编写更少的样板文件，同时维护代码干净和灵活同时进行扩展。它帮助研究人员更多地专注于解决问题，而不是编写工程代码。

我从两年前就开始使用PyTorch了，我从0.3.0版本开始使用。在我使用PyTorch之前，我使用Keras作为我的深度学习框架，但后来我开始切换到PyTorch，原因有几个。如果你想知道我的原因，看看下面这篇文章：https://medium.com/swlh/why-i...。

由于我一直在使用PyTorch，所以我需要牺牲在Keras中只用几行简单的行代码就可以进行训练的乐趣，而编写自己的训练代码。它有优点也有缺点，但是我选择PyTorch编写代码的方式来获得对我的训练代码的更多控制。但每当我想在深度学习中尝试一些新的模型时，就意味着我每次都需要编写训练和评估代码。

所以，我决定建立我自己的库，我称之为torchwisdom，但我陷入了困境，因为我仍在为我的公司构建OCR全pipeline系统。所以，我试图找到另一个解决方案，然后我找到了PyTorch Lightning，在我看到代码后，它让我一见钟情。

因此，我将在本文中介绍的内容是安装、基本的代码比较以及通过示例进行比较，这些示例是我自己通过从pytorch lightning site获取的，一些代码自己创建的。最后是本文的结论。

安装

好的，让我们从安装pytorch-lighting开始，这样你就可以跟着我一起做了。你可以使用pip或者conda安装pytorch lightning。

pip install

pip install pytorch-lightning

conda install

conda install pytorch-lightning -c conda-forge

对我来说，我更喜欢用anaconda作为我的python解释器，它对于深度学习和数据科学的人来说更完整。从第一次安装开始，它就自带了许多标准机器学习和数据处理库包。

基本代码的比较

在我们进入代码之前，我想让你看看下面的图片。下面有2张图片解释了pytorch和pytorch lightning在编码、建模和训练上的区别。在左边，你可以看到，pytorch需要更多的代码行来创建模型和训练。

有了pytorch lightning，代码就变成了Lightning模块的内部，所有的训练工程代码都被pytorch lightning解决了。但是你需要在一定程度上定制你的训练步骤，如下面的示例代码所示。

对于训练代码，你只需要3行代码，第一行是用于实例化模型类，第二行是用于实例化Trainer类，第三行是用于训练模型。

这个例子是用pytorch lightning训练的一种方法。当然，你可以对pytorch进行自定义风格的编码，因为pytorch lightning具有不同程度的灵活性。你想看吗？让我们继续。

通过例子进行比较

好了，在完成安装之后，让我们开始编写代码。要做的第一件事是导入需要使用的所有库。在此之后，你需要构建将用于训练的数据集和数据加载器。

# import all you needimport osimport torchimport torchvisionimport torch.nn as nnimport torch.nn.functional as Ffrom torch.utils.data import DataLoader, random_splitfrom torchvision.datasets import MNISTfrom torchvision import datasets, transformsimport pytorch_lightning as plfrom pytorch_lightning import Trainerfrom pytorch_lightning.core.lightning import LightningModule# transforms# prepare transforms standard to MNISTtransform=transforms.Compose([transforms.ToTensor(),                              transforms.Normalize((0.1307,), (0.3081,))])# datamnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)mnist_train_loader = DataLoader(mnist_train, batch_size=64)

正如上面看到的代码，我们使用来自torchvision的MNIST数据集，并使用torch.utils.DataLoader创建数据加载器。现在，在下面的代码中，我们使网络与28x28像素的MNIST数据集想匹配。第一层有128个隐藏节点，第二层有256个隐藏节点，第三层为输出层，有10个类作为输出。

# build your modelclass CustomMNIST(LightningModule):    def __init__(self):        super().__init__()        # mnist images are (1, 28, 28) (channels, width, height)        self.layer1 = torch.nn.Linear(28 * 28, 128)        self.layer2 = torch.nn.Linear(128, 256)        self.layer3 = torch.nn.Linear(256, 10)    def forward(self, x):        batch_size, channels, width, height = x.size()        # (b, 1, 28, 28) -> (b, 1*28*28)        x = x.view(batch_size, -1)        x = self.layer1(x)        x = torch.relu(x)        x = self.layer2(x)        x = torch.relu(x)        x = self.layer3(x)        x = torch.log_softmax(x, dim=1)        return x    def training_step(self, batch, batch_idx):        data, target = batch        logits = self.forward(data)        loss = F.nll_loss(logits, target)        return {'loss': loss}    def configure_optimizers(self):        return torch.optim.Adam(self.parameters(), lr=1e-3)# train your modelmodel = CustomMNIST()trainer = Trainer(max_epochs=5, gpus=1)

如果你在上面的gist代码中看到第27和33行，你会看到_training\_step_和_configure\_optimators_方法，它覆盖了在第2行中扩展的类_LightningModule_中的方法。这使得pytorch中标准的_nn.Module_不同于_LightningModule_，它有一些方法使它与第39行中的_Trainer_类兼容。

现在，让我们尝试另一种方法来编写代码。假设你必须编写一个库，或者希望其他人使用纯pytorch编写的库。你该怎样使用pytorch lightning？

下面的代码有两个类，第一个类使用标准的pytorch的nn.Module作为其父类。它是按照标准pytorch模块中通常编写的方式编写的，但是看第30行，有一个名为_ExtendMNIST_的类继承了两个类。这两个类由StandardMNIST类和LightningModule类组合在一起。这就是我喜欢python的地方，一个类可以有多个父类。

# build your modelclass StandardMNIST(nn.Module):    def __init__(self):        super().__init__()        # mnist images are (1, 28, 28) (channels, width, height)        self.layer1 = torch.nn.Linear(28 * 28, 128)        self.layer2 = torch.nn.Linear(128, 256)        self.layer3 = torch.nn.Linear(256, 10)        def forward(self, x):        batch_size, channels, width, height = x.size()                # (b, 1, 28, 28) -> (b, 1*28*28)        x = x.view(batch_size, -1)                x = self.layer1(x)        x = torch.relu(x)                x = self.layer2(x)        x = torch.relu(x)                x = self.layer3(x)        x = torch.log_softmax(x, dim=1)                return x# extend StandardMNIST and LightningModule at the same time# this is what I like from python, extend two class at the same timeclass ExtendMNIST(StandardMNIST, LightningModule):    def __init__(self):        super().__init__()          def training_step(self, batch, batch_idx):        data, target = batch        logits = self.forward(data)        loss = F.nll_loss(logits, target)        return {'loss': loss}        def configure_optimizers(self):        return torch.optim.Adam(self.parameters(), lr=1e-3)# run the trainingmodel = ExtendMNIST()trainer = Trainer(max_epochs=5, gpus=1)trainer.fit(model, mnist_train_loader)

如果你看到ExtendMNIST类中的代码，你会看到它只是覆盖了LightningModule类。使用这种编写代码的方法，你可以扩展以前编写的任何其他模型，而无需更改它，并且仍然可以使用pytorch lightning库。

那么，你能在训练时给我看一下结果吗？好，让我们看看它在训练时是什么样子。