Self.state is always empty in torch.optim.Adam

jonbakerfish: <p>I want to resume training and try to load the state_dict of an Adam optimizer, however I find that the loss will increase for a few hundred iterations before going down again.<br>
Later I found that the <code>self.state</code> in torch.optim.Adam is always empty which causes <a href="https://github.com/pytorch/pytorch/blob/39f94619ec3bf827567716e6c9b5003ef20036f4/torch/optim/adam.py#L71" rel="nofollow noopener">state initialization for each optimizer.step() call</a>. Shall we modify the Adam code by adding <code>self.state[p] = state</code> in the end of the loop?</p>
<p>I’m using ignite for training and <a href="https://pytorch.org/ignite/handlers.html#ignite.handlers.ModelCheckpoint" rel="nofollow noopener">its ModelCheckpoint</a> for saving model’s and optimizer’s <code>state_dict()</code>. I found all the saved optimizers’ <code>self.state</code> is an empty dict().</p>

ptrblck: <p>If I load the <code>state_dict</code> after re-initializing the optimizer, <code>optimizer.state</code> is set again:</p>
<pre><code class="lang-python"># Create dummy model and optimizer
model = nn.Linear(10, 20)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Create dummy backward pass
out = model(torch.randn(1, 10))
out.mean().backward()
optimizer.step()

# Check state and store state_dict
print(optimizer.state)
state_dict = optimizer.state_dict()

# Re-initialize and check state
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
print(optimizer.state)

# Load and check state
optimizer.load_state_dict(state_dict)
print(optimizer.state)
</code></pre>
<p>Could you post the code you are using to store and load the optimizer?</p>

jonbakerfish: <p>It seems that there is a bug in my training code. I shouldn’t pass <code>obj.state_dict()</code> to the <a href="https://pytorch.org/ignite/handlers.html#ignite.handlers.ModelCheckpoint" rel="nofollow noopener"><code>ModelCheckpoint</code></a>:</p>
<pre><code>checkpointer = ModelCheckpoint(output_dir, cfg.MODEL.NAME, checkpoint_period, n_saved=10, require_empty=False)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpointer, {'model': model.state_dict(),
                                                                 'optimizer': optimizer.state_dict(),
                                                                 'scheduler': scheduler.state_dict()})
</code></pre>
<p>Instead, I should pass the <code>obj</code> itself:</p>
<pre><code>checkpointer = ModelCheckpoint(output_dir, cfg.MODEL.NAME, checkpoint_period, n_saved=10, require_empty=False)
trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpointer, {'model': model,
                                                                 'optimizer': optimizer,
                                                                 'scheduler': scheduler})
</code></pre>
<p>The <code>model.state_dict()</code> works in both situations, but for optimizer and scheduler only the initial state will be saved. The reason might be the different implementations of <code>state_dict()</code> in <a href="https://github.com/pytorch/pytorch/blob/f6fda4409be706c2ecceaf55f44d7c75a2f22e6e/torch/nn/modules/module.py#L604" rel="nofollow noopener"><code>nn.Module</code></a>, <a href="https://github.com/pytorch/pytorch/blob/f6fda4409be706c2ecceaf55f44d7c75a2f22e6e/torch/optim/optimizer.py#L72" rel="nofollow noopener"><code>optim.Optimizer</code></a> and <a href="https://github.com/pytorch/pytorch/blob/f6fda4409be706c2ecceaf55f44d7c75a2f22e6e/torch/optim/lr_scheduler.py#L90" rel="nofollow noopener"><code>optim._LRScheduler</code></a>.</p>

MariosOreo: <p>Hi ptrblck,</p>
<p>I tried to output <code>opt.state</code> but it is always empty as this</p>
<pre><code class="lang-auto">defaultdict(&lt;class 'dict'&gt;, {})
</code></pre>
<p>What dose <code>state</code> used for?<img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=6" title=":thinking:" class="emoji" alt=":thinking:"></p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p><strong>TL;DR</strong> - Can you save <code>optim.state</code> to load in after a checkpoint? (Instead of it being re-initialized)</p>
<p>Apologizes for opening this post again (after 4 years), but is there a way to save the <code>optim.state</code> when saving the <code>state_dict</code>? Instead of it being re-initialized after re-initializing the class?</p>
<p>The case I have in mind is that, if I train some model with a custom optimizer that uses <code>self.state</code> in some way and the model is interrupted and has to resume from a checkpoint/<code>torch.save</code> file. The behavior won’t be consistent between a model, which ran without interruption and the one that did.</p>

ptrblck: <p>Yes, you can store the <code>optimizer.state_dict()</code> and load it via <code>optimizer.load_state_dict(state_dict)</code> as is also done for the model.</p>