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

autograd
1

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):
       self.net = nn.Sequential(...)
       self.optimizer = {‘optimizer_a’: optim.Adam(self.parameters())}
   
...

  def optimize(self, loss):
       self.optimizers['optimizer_a'].zero_grad()
       loss.backward()
       self.optimizers['optimizer_a'].step()

model = EgModel()
...
model.optimise(loss)

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
2

@Giobio
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!
–SEH

4

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])                                                                
         self.optimizers['optimizer_G'].zero_grad()                                                                                 
         gen_loss.backward()                                                                                            
         self.optimizers['optimizer_G'].step()                                                                          
                                                                                                                        
         disc_loss = sum([losses[x] for x in disc_losses])  
         self.optimizers['optimizer_D'].zero_grad()                                                                                                                                 
         disc_loss.backward()                                                                                           
         self.optimizers['optimizer_D'].step()
5

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: