Function test_sgd_complex

tests/python/relax/test_training_optimizer.py:96–134

()

def test_sgd_complex():
    x = relax.Var("x", R.Tensor((3, 3), "float32"))
    y = relax.Var("y", R.Tensor((3,), "float32"))
    sgd = SGD(0.01, 0.02).init([x, y]).get_function()

    @R.function
    def sgd_expected(
        params: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
        gradients: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
        optim_states: R.Tuple(R.Tensor((), "int64")),
    ) -> R.Tuple(
        R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
        R.Tuple(R.Tensor((), "int64")),
    ):
        R.func_attr({"global_symbol": "SGD"})
        with R.dataflow():
            num_steps: R.Tensor((), "int64") = optim_states[0]
            num_steps_new: R.Tensor((), "int64") = R.add(num_steps, R.const(1, "int64"))
            x: R.Tensor((3, 3), "float32") = params[0]
            x_grad: R.Tensor((3, 3), "float32") = gradients[0]
            lv: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.02, "float32"), x)
            x_grad_new: R.Tensor((3, 3), "float32") = R.add(lv, x_grad)
            lv1: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.01, "float32"), x_grad_new)
            x_new: R.Tensor((3, 3), "float32") = R.subtract(x, lv1)
            y: R.Tensor((3,), "float32") = params[1]
            y_grad: R.Tensor((3,), "float32") = gradients[1]
            lv2: R.Tensor((3,), "float32") = R.multiply(R.const(0.02, "float32"), y)
            y_grad_new: R.Tensor((3,), "float32") = R.add(lv2, y_grad)
            lv3: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), y_grad_new)
            y_new: R.Tensor((3,), "float32") = R.subtract(y, lv3)
            params_new: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")) = (
                x_new,
                y_new,
            )
            optim_states_new: R.Tuple(R.Tensor((), "int64")) = (num_steps_new,)
            R.output(params_new, optim_states_new)
        return (params_new, optim_states_new)

    assert_structural_equal(sgd, sgd_expected)


def test_momentum_sgd_simple():