Access the optimizer state_dict

afshin67 · January 10, 2022, 4:45pm

In Python one can get the the underlying data of an optimizer by calling the optimizer.state_dict(). For example for the Adam optimizer, it returns the number of steps, exponential moving average of gradient values, and exponential moving average of squared gradient values. I was trying to access the same values in the C++ version of torch, but it seems that it is not straight-forward to get them. I looked into the Adam optimizer implementation at:

github.com

pytorch/pytorch/blob/master/torch/csrc/api/src/optim/adam.cpp

#include <torch/optim/adam.h>

#include <torch/csrc/autograd/variable.h>
#include <torch/nn/module.h>
#include <torch/serialize/archive.h>
#include <torch/utils.h>

#include <ATen/ATen.h>
#include <c10/util/irange.h>

#include <cmath>
#include <functional>

namespace torch {
namespace optim {

AdamOptions::AdamOptions(double lr) : lr_(lr) {}

bool operator==(const AdamOptions& lhs, const AdamOptions& rhs) {
  return (lhs.lr() == rhs.lr()) &&

This file has been truncated. show original

But, still it is not clear to me how to get those.

I looked into optimizer->state() which can be obtained by:
ska::flat_hash_map<std::string, std::unique_ptr<torch::optim::OptimizerParamState>>& state_ = optimizer->state();
and it does not look like what I was expecting.

Does anyone know how can I access the number of steps, exponential moving average of gradient values, and exponential moving average of squared gradient values in C++ API?