Attaching an optimizer as an attribute to a nn.Module instance

Hello everyone,

I am building a modular framework for a series of tests. I was wondering if adding an optimizer instance to a nn.Module would cause any issues. Here is an example:

class EgModel(nn.Module):
   def __init__(self): = nn.Sequential(...)
       self.optimizer = {‘optimizer_a’: optim.Adam(self.parameters())}

  def optimize(self, loss):

model = EgModel()

The code seems to be running but I am worried that certain things might not be working as they should under the hood.

The idea would be to add more optimizers in the context of a GAN.

Thank you for your time !

1 Like

What’s the driver for storing the optimizers within your nn.Module class? I don’t think that would break anything, but it seems unrelated to the ability to use multiple optimizers. For example, you could probably handle it with a custom optimizer class as described here:

There’s some other good in that post as well. Hope that helps!

Thanks for reaching out ! To be more precise, my models inherit from a BaseClass that itself is a combination of nn.Module and an Abstract Class.
The model is then fed in a pipeline that can quickly test different configurations of the underlying model. I wanted the optimizer to be defined here so that when a new architecture is implemented it could go into the pipeline independently of if it is a GAN or simple DCNN.

Here is an example:

def optimize(self, losses):                                                                                        
         gen_losses = ['G_Trick', 'L1_img', 'VGG_img', 'L1_probes']                                                     
         disc_losses = ['D_Fakes', 'D_Reals']                                                                           
         gen_loss = sum([losses[x] for x in gen_losses])                                                                
         disc_loss = sum([losses[x] for x in disc_losses])  

Gotcha! Hopefully, some people with more experience than I will comment, but I can’t think of a reason it shouldn’t work. I’d say go for it… :wink: