使用 Pytorch Lightning 训练神经网络

介绍：

PyTorch Lightning 是一个为 PyTorch 提供高级接口的库。 PyTorch 的问题在于，每次开始一个项目时，您都必须重写那些训练和测试循环。 PyTorch Lightning 不仅减少了样板代码，还提供了在训练神经网络时可能会派上用场的附加功能，从而解决了这个问题。我喜欢 Lightning 的一件事是代码非常有条理和可重用，不仅如此，它还减少了训练和测试循环，同时保留了 PyTorch 众所周知的灵活性。一旦你学会了如何使用它，你就会看到它的代码与 PyTorch 的代码有多相似。

安装 PyTorch 闪电：

安装 Lightning 与在Python安装任何其他库相同。

pip install pytorch-lightning

或者，如果您想在 conda 环境中安装它，您可以使用以下命令：-

conda install -c conda-forge pytorch-lightning

PyTorch 闪电模型格式：

如果你曾经使用过 PyTorch 你一定知道定义 PyTorch 模型遵循以下格式

from torch import nn

class model(nn.Module):
    def __init__(self):
        # Define Model Here
        
    def forward(self, x):
        # Define Forward Pass Here

这就是我们在 PyTorch 中定义模型的方式，在定义循环之后，我们通常在类外定义损失、优化器和训练。在 PyTorch Lightning 中，定义模型的方式类似，只是我们在模型本身中添加了损失、优化器和训练步骤。要定义闪电模型，我们遵循以下格式：-

import pytorch-lightning as pl

class model(pl.LightningModule):
    def __init__(self):
        # Define Model Here
        
    def forward(self, x):
        # Define Forward Pass Here
    
    def configure_optimizers(self):
       # Define Optimizer Here
       
    def training_step(self, train_batch, batch_idx):
        # Define Training loop steps here
        
    def validation_step(self, valid_batch, batch_idx):
        # Define Validation loop steps here

注意：上述函数的名称应完全相同。

训练我们的神经网络：

加载我们的数据：

在本教程中，我们将使用 MNIST 数据集，因此我们将首先加载我们的数据并在之后定义模型。要为 Lightning 模型加载数据，您可以像在 PyTorch 中一样定义 DataLoaders 并在 pl.Trainer()函数传递训练数据加载器和验证数据加载器，或者您可以使用 LightingDataModule 做同样的事情，除了现在您在Python执行这些步骤班级。要创建数据加载器，我们遵循以下步骤：-

通过创建 DataLoaders 加载数据：

from torchvision import datasets,transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.ToTensor()
])

train = datasets.MNIST('',train = True, download = True, transform=transform)
test = datasets.MNIST('',train = False, download = True, transform=transform)

trainloader = DataLoader(train, batch_size= 32, shuffle=True)
testloader = DataLoader(test, batch_size= 32, shuffle=True)

要创建LightningDataModule，我们遵循以下步骤：-

通过创建 LightningDataModule 加载数据：

import pytorch-lightning as pl
from torchvision import datasets,transforms
from torch.utils.data import DataLoader

class Data(pl.LightningDataModule):
    def prepare_data(self):
        transform=transforms.Compose([
            transforms.ToTensor()
        ])
      
        self.train_data = datasets.MNIST('', train=True, download=True, transform=transform)
        self.test_data = datasets.MNIST('', train=False, download=True, transform=transform)

    def train_dataloader(self):
        return DataLoader(self.train_data, batch_size= 32, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.test_data, batch_size= 32, shuffle=True)

注意：上述函数的名称应完全相同。

这就是您创建闪电数据模块的方式。创建数据加载器可能会变得混乱，这就是为什么最好以数据模块的形式对数据集进行分组。

定义我们的神经网络

在 PyTorch 照明中定义模型与在 PyTorch 中定义模型几乎相同，除了现在我们将模型类中的所有内容都放在一起。

from torch import nn
import pytorch_lightning as pl
import torch.nn.functional as F
from torch.optim import SGD

class model(pl.LightningModule):
    def __init__(self):
        super(model,self).__init__()
        self.fc1 = nn.Linear(28*28,256)
        self.fc2 = nn.Linear(256,128)
        self.out = nn.Linear(128,10)
        self.lr = 0.01
        self.loss = nn.CrossEntropyLoss()
    
    def forward(self,x):
        batch_size, _, _, _ = x.size()
        x = x.view(batch_size,-1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.out(x)
    
    def configure_optimizers(self):
        return SGD(self.parameters(),lr = self.lr)
    
    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        logits = self.forward(x)
        loss = self.loss(logits,y)
        return loss
    
    def validation_step(self, valid_batch, batch_idx):
        x, y = valid_batch
        logits = self.forward(x)
        loss = self.loss(logits,y)

我们将进一步讨论 training_step() 与 Pytorch 中训练循环中的步骤有何不同以及 Lightning 模型和 Pytorch 模型之间的其他差异。

训练我们的模型

要在 Pytorch 中训练模型，您首先必须编写训练循环，但 Lightning 中的 Trainer 类使任务更容易。在 Lightning 中训练模型：-

# Create Model Object
clf = model()
# Create Data Module Object
mnist = Data()
# Create Trainer Object
trainer = pl.Trainer(gpus=1,accelerator='dp',max_epochs=5)
trainer.fit(clf,mnist)

注意： `dp` 是 DataParallel（在同一台机器的 GPU 之间拆分批处理）。

注意：如果您通过创建数据加载器加载了数据，您可以通过trainer.fit(clf,trainloader,testloader)来拟合训练器。

PyTorch 模型和闪电模型的区别：

我们可以看到PyTorch和闪电模型之间的第一个区别是模型类继承的类：-

火炬

class model(nn.Module):

PyTorch -闪电

class model(pl.LightningModule):

__init__() 方法

在Pytorch和 Lightning 模型中，我们使用 __init__() 方法来定义我们的层，因为在 Lightning 中，我们将所有东西放在一起，我们还可以定义其他超参数，例如优化器的学习率和损失函数。

火炬

def __init__(self):
    super(model,self).__init__()
    self.fc1 = nn.Linear(28*28,256)
    self.fc2 = nn.Linear(256,128)
    self.out = nn.Linear(128,10)

Pytorch-闪电

def __init__(self):
    super(model,self).__init__()
    self.fc1 = nn.Linear(28*28,256)
    self.fc2 = nn.Linear(256,128)
    self.out = nn.Linear(128,10)
    self.lr = 0.01
    self.loss = nn.CrossEntropyLoss()

转发（）方法：

在Pytorch和 Lightning 模型中，我们都使用 forward() 方法来定义我们的前向传递，因此两者都是相同的。

PyTorch和PyTorch -Lightning

def forward(self,x):
    batch_size, _, _, _ = x.size()
    x = x.view(batch_size,-1)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    return self.out(x)

定义优化器：

在PyTorch 中，我们通常通过直接创建对象来定义优化器，但在PyTorch -lightning 中，我们在configure_optimizers()方法下定义优化器。另一件要注意的事情是，在PyTorch 中，我们将模型对象参数作为优化器的参数传递，但在 Lightning 中，我们传递self.parameters()，因为类是模型本身。

火炬

from torch.optim import SGD
clf = model()    # Pytorch Model Object
optimizer = SGD(clf.parameters(),lr=0.01)

PyTorch -闪电

def configure_optimizers(self):
    return SGD(self.parameters(),lr = self.lr)

注意：您也可以在 Lightning 中创建多个优化器。

训练循环（步骤）：

说这就是 Lightning 从PyTorch 中脱颖而出的原因并没有错。在 PyTorch 中，我们定义了完整的训练循环，而在 Lightning 中，我们使用 Trainer() 来完成这项工作。但是我们仍然定义了训练时将要执行的步骤。

火炬

epochs = 5

for i in range(epochs):
    train_loss = 0.0
    for data,label in trainloader:
        if is_gpu:
            data, label = data.cuda(), label.cuda()
        output = model(data)
        optimizer.zero_grad()
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item() * data.size(0)
    print(f'Epoch: {i+1} / {epochs} \t\t\t Training Loss:{train_loss/len(trainloader)}')

PyTorch-闪电

def training_step(self, train_batch, batch_idx):
    x, y = train_batch
    logits = self.forward(x)
    loss = self.loss(logits,y)
    return loss

看看我们如何在训练步骤中只编写必要的步骤（粗体）。

代码

Python3

import torch
from torch import nn
import pytorch_lightning as pl
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.optim import SGD
  
  
class model(pl.LightningModule):
    def __init__(self):
        super(model, self).__init__()
        self.fc1 = nn.Linear(28*28, 256)
        self.fc2 = nn.Linear(256, 128)
        self.out = nn.Linear(128, 10)
        self.lr = 0.01
        self.loss = nn.CrossEntropyLoss()
  
    def forward(self, x):
        batch_size, _, _, _ = x.size()
        x = x.view(batch_size, -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.out(x)
  
    def configure_optimizers(self):
        return torch.optim.SGD(self.parameters(), lr=self.lr)
  
    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        logits = self.forward(x)
        loss = self.loss(logits, y)
        return loss
  
    def validation_step(self, valid_batch, batch_idx):
        x, y = valid_batch
        logits = self.forward(x)
        loss = self.loss(logits, y)
  
  
class Data(pl.LightningDataModule):
    def prepare_data(self):
        transform = transforms.Compose([
            transforms.ToTensor()
        ])
  
        self.train_data = datasets.MNIST(
            '', train=True, download=True, transform=transform)
        self.test_data = datasets.MNIST(
            '', train=False, download=True, transform=transform)
  
    def train_dataloader(self):
        return DataLoader(self.train_data, batch_size=32, shuffle=True)
  
    def val_dataloader(self):
        return DataLoader(self.test_data, batch_size=32, shuffle=True)
  
  
clf = model()
mnist = Data()
trainer = pl.Trainer(gpus=1, accelerator='dp', max_epochs=5)
trainer.fit(clf, mnist)