¿PyTorch Lightning – LightningModule para ModuleList / ModuleDict?
PyTorch tiene ModuleList
y ModuleDict
. ¿Cómo puedo crear equivalentes en PyTorch Lightning para obtener todas las propiedades útiles de los LightningModule
?
Si tienes alguna pregunta, puedes hacerla a continuación o ingresar lo que estás buscando.
davy.ai
Para crear equivalentes de
ModuleList
yModuleDict
en PyTorch Lightning, necesitamos utilizar los objetosnn.ModuleList
ynn.ModuleDict
, los cuales son subclases denn.Module
.Podemos crear un equivalente de
nn.ModuleList
definiendo un atributoList
en nuestra claseLightningModule
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 atributoList
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 atributoDict
en nuestra claseLightningModule
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 atributoDict
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:
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 denn.ModuleDict
con dos capas, donde la primera capa se llama'layer1'
y la segunda capa se llama'layer2'
. En el métodoforward
, iteramos sobremodule_list
y pasamosx
por cada capa, y luego accedemos y pasamosx
a través de las dos capas enmodule_dict
usando sus respectivas llaves.