Creating and accessing a custom state in optimizer

sdg91: <p>I am calculating the norm of the gradients in current step during optimisation and would like to access it in the next step.I have tried to access it through optimizer state_dict(). Below is my code for the same.</p>
<pre><code>for e in range(epochs):
correct, total, epoch_loss = 0, 0, 0.0
for batch_idx, (inputs, target) in enumerate(loader):
    # Initialize grad_prev as a tensor of zeros                 
    if 'prev_grad_norm' not in optimizer.state_dict():
        prev_grad_norm = torch.Tensor(100)
        optimizer_state = optimizer.state_dict()
        optimizer_state['prev_grad_norm'] = prev_grad_norm
        optimizer.load_state_dict(optimizer_state)
        print("1. Norm in Previous Epoch:", grad_prev)
        print("\n")
    else:
        prev_grad_norm = optimizer.state_dict()['prev_grad_norm']
        print("Norm in Previous Epoch:", prev_grad_norm)

    optimizer.zero_grad()
    data = inputs.view(inputs.shape[0], -1)
    outputs = net(data)
    loss = criterion(outputs, labels)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=1.0)
    for param in net.parameters():
      if param.grad is not None:
          param.grad += torch.randn_like(param.grad) * noise_scale 
    optimizer.step()
    current_norm = torch.max(torch.tensor([torch.norm(p.grad, 2) for p in net.parameters()]))
    print("2. Norm in current Epoch:", current_norm)
    print("\n")
    # Metrics
    epoch_loss += loss
    total += labels.size(0)
    correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
    print("3. Norm updated in current Epoch:", current_norm)
    print("\n")
    optimizer.state_dict()['prev_grad_norm'] = current_norm
    print(optimizer.state_dict())
    print("\n")
    optimizer.step()
</code></pre>
<p>Though I see that norm is being calculated at the current step, but when I try to access it through ‘prev_grad_norm’ in the next step, it is showing as None. My purpose of accessing it from previous step is to use that as threshold for clipping in next step.</p>
<p>Why the state is not being passed here?</p>

ptrblck: <p>I don’t think you can store arbitrary data into the optimizer’s <code>state_dict</code> as it should return the <code>state</code> and <code>param_groups</code> keys as seen <a href="https://github.com/pytorch/pytorch/blob/157c869026bd0aa866e0138d5ed57d09966863fc/torch/optim/optimizer.py#L340-L369">here</a>, so you might need to store it in a custom <code>state</code> class if possible.</p>

sdg91: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> , can you please share any documentation for custom state class implementation? I am fine to access previous norm by any method.Thanks!</p>

ptrblck: <p>I was simply referring to a custom <code>dict</code> that users often save for more hyperparameter and training information. E.g. you will find code such as:</p>
<pre><code class="lang-python">state = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "epochs": nb_epochs,
    "any_additional_information": value
}
</code></pre>
<p>I thus suggested to store the value also in a <code>dict</code> or another class you like instead of the optimizer’s <code>state_dict</code>.</p>