Dev adam graph conf

python/oneflow/nn/optimizer/adam.py

@@ -14,6 +14,7 @@
 limitations under the License.
 """
 import collections
+import math
 from typing import Callable, Dict, Iterator, List, Tuple, Union
 
 import oneflow as flow
@@ -138,3 +139,38 @@ def step(self, closure: Callable = None):
                     self._op(param, param.grad, m_tensor, v_tensor, **kwargs)
             self._state["step"] = self._state["step"] + 1
             return loss
+
+    def generate_conf_for_graph(self, train_conf, vars_conf):
+        for param_group in self.param_groups:
+            optimizer_conf = train_conf.mutable_optimizer_conf().Add()
+
+            lr = param_group["lr"]
+            scale = param_group["scale"]
+            l2 = param_group["weight_decay"]
+            beta1 = param_group["betas"][0]
+            beta2 = param_group["betas"][1]
+
+            epsilon = param_group["eps"]
+            # TODO(): optimizer_conf need to have loss_scale_factor field to support multi scale factor
+            base_scale = train_conf.loss_scale_factor()
+            assert math.isclose(base_scale, 1, rel_tol=1e-4) or math.isclose(
+                scale, base_scale, rel_tol=1e-4
+            ), "nn.Graph only support one scale factor at the moment, base_scale {} vs scale {}".format(
+                base_scale, scale
+            )
+
+            train_conf.set_loss_scale_factor(scale)
+            optimizer_conf.set_base_learning_rate(lr)
+
+            optimizer_conf.mutable_adam_conf().set_beta1(beta1)
+            optimizer_conf.mutable_adam_conf().set_beta2(beta2)
+            optimizer_conf.mutable_adam_conf().set_epsilon(epsilon)
+            optimizer_conf.mutable_adam_conf().set_do_bias_correction(
+                False
+            )  # TODO(zzk): Check this option
+
+            for param in param_group.parameters:
+                vars_conf[param].l2 = l2
+                if not param.requires_grad:
+                    continue
+                optimizer_conf.add_variable_op_names(vars_conf[param].name)

python/oneflow/test/graph/test_graph_optimizer.py

@@ -26,7 +26,7 @@
 
 @flow.unittest.skip_unless_1n1d()
 class TestGraphOptimizer(flow.unittest.TestCase):
-    def test_optimizer(test_case):
+    def test_sgd_optimizer(test_case):
         class CustomModule(flow.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -72,6 +72,54 @@ def build(self, x):
         print("g.config.proto: \n", g.config.proto)
         print("graph proto: \n", g._graph_proto)
 
+    def test_adam_optimizer(test_case):
+        class CustomModule(flow.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.para0 = flow.nn.Parameter(flow.Tensor(10, 4))
+
+            def forward(self, x):
+                x = flow.F.matmul(x, self.para0)
+                return x
+
+        m = CustomModule()
+        learning_rate = 0.1
+        betas = [0.9, 0.999]
+        scale = 0.3
+        weight_decay = 0.7
+        eps = 1e-8
+        adam0 = flow.optim.Adam(
+            [
+                {
+                    "params": [m.para0],
+                    "lr": learning_rate,
+                    "betas": betas,
+                    "eps": eps,
+                    "weight_decay": weight_decay,
+                    "scale": scale,
+                }
+            ]
+        )
+
+        class CustomGraph0(flow.nn.Graph):
+            def __init__(self):
+                super().__init__()
+                self.m = m
+                self.add_optimizer("adam0", adam0)
+
+            def build(self, x):
+                out = self.m(x)
+                out.backward()
+                return out
+
+        g = CustomGraph0()
+        x = flow.Tensor(4, 10)
+        flow.nn.init.uniform_(x, a=-1.0, b=1.0)
+        z = g._compile(x)
+        print("repr(g): \n", repr(g))
+        print("g.config.proto: \n", g.config.proto)
+        print("graph proto: \n", g._graph_proto)
+
     def test_multi_optimizer_conf(test_case):
         class CustomModule(flow.nn.Module):
             def __init__(self):