Entrenamiento de redes neuronales usando Pytorch Lightning

Introducción:

PyTorch Lightning es una biblioteca que proporciona una interfaz de alto nivel para PyTorch. El problema con PyTorch es que cada vez que inicias un proyecto tienes que reescribir esos ciclos de entrenamiento y prueba. PyTorch Lightning soluciona el problema no solo al reducir el código repetitivo, sino también al proporcionar una funcionalidad adicional que puede resultar útil al entrenar sus redes neuronales. Una de las cosas que me encantan de Lightning es que el código está muy organizado y es reutilizable, y no solo eso, sino que reduce el ciclo de entrenamiento y prueba al tiempo que conserva la flexibilidad por la que PyTorch es conocido. Y una vez que aprenda a usarlo, verá cuán similar es el código al de PyTorch.

Instalación de PyTorch Lightning:

La instalación de Lightning es igual que la de cualquier otra biblioteca en python.

pip install pytorch-lightning

o si desea instalarlo en un entorno conda, puede usar el siguiente comando:-

conda install -c conda-forge pytorch-lightning

Formato del modelo PyTorch Lightning:

Si alguna vez ha usado PyTorch, debe saber que la definición del modelo PyTorch sigue el siguiente formato

from torch import nn

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

Así es como definimos un modelo en PyTorch ahora, después de definir el bucle, generalmente definimos pérdida, optimizador y entrenamiento fuera de la clase. En PyTorch Lightning, la forma de definir el modelo es similar excepto por el hecho de que agregamos los pasos de pérdida, optimización y entrenamiento en el propio modelo. Para definir un modelo de rayo seguimos el siguiente formato:-

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

Nota: Los nombres de las funciones anteriores deben ser exactamente iguales.

Entrenando nuestra Red Neuronal:

Cargando nuestros datos:

Para este tutorial, usaremos el conjunto de datos MNIST, por lo que comenzaremos cargando nuestros datos y luego definiendo el modelo. Para cargar datos para Lightning Model, puede definir DataLoaders como lo hace en PyTorch y pasar el cargador de datos de entrenamiento y el cargador de datos de validación en la función pl.Trainer() o puede usar LightingDataModule que hace lo mismo, excepto que ahora realiza los pasos en un python clase. Para crear cargadores de datos seguimos el siguiente paso:-

Carga de datos mediante la creación de cargadores de datos:

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)

Para crear LightningDataModule seguimos los siguientes pasos:-

Carga de datos mediante la creación de 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)

Nota: Los nombres de las funciones anteriores deben ser exactamente iguales.

Así es como se crea el módulo de datos Lightning. La creación de cargadores de datos puede complicarse, por eso es mejor agrupar el conjunto de datos en forma de Módulo de datos.

Definición de nuestra red neuronal

Definir el modelo en la iluminación de PyTorch es más o menos lo mismo que en PyTorch, excepto que ahora estamos agrupando todo dentro de nuestra clase de modelo.

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)

Discutiremos más a fondo cómo training_step() difiere de los pasos en Training Loop en Pytorch y otras diferencias entre el modelo Lightning y el modelo Pytorch.

Entrenando Nuestro Modelo

Para entrenar el modelo en Pytorch, primero debe escribir el ciclo de entrenamiento, pero la clase Trainer en Lightning facilita las tareas. Para entrenar el modelo en 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)

Nota: `dp` es DataParallel (lote dividido entre GPU de la misma máquina).

Nota: si ha cargado datos mediante la creación de cargadores de datos, puede ajustar trainer by trainer.fit(clf,trainloader,testloader) .

Diferencia entre el modelo PyTorch y el modelo Lightning:

Como podemos ver, la primera diferencia entre PyTorch y el modelo Lightning es la clase que hereda la clase del modelo:

PyTorch

class model(nn.Module):

PyTorch -Relámpago

class model(pl.LightningModule):

método __init__()

Tanto en Pytorch como en Lightning Model, usamos el método __init__() para definir nuestras capas, ya que en Lightning juntamos todo, también podemos definir otros hiperparámetros como la tasa de aprendizaje para el optimizador y la función de pérdida.

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)

Pytorch-Rayo

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()

adelante() método:

Tanto en Pytorch como en Lightning Model, usamos el método forward() para definir nuestro pase hacia adelante, por lo tanto, es el mismo para ambos.

PyTorch y PyTorch -Relámpago

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)

Definición de Optimizador:

En PyTorch, generalmente definimos nuestros optimizadores creando directamente su objeto, pero en PyTorch -lightning definimos nuestros optimizadores bajo el método configure_optimizers() . Otra cosa a tener en cuenta es que en PyTorch pasamos los parámetros del objeto del modelo como argumentos para el optimizador, pero en lightning, pasamos self.parameters() ya que la clase es el modelo en sí.

PyTorch 

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

PyTorch -Relámpago

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

Nota: También puede crear múltiples optimizadores en Lightning.

Bucle de entrenamiento (paso):

No estará mal decir que esto es lo que hace que Lightning se destaque de PyTorch . En PyTorch definimos el ciclo de entrenamiento completo mientras que en lightning usamos Trainer() para hacer el trabajo. Pero igual definimos los pasos que se van a ejecutar durante el entrenamiento.

PyTorch

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)}')

Rayo PyTorch

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


Vea cómo en los pasos de entrenamiento solo escribimos los pasos necesarios (en negrita). 

Código

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)

Publicación traducida automáticamente

Artículo escrito por herumbshandilya y traducido por Barcelona Geeks. The original can be accessed here. Licence: CCBY-SA

Deja una respuesta

Tu dirección de correo electrónico no será publicada. Los campos obligatorios están marcados con *