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

Giobio: <p>Hello everyone,</p>
<p>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:</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>The code seems to be running but I am worried that certain things might not be working as they should under the hood.</p>
<p>The idea would be to add more optimizers in the context of a GAN.</p>
<p>Thank you for your time !</p>

shartzog: <p><a class="mention" href="/u/giobio">@Giobio</a><br>
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:</p>
<aside class="quote" data-post="7" data-topic="11085">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/matthew_zeng/48/1068_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/two-optimizers-for-one-model/11085/7">Two optimizers for one model</a> 
  </div>
  <blockquote>
    Use the basic knowledge of software engineering. 

class MultipleOptimizer(object):
    def __init__(*op):
        self.optimizers = op

    def zero_grad(self):
        for op in self.optimizers:
            op.zero_grad()

    def step(self):
        for op in self.optimizers:
            op.step()


opt = MultipleOptimizer(optimizer1(params1, lr=lr1), 
                        optimizer2(params2, lr=lr2))

loss.backward()
opt.zero_grad()
opt.step()
  </blockquote>
</aside>

<p>There’s some other good in that post as well.  Hope that helps!<br>
–SEH</p>

Giobio: <p>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.<br>
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.</p>
<p>Here is an example:</p>
<pre><code class="lang-auto">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()
</code></pre>

shartzog: <p>Gotcha!  Hopefully, some people with more experience than I will comment, but I can’t think of a reason it <em>shouldn’t</em> work.  I’d say go for it…  <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>