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_adam_optim.py

@@ -0,0 +1,128 @@
+import unittest
+from collections import OrderedDict
+
+import numpy as np
+
+import oneflow as flow
+import oneflow.unittest
+
+
+@flow.unittest.skip_unless_1n1d()
+def compare_with_numpy_adam(
+    test_case,
+    device,
+    x_shape,
+    scale,
+    learning_rate,
+    train_iters,
+    betas,
+    weight_decay,
+    eps,
+):
+    random_grad_seq = []
+    for _ in range(train_iters):
+        random_grad_seq.append(np.random.uniform(size=x_shape).astype(np.float32))
+    init_value = np.random.uniform(size=x_shape).astype(np.float32)
+
+    class CustomModule(flow.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.para0 = flow.nn.Parameter(flow.Tensor(init_value, device=flow.device(device)))
+
+        def forward(self, mask):
+            return self.para0 * mask
+
+
+    simp_module = CustomModule()
+    simp_module.to("cuda")
+    simp_module.train()
+
+    adam0 = flow.optim.Adam(
+            [
+                {
+                    "params": simp_module.parameters(),
+                    "lr": learning_rate,
+                    "betas": betas,
+                    "eps": eps,
+                    "weight_decay": weight_decay,
+                    "scale": scale,
+                }
+            ]
+        )
+
+    class CustomAdamGraph(flow.nn.Graph):
+        def __init__(self):
+            super().__init__()
+            self.m = simp_module
+            self.add_optimizer("adam", adam0)
+
+        def build(self, mask_tensor):
+            loss = flow.sum(self.m(mask_tensor))
+            loss.backward()
+            return loss
+
+    of_res_list = []
+    adam_graph = CustomAdamGraph()
+
+    for i in range(train_iters):
+        mask_tensor = flow.Tensor(
+                random_grad_seq[i], requires_grad=False, device=flow.device(device)
+            )
+        adam_x = adam_graph(mask_tensor)
+
+        of_res_list.append(simp_module.para0.numpy())
+
+    np_res_list = []
+
+    def train_by_numpy():
+        x = init_value
+        vt = np.zeros_like(x)
+        st = np.zeros_like(x)
+        beta1 = betas[0]
+        beta2 = betas[1]
+
+        def np_train_one_iter(grad):
+            grad = grad * scale + weight_decay * x
+            v = beta1 * vt + (1 - beta1) * grad
+            s = beta2 * st + (1 - beta2) * grad * grad
+            param = x - learning_rate * (v / (np.sqrt(s) + eps))
+            return (param, v, s)
+
+        for i in range(train_iters):
+            (x, vt, st) = np_train_one_iter(random_grad_seq[i])
+            np_res_list.append(x)
+        return x
+
+    train_by_numpy()
+
+    test_case.assertTrue(
+        np.allclose(of_res_list, np_res_list, rtol=0.001, atol=0.001)
+    )
+
+
+@flow.unittest.skip_unless_1n1d()
+class TestAdam(flow.unittest.TestCase):
+    def test_adam1(test_case):
+        compare_with_numpy_adam(test_case, 
+                                device="cuda", 
+                                x_shape=(1, ), 
+                                scale=1.0, 
+                                learning_rate=1, 
+                                train_iters=10, 
+                                betas=(0.99, 0.9), 
+                                weight_decay=0.0, 
+                                eps=1e-8)
+
+    def test_adam2(test_case):
+        compare_with_numpy_adam(test_case, 
+                                device="cuda", 
+                                x_shape=(1, ), 
+                                scale=0.8, 
+                                learning_rate=1, 
+                                train_iters=10, 
+                                betas=(0.99, 0.9), 
+                                weight_decay=0.0005, 
+                                eps=1e-8)
+
+if __name__ == "__main__":
+    unittest.main()

python/oneflow/test/graph/test_graph_sgd_optim.py

@@ -0,0 +1,118 @@
+import unittest
+from collections import OrderedDict
+
+import numpy as np
+
+import oneflow as flow
+import oneflow.unittest
+
+
+@flow.unittest.skip_unless_1n1d()
+def compare_with_numpy_sgd(
+    test_case,
+    device,
+    x_shape,
+    scale,
+    learning_rate,
+    train_iters,
+    momentum,
+    weight_decay,
+    eps,
+):
+    random_grad_seq = []
+    for _ in range(train_iters):
+        random_grad_seq.append(np.random.uniform(size=x_shape).astype(np.float32))
+    init_value = np.random.uniform(size=x_shape).astype(np.float32)
+
+    class CustomModule(flow.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.para0 = flow.nn.Parameter(flow.Tensor(init_value, device=flow.device(device)))
+
+        def forward(self, mask):
+            return self.para0 * mask
+
+
+    simp_module = CustomModule()
+    simp_module.to("cuda")
+    simp_module.train()
+
+    sgd0 = flow.optim.SGD(
+            [
+                {
+                    "params": simp_module.parameters(),
+                    "lr": learning_rate,
+                    "momentum": momentum,
+                    "weight_decay": weight_decay,
+                    "scale": scale,
+                }
+            ]
+        )
+
+    class CustomSGDGraph(flow.nn.Graph):
+        def __init__(self):
+            super().__init__()
+            self.m = simp_module
+            self.add_optimizer("sgd", sgd0)
+
+        def build(self, mask_tensor):
+            loss = flow.sum(self.m(mask_tensor))
+            loss.backward()
+            return loss
+
+    of_res_list = []
+    sgd_graph = CustomSGDGraph()
+    for i in range(train_iters):
+        mask_tensor = flow.Tensor(
+                random_grad_seq[i], requires_grad=False, device=flow.device(device)
+            )
+        sgd_x = sgd_graph(mask_tensor)
+        of_res_list.append(simp_module.para0.numpy())
+
+    np_res_list = []
+    def train_by_numpy():
+        x = init_value
+        vt = np.zeros_like(x)
+
+        def np_train_one_iter(grad):
+            grad = grad * scale + weight_decay * x
+            v = momentum * vt - learning_rate * grad
+            param = x + v
+            return (param, v)
+
+        for i in range(train_iters):
+            (x, vt) = np_train_one_iter(random_grad_seq[i])
+            np_res_list.append(x)
+
+    train_by_numpy()
+
+    test_case.assertTrue(
+        np.allclose(np_res_list, of_res_list, rtol=0.001, atol=0.001)
+    )
+
+
+@flow.unittest.skip_unless_1n1d()
+class TestSGD(flow.unittest.TestCase):
+    def test_sgd1(test_case):
+        compare_with_numpy_sgd(test_case, 
+                                device="cuda", 
+                                x_shape=(1, ), 
+                                scale=1.0, 
+                                learning_rate=1, 
+                                momentum=0.9, 
+                                train_iters=10, 
+                                weight_decay=0.0, 
+                                eps=1e-8)
+    def test_sgd2(test_case):
+        compare_with_numpy_sgd(test_case, 
+                                device="cuda", 
+                                x_shape=(1, ), 
+                                scale=1.0, 
+                                learning_rate=0.01, 
+                                momentum=0.0, 
+                                train_iters=10, 
+                                weight_decay=0.0005, 
+                                eps=1e-8)
+
+if __name__ == "__main__":
+    unittest.main()