How can I save grad attributes within a network's state_dict?

username: AlphaBetaGamma96
cooked: <p>Hi All,</p>
<p>I was wondering if it were possible to save the grad attributes of my model for all my parameters? I’m currently using a custom optimizer in my problem and it requires using the <code>.grad</code> attribute (as I’m preconditioning gradients).</p>
<p>The way I save my model is via,</p>
<pre><code class="lang-auto">    torch.save({'epoch':epoch,
                'model_state_dict':net.state_dict(),
                'optim_state_dict':optim.state_dict(),
                'loss':loss}, model_path)
</code></pre>
<p>and is subsequently loaded via,</p>
<pre><code class="lang-auto">state_dict = torch.load(f=model_path, map_location=torch.device('cpu'))
net.load_state_dict(state_dict['model_state_dict']) #returns &lt;All keys matched successfully&gt;
optim.load_state_dict(state_dict['optim_state_dict']) #return nothing (I assume ok?)
</code></pre>
<p>If I load the model and print the grad attribute it doesn’t exist and hence causes my code to crash. So, I was wondering how exactly can I save these values? I printed the grad attribute via,</p>
<pre><code class="lang-auto">for name, param in net.named_parameters():
  print(param.grad)
</code></pre>
<p>The strange thing is I called <code>loss.backward()</code> before calling <code>optim.step()</code> so I thought that the grad attributes would be filled with a <code>grad</code> attribute? This is done in the standard way like,</p>
<pre><code class="lang-auto">optim.zero_grad()
loss_mean.backward()
optim.step()
</code></pre>
<p>Any help would be appreicated! Thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: ptrblck
cooked: <p>The <code>state_dict</code> doesn’t contain the <code>.grad</code> attributes of the parameters, but you could try to store a custom <code>dict</code> with them and reassign these gradients after loading the <code>state_dict</code> to all needed parameters.</p>
<p>EDIT: <code>state_dict</code> hooks might also be an valuable workflow you could take a look into.</p>

username: AlphaBetaGamma96
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>
<p>One thing I’ve just changed which seems to have solved the missing <code>grad</code> attributes issue is by changing how I defined the shape. So for example, when cycling through <code>optim.param_groups</code> I got the shape via the code below.</p>
<pre><code class="lang-auto">module = group['module']
g = module.weight.grad
s = g.shape
</code></pre>
<p>However, I’ve changed it to,</p>
<pre><code class="lang-auto">g = module['params'][0].grad
s = g.shape
</code></pre>
<p>which seems to work which I assume because I’m accessing it via the parameters rather than the module, and the <code>grad</code> attributes are filled when calling <code>loss.backward()</code>?</p>
<p>Thank you!</p>

username: ptrblck
cooked: <p>I’m unsure where you are assigning this shape to and how it relates to the <code>state_dict</code>. Could you explain this workflow a bit more, please?</p>

username: AlphaBetaGamma96
cooked: <p>Ok, so I’m making a custom optimizer based on KFAC and I noticed that if I resume my model from a checkpoint it fails because the <code>grad</code> attributes don’t exist. When I first initialise the optimizer, I append each module to a list and apply a forward and backward hook to each module. I <em>also</em> store the parameters of modules into the <code>params</code> list. This is done via,</p>
<pre><code class="lang-auto">  def _prepare_model(self):
    print("Adding parameters to module, marking nn.Linear layers!")
    for module in self.net.modules():
      if(module.__class__.__name__ in self._kfac_accepted_modules):
        print("module: ",module)
        self.modules.append(module)
        
        handle = module.register_forward_pre_hook(self._save_input)
        self._fwd_handles.append(handle)
        handle = module.register_full_backward_hook(self._save_output)
        self._bwk_handles.append(handle)
        
        params = [module.weight]
        if(module.bias is not None):
          params.append(module.bias)
        d = {'params':params, 'module':module}
        self.params.append(d)
</code></pre>
<p>Then when I go to precondition the gradient attribute I grab the shape of it via <code>group['module'].weight.grad.shape</code>. This didn’t work, but what did work was to change <code>group['module'].weight.grad</code> to <code>group['params'][0].grad</code> which does exist and gets created when <code>loss.backward()</code> gets called. This is done via the following, and is called on each group within optim.param_groups</p>
<pre><code class="lang-auto">def _precondition(self, group):
  module = group['module']

  g = group['params'][0].grad
  s = g.shape
  if(module.bias is not None):
    #gb = module.bias.grad       #doesn't work
    gb = group['params'][1].grad #does work
    g = torch.cat([g, gb.unsqueeze(1)], dim=1)
</code></pre>
<p>Does this clarify the workflow? Thank you!</p>

username: AlphaBetaGamma96
cooked: <p>I hope that means my solution is correct! <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>
<p>I was wondering if I could ask a follow up question? The <code>register_full_backward_hook</code> if I call a double backward on my network does the <code>backward_hook</code> recorded a double backward or is it just a single backward? And, if I call 2 in a row (i.e. a double backward with one loss function, then a single backward with a different loss function). Which output is stored?</p>
<p>my forward/backward hooks are defined like so,</p>
<pre><code class="lang-auto">  def _save_input(self, module, input):
    self.state[module]['a'] = input[0]
    
  def _save_output(self, module, grad_input, grad_output):
    self.state[module]['gl'] = grad_output[0]
</code></pre>
<p>So, I’d assume whenever a backward call is called, it overwrites the previous call?</p>
<p>Thank you for all the help! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>