¿PyTorch Lightning – LightningModule para ModuleList / ModuleDict?

Answer

1. 

   davy.ai

   July 17, 2023 at 10:26 pm

   Para crear equivalentes de ModuleList y ModuleDict en PyTorch Lightning, necesitamos utilizar los objetos nn.ModuleList y nn.ModuleDict, los cuales son subclases de nn.Module.

   Podemos crear un equivalente de nn.ModuleList definiendo un atributo List en nuestra clase LightningModule e inicializándolo con una lista vacía. Luego, podemos agregar cualquier módulo que deseemos a esta lista durante el método __init__. El atributo List también se puede acceder como una lista normal de Python, lo que nos permite iterar sobre su contenido y llamar a sus métodos según sea necesario.

   De manera similar, podemos crear un equivalente de nn.ModuleDict definiendo un atributo Dict en nuestra clase LightningModule e inicializándolo con un diccionario vacío. Luego, podemos agregar cualquier módulo que deseemos a este diccionario durante el método __init__, donde las llaves del diccionario son los nombres de los módulos. El atributo Dict se puede acceder como un diccionario normal de Python, lo que nos permite acceder a los módulos individuales por sus nombres de llave y llamar a sus métodos según sea necesario.

   Aquí hay un ejemplo de cómo podría verse:

   import torch
   import torch.nn as nn
   import pytorch_lightning as pl
   
   class MyModule(pl.LightningModule):
       def __init__(self):
           super().__init__()
   
           self.module_list = nn.ModuleList([])
           self.module_list.append(nn.Linear(10, 5))
           self.module_list.append(nn.ReLU())
   
           self.module_dict = nn.ModuleDict({
               'layer1': nn.Linear(5, 2),
               'layer2': nn.Sigmoid()
           })
   
       def forward(self, x):
           for layer in self.module_list:
               x = layer(x)
   
           x = self.module_dict['layer1'](x)
           x = self.module_dict['layer2'](x)
   
           return x
   

   En este ejemplo, creamos un equivalente de nn.ModuleList con dos capas: una capa lineal y una capa de ReLU. También creamos un equivalente de nn.ModuleDict con dos capas, donde la primera capa se llama 'layer1' y la segunda capa se llama 'layer2'. En el método forward, iteramos sobre module_list y pasamos x por cada capa, y luego accedemos y pasamos x a través de las dos capas en module_dict usando sus respectivas llaves.

Comments are closed.