chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,604 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
"""Unit tests for relax optimizer APIs."""
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.ir.base import assert_structural_equal
|
||||
from tvm.relax.training.optimizer import SGD, Adam, MomentumSGD
|
||||
from tvm.script.parser import relax as R
|
||||
|
||||
|
||||
def test_optimizer_error():
|
||||
x1 = relax.Var("x1", R.Tensor((3, 3), "float32"))
|
||||
x2 = relax.Var("x2", R.Tensor((3, 3), "float64"))
|
||||
x3 = relax.Var("x3", R.Tuple([R.Tensor((3, 3), "float32")]))
|
||||
x4 = relax.Var("x4", R.Tensor((3, 3), "int64"))
|
||||
x5 = relax.Tuple([x1])
|
||||
|
||||
# fine cases
|
||||
SGD(0.01).init(x1)
|
||||
SGD(0.01).init([x1])
|
||||
assert SGD(0.01).init([x2]).dtype == "float64"
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
SGD(0.01).init([x1, x1])
|
||||
with pytest.raises(ValueError):
|
||||
SGD(0.01).init([x1, x2])
|
||||
with pytest.raises(ValueError):
|
||||
SGD(0.01).init(x3)
|
||||
with pytest.raises(ValueError):
|
||||
SGD(0.01).init(x4)
|
||||
with pytest.raises(ValueError):
|
||||
SGD(0.01).init(x5)
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match="Please call init\\(\\) for the optimizer before calling get_function\\(\\)",
|
||||
):
|
||||
SGD(0.01).get_function()
|
||||
|
||||
|
||||
def test_sgd_simple():
|
||||
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
||||
y = relax.Var("y", R.Tensor((3,), "float32"))
|
||||
sgd = SGD(0.01).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"})
|
||||
# block 0
|
||||
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.01, "float32"), x_grad)
|
||||
x_new: R.Tensor((3, 3), "float32") = R.subtract(x, lv)
|
||||
y: R.Tensor((3,), "float32") = params[1]
|
||||
y_grad: R.Tensor((3,), "float32") = gradients[1]
|
||||
lv1: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), y_grad)
|
||||
y_new: R.Tensor((3,), "float32") = R.subtract(y, lv1)
|
||||
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_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():
|
||||
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
||||
y = relax.Var("y", R.Tensor((3,), "float32"))
|
||||
msgd = MomentumSGD(0.01, 0.9).init([x, y]).get_function()
|
||||
|
||||
@R.function
|
||||
def msgd_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.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")
|
||||
),
|
||||
) -> R.Tuple(
|
||||
R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
R.Tuple(R.Tensor((), "int64"), R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
):
|
||||
R.func_attr({"global_symbol": "MomentumSGD"})
|
||||
# block 0
|
||||
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]
|
||||
x_v: R.Tensor((3, 3), "float32") = optim_states[1]
|
||||
lv: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.9, "float32"), x_v)
|
||||
x_v_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_v_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]
|
||||
y_v: R.Tensor((3,), "float32") = optim_states[2]
|
||||
lv2: R.Tensor((3,), "float32") = R.multiply(R.const(0.9, "float32"), y_v)
|
||||
y_v_new: R.Tensor((3,), "float32") = R.add(lv2, y_grad)
|
||||
lv3: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), y_v_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"), R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")
|
||||
) = (num_steps_new, x_v_new, y_v_new)
|
||||
R.output(params_new, optim_states_new)
|
||||
return (params_new, optim_states_new)
|
||||
|
||||
assert_structural_equal(msgd, msgd_expected)
|
||||
|
||||
|
||||
def test_momentum_sgd_complex():
|
||||
lr, mom, damp, wd, nest = 0.01, 0.9, 0.85, 0.02, False
|
||||
|
||||
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
||||
y = relax.Var("y", R.Tensor((3,), "float32"))
|
||||
msgd = MomentumSGD(lr, mom, damp, wd, nest).init([x, y]).get_function()
|
||||
|
||||
@R.function
|
||||
def msgd_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.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")
|
||||
),
|
||||
) -> R.Tuple(
|
||||
R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
R.Tuple(R.Tensor((), "int64"), R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
):
|
||||
R.func_attr({"global_symbol": "MomentumSGD"})
|
||||
# block 0
|
||||
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]
|
||||
x_v: R.Tensor((3, 3), "float32") = optim_states[1]
|
||||
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.9, "float32"), x_v)
|
||||
lv2: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.15, "float32"), x_grad_new)
|
||||
x_v_new: R.Tensor((3, 3), "float32") = R.add(lv1, lv2)
|
||||
lv3: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.01, "float32"), x_v_new)
|
||||
x_new: R.Tensor((3, 3), "float32") = R.subtract(x, lv3)
|
||||
y: R.Tensor((3,), "float32") = params[1]
|
||||
y_grad: R.Tensor((3,), "float32") = gradients[1]
|
||||
y_v: R.Tensor((3,), "float32") = optim_states[2]
|
||||
lv4: R.Tensor((3,), "float32") = R.multiply(R.const(0.02, "float32"), y)
|
||||
y_grad_new: R.Tensor((3,), "float32") = R.add(lv4, y_grad)
|
||||
lv5: R.Tensor((3,), "float32") = R.multiply(R.const(0.9, "float32"), y_v)
|
||||
lv6: R.Tensor((3,), "float32") = R.multiply(R.const(0.15, "float32"), y_grad_new)
|
||||
y_v_new: R.Tensor((3,), "float32") = R.add(lv5, lv6)
|
||||
lv7: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), y_v_new)
|
||||
y_new: R.Tensor((3,), "float32") = R.subtract(y, lv7)
|
||||
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"), R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")
|
||||
) = (num_steps_new, x_v_new, y_v_new)
|
||||
R.output(params_new, optim_states_new)
|
||||
return (params_new, optim_states_new)
|
||||
|
||||
assert_structural_equal(msgd, msgd_expected)
|
||||
|
||||
|
||||
def test_momentum_sgd_nesterov():
|
||||
lr, mom, damp, wd, nest = 0.01, 0.9, 0.85, 0.02, True
|
||||
|
||||
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
||||
y = relax.Var("y", R.Tensor((3,), "float32"))
|
||||
msgd = MomentumSGD(lr, mom, damp, wd, nest).init([x, y]).get_function()
|
||||
|
||||
@R.function
|
||||
def msgd_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.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")
|
||||
),
|
||||
) -> R.Tuple(
|
||||
R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
R.Tuple(R.Tensor((), "int64"), R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
):
|
||||
R.func_attr({"global_symbol": "MomentumSGD"})
|
||||
# block 0
|
||||
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]
|
||||
x_v: R.Tensor((3, 3), "float32") = optim_states[1]
|
||||
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.9, "float32"), x_v)
|
||||
lv2: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.15, "float32"), x_grad_new)
|
||||
x_v_new: R.Tensor((3, 3), "float32") = R.add(lv1, lv2)
|
||||
lv3: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.9, "float32"), x_v_new)
|
||||
x_g_nest: R.Tensor((3, 3), "float32") = R.add(x_grad_new, lv3)
|
||||
lv4: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.01, "float32"), x_g_nest)
|
||||
x_new: R.Tensor((3, 3), "float32") = R.subtract(x, lv4)
|
||||
y: R.Tensor((3,), "float32") = params[1]
|
||||
y_grad: R.Tensor((3,), "float32") = gradients[1]
|
||||
y_v: R.Tensor((3,), "float32") = optim_states[2]
|
||||
lv5: R.Tensor((3,), "float32") = R.multiply(R.const(0.02, "float32"), y)
|
||||
y_grad_new: R.Tensor((3,), "float32") = R.add(lv5, y_grad)
|
||||
lv6: R.Tensor((3,), "float32") = R.multiply(R.const(0.9, "float32"), y_v)
|
||||
lv7: R.Tensor((3,), "float32") = R.multiply(R.const(0.15, "float32"), y_grad_new)
|
||||
y_v_new: R.Tensor((3,), "float32") = R.add(lv6, lv7)
|
||||
lv8: R.Tensor((3,), "float32") = R.multiply(R.const(0.9, "float32"), y_v_new)
|
||||
y_g_nest: R.Tensor((3,), "float32") = R.add(y_grad_new, lv8)
|
||||
lv9: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), y_g_nest)
|
||||
y_new: R.Tensor((3,), "float32") = R.subtract(y, lv9)
|
||||
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"), R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")
|
||||
) = (num_steps_new, x_v_new, y_v_new)
|
||||
R.output(params_new, optim_states_new)
|
||||
return (params_new, optim_states_new)
|
||||
|
||||
assert_structural_equal(msgd, msgd_expected)
|
||||
|
||||
|
||||
def test_adam_simple():
|
||||
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
||||
y = relax.Var("y", R.Tensor((3,), "float32"))
|
||||
adam = Adam(0.01).init([x, y]).get_function()
|
||||
|
||||
@R.function
|
||||
def adam_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.Tensor((), "float32"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
),
|
||||
) -> R.Tuple(
|
||||
R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
R.Tuple(
|
||||
R.Tensor((), "int64"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
),
|
||||
):
|
||||
R.func_attr({"global_symbol": "Adam"})
|
||||
# block 0
|
||||
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"))
|
||||
lv: R.Tensor((), "float32") = optim_states[1]
|
||||
beta1_prod: R.Tensor((), "float32") = R.multiply(lv, R.const(0.9, "float32"))
|
||||
lv1: R.Tensor((), "float32") = optim_states[2]
|
||||
beta2_prod: R.Tensor((), "float32") = R.multiply(lv1, R.const(0.999, "float32"))
|
||||
x: R.Tensor((3, 3), "float32") = params[0]
|
||||
x_grad: R.Tensor((3, 3), "float32") = gradients[0]
|
||||
x_m: R.Tensor((3, 3), "float32") = optim_states[3]
|
||||
x_v: R.Tensor((3, 3), "float32") = optim_states[5]
|
||||
lv2: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.9, "float32"), x_m)
|
||||
lv3: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.1, "float32"), x_grad)
|
||||
x_m_new: R.Tensor((3, 3), "float32") = R.add(lv2, lv3)
|
||||
lv4: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.999, "float32"), x_v)
|
||||
lv5: R.Tensor((3, 3), "float32") = R.multiply(x_grad, x_grad)
|
||||
lv6: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.001, "float32"), lv5)
|
||||
x_v_new: R.Tensor((3, 3), "float32") = R.add(lv4, lv6)
|
||||
lv7: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta1_prod)
|
||||
x_m_hat: R.Tensor((3, 3), "float32") = R.divide(x_m_new, lv7)
|
||||
lv8: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta2_prod)
|
||||
x_v_hat: R.Tensor((3, 3), "float32") = R.divide(x_v_new, lv8)
|
||||
lv9: R.Tensor((3, 3), "float32") = R.sqrt(x_v_hat)
|
||||
lv10: R.Tensor((3, 3), "float32") = R.add(lv9, R.const(1e-08, "float32"))
|
||||
lv11: R.Tensor((3, 3), "float32") = R.divide(x_m_hat, lv10)
|
||||
lv12: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.01, "float32"), lv11)
|
||||
x_new: R.Tensor((3, 3), "float32") = R.subtract(x, lv12)
|
||||
y: R.Tensor((3,), "float32") = params[1]
|
||||
y_grad: R.Tensor((3,), "float32") = gradients[1]
|
||||
y_m: R.Tensor((3,), "float32") = optim_states[4]
|
||||
y_v: R.Tensor((3,), "float32") = optim_states[6]
|
||||
lv13: R.Tensor((3,), "float32") = R.multiply(R.const(0.9, "float32"), y_m)
|
||||
lv14: R.Tensor((3,), "float32") = R.multiply(R.const(0.1, "float32"), y_grad)
|
||||
y_m_new: R.Tensor((3,), "float32") = R.add(lv13, lv14)
|
||||
lv15: R.Tensor((3,), "float32") = R.multiply(R.const(0.999, "float32"), y_v)
|
||||
lv16: R.Tensor((3,), "float32") = R.multiply(y_grad, y_grad)
|
||||
lv17: R.Tensor((3,), "float32") = R.multiply(R.const(0.001, "float32"), lv16)
|
||||
y_v_new: R.Tensor((3,), "float32") = R.add(lv15, lv17)
|
||||
lv18: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta1_prod)
|
||||
y_m_hat: R.Tensor((3,), "float32") = R.divide(y_m_new, lv18)
|
||||
lv19: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta2_prod)
|
||||
y_v_hat: R.Tensor((3,), "float32") = R.divide(y_v_new, lv19)
|
||||
lv20: R.Tensor((3,), "float32") = R.sqrt(y_v_hat)
|
||||
lv21: R.Tensor((3,), "float32") = R.add(lv20, R.const(1e-08, "float32"))
|
||||
lv22: R.Tensor((3,), "float32") = R.divide(y_m_hat, lv21)
|
||||
lv23: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), lv22)
|
||||
y_new: R.Tensor((3,), "float32") = R.subtract(y, lv23)
|
||||
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"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
) = (num_steps_new, beta1_prod, beta2_prod, x_m_new, y_m_new, x_v_new, y_v_new)
|
||||
R.output(params_new, optim_states_new)
|
||||
return (params_new, optim_states_new)
|
||||
|
||||
assert_structural_equal(adam, adam_expected)
|
||||
|
||||
|
||||
def test_adam_complex():
|
||||
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
||||
y = relax.Var("y", R.Tensor((3,), "float32"))
|
||||
adam = Adam(0.01, (0.8, 0.85), 1e-7, 0.1).init([x, y]).get_function()
|
||||
|
||||
@R.function
|
||||
def adam_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.Tensor((), "float32"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
),
|
||||
) -> R.Tuple(
|
||||
R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
|
||||
R.Tuple(
|
||||
R.Tensor((), "int64"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
),
|
||||
):
|
||||
R.func_attr({"global_symbol": "Adam"})
|
||||
# block 0
|
||||
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"))
|
||||
lv: R.Tensor((), "float32") = optim_states[1]
|
||||
beta1_prod: R.Tensor((), "float32") = R.multiply(lv, R.const(0.8, "float32"))
|
||||
lv1: R.Tensor((), "float32") = optim_states[2]
|
||||
beta2_prod: R.Tensor((), "float32") = R.multiply(lv1, R.const(0.85, "float32"))
|
||||
x: R.Tensor((3, 3), "float32") = params[0]
|
||||
x_grad: R.Tensor((3, 3), "float32") = gradients[0]
|
||||
x_m: R.Tensor((3, 3), "float32") = optim_states[3]
|
||||
x_v: R.Tensor((3, 3), "float32") = optim_states[5]
|
||||
lv2: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.1, "float32"), x)
|
||||
x_grad_new: R.Tensor((3, 3), "float32") = R.add(lv2, x_grad)
|
||||
lv3: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.8, "float32"), x_m)
|
||||
lv4: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.2, "float32"), x_grad_new)
|
||||
x_m_new: R.Tensor((3, 3), "float32") = R.add(lv3, lv4)
|
||||
lv5: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.85, "float32"), x_v)
|
||||
lv6: R.Tensor((3, 3), "float32") = R.multiply(x_grad_new, x_grad_new)
|
||||
lv7: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.15, "float32"), lv6)
|
||||
x_v_new: R.Tensor((3, 3), "float32") = R.add(lv5, lv7)
|
||||
lv8: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta1_prod)
|
||||
x_m_hat: R.Tensor((3, 3), "float32") = R.divide(x_m_new, lv8)
|
||||
lv9: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta2_prod)
|
||||
x_v_hat: R.Tensor((3, 3), "float32") = R.divide(x_v_new, lv9)
|
||||
lv10: R.Tensor((3, 3), "float32") = R.sqrt(x_v_hat)
|
||||
lv11: R.Tensor((3, 3), "float32") = R.add(lv10, R.const(1e-07, "float32"))
|
||||
lv12: R.Tensor((3, 3), "float32") = R.divide(x_m_hat, lv11)
|
||||
lv13: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.01, "float32"), lv12)
|
||||
x_new: R.Tensor((3, 3), "float32") = R.subtract(x, lv13)
|
||||
y: R.Tensor((3,), "float32") = params[1]
|
||||
y_grad: R.Tensor((3,), "float32") = gradients[1]
|
||||
y_m: R.Tensor((3,), "float32") = optim_states[4]
|
||||
y_v: R.Tensor((3,), "float32") = optim_states[6]
|
||||
lv14: R.Tensor((3,), "float32") = R.multiply(R.const(0.1, "float32"), y)
|
||||
y_grad_new: R.Tensor((3,), "float32") = R.add(lv14, y_grad)
|
||||
lv15: R.Tensor((3,), "float32") = R.multiply(R.const(0.8, "float32"), y_m)
|
||||
lv16: R.Tensor((3,), "float32") = R.multiply(R.const(0.2, "float32"), y_grad_new)
|
||||
y_m_new: R.Tensor((3,), "float32") = R.add(lv15, lv16)
|
||||
lv17: R.Tensor((3,), "float32") = R.multiply(R.const(0.85, "float32"), y_v)
|
||||
lv18: R.Tensor((3,), "float32") = R.multiply(y_grad_new, y_grad_new)
|
||||
lv19: R.Tensor((3,), "float32") = R.multiply(R.const(0.15, "float32"), lv18)
|
||||
y_v_new: R.Tensor((3,), "float32") = R.add(lv17, lv19)
|
||||
lv20: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta1_prod)
|
||||
y_m_hat: R.Tensor((3,), "float32") = R.divide(y_m_new, lv20)
|
||||
lv21: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta2_prod)
|
||||
y_v_hat: R.Tensor((3,), "float32") = R.divide(y_v_new, lv21)
|
||||
lv22: R.Tensor((3,), "float32") = R.sqrt(y_v_hat)
|
||||
lv23: R.Tensor((3,), "float32") = R.add(lv22, R.const(1e-07, "float32"))
|
||||
lv24: R.Tensor((3,), "float32") = R.divide(y_m_hat, lv23)
|
||||
lv25: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), lv24)
|
||||
y_new: R.Tensor((3,), "float32") = R.subtract(y, lv25)
|
||||
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"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
R.Tensor((3, 3), "float32"),
|
||||
R.Tensor((3,), "float32"),
|
||||
) = (num_steps_new, beta1_prod, beta2_prod, x_m_new, y_m_new, x_v_new, y_v_new)
|
||||
R.output(params_new, optim_states_new)
|
||||
return (params_new, optim_states_new)
|
||||
|
||||
assert_structural_equal(adam, adam_expected)
|
||||
|
||||
|
||||
def test_adam_float64():
|
||||
x = relax.Var("x", R.Tensor((3, 3), "float64"))
|
||||
y = relax.Var("y", R.Tensor((3,), "float64"))
|
||||
adam = Adam(0.01, (0.8, 0.85), 1e-7, 0.1).init([x, y]).get_function()
|
||||
|
||||
@R.function
|
||||
def adam_expected(
|
||||
params: R.Tuple(R.Tensor((3, 3), "float64"), R.Tensor((3,), "float64")),
|
||||
gradients: R.Tuple(R.Tensor((3, 3), "float64"), R.Tensor((3,), "float64")),
|
||||
optim_states: R.Tuple(
|
||||
R.Tensor((), "int64"),
|
||||
R.Tensor((), "float64"),
|
||||
R.Tensor((), "float64"),
|
||||
R.Tensor((3, 3), "float64"),
|
||||
R.Tensor((3,), "float64"),
|
||||
R.Tensor((3, 3), "float64"),
|
||||
R.Tensor((3,), "float64"),
|
||||
),
|
||||
) -> R.Tuple(
|
||||
R.Tuple(R.Tensor((3, 3), "float64"), R.Tensor((3,), "float64")),
|
||||
R.Tuple(
|
||||
R.Tensor((), "int64"),
|
||||
R.Tensor((), "float64"),
|
||||
R.Tensor((), "float64"),
|
||||
R.Tensor((3, 3), "float64"),
|
||||
R.Tensor((3,), "float64"),
|
||||
R.Tensor((3, 3), "float64"),
|
||||
R.Tensor((3,), "float64"),
|
||||
),
|
||||
):
|
||||
R.func_attr({"global_symbol": "Adam"})
|
||||
# block 0
|
||||
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"))
|
||||
lv: R.Tensor((), "float64") = optim_states[1]
|
||||
beta1_prod: R.Tensor((), "float64") = R.multiply(lv, R.const(0.8, "float64"))
|
||||
lv1: R.Tensor((), "float64") = optim_states[2]
|
||||
beta2_prod: R.Tensor((), "float64") = R.multiply(lv1, R.const(0.85, "float64"))
|
||||
x: R.Tensor((3, 3), "float64") = params[0]
|
||||
x_grad: R.Tensor((3, 3), "float64") = gradients[0]
|
||||
x_m: R.Tensor((3, 3), "float64") = optim_states[3]
|
||||
x_v: R.Tensor((3, 3), "float64") = optim_states[5]
|
||||
lv2: R.Tensor((3, 3), "float64") = R.multiply(R.const(0.1, "float64"), x)
|
||||
x_grad_new: R.Tensor((3, 3), "float64") = R.add(lv2, x_grad)
|
||||
lv3: R.Tensor((3, 3), "float64") = R.multiply(R.const(0.8, "float64"), x_m)
|
||||
lv4: R.Tensor((3, 3), "float64") = R.multiply(R.const(0.2, "float64"), x_grad_new)
|
||||
x_m_new: R.Tensor((3, 3), "float64") = R.add(lv3, lv4)
|
||||
lv5: R.Tensor((3, 3), "float64") = R.multiply(R.const(0.85, "float64"), x_v)
|
||||
lv6: R.Tensor((3, 3), "float64") = R.multiply(x_grad_new, x_grad_new)
|
||||
lv7: R.Tensor((3, 3), "float64") = R.multiply(R.const(0.15, "float64"), lv6)
|
||||
x_v_new: R.Tensor((3, 3), "float64") = R.add(lv5, lv7)
|
||||
lv8: R.Tensor((), "float64") = R.subtract(R.const(1, "float64"), beta1_prod)
|
||||
x_m_hat: R.Tensor((3, 3), "float64") = R.divide(x_m_new, lv8)
|
||||
lv9: R.Tensor((), "float64") = R.subtract(R.const(1, "float64"), beta2_prod)
|
||||
x_v_hat: R.Tensor((3, 3), "float64") = R.divide(x_v_new, lv9)
|
||||
lv10: R.Tensor((3, 3), "float64") = R.sqrt(x_v_hat)
|
||||
lv11: R.Tensor((3, 3), "float64") = R.add(lv10, R.const(1e-07, "float64"))
|
||||
lv12: R.Tensor((3, 3), "float64") = R.divide(x_m_hat, lv11)
|
||||
lv13: R.Tensor((3, 3), "float64") = R.multiply(R.const(0.01, "float64"), lv12)
|
||||
x_new: R.Tensor((3, 3), "float64") = R.subtract(x, lv13)
|
||||
y: R.Tensor((3,), "float64") = params[1]
|
||||
y_grad: R.Tensor((3,), "float64") = gradients[1]
|
||||
y_m: R.Tensor((3,), "float64") = optim_states[4]
|
||||
y_v: R.Tensor((3,), "float64") = optim_states[6]
|
||||
lv14: R.Tensor((3,), "float64") = R.multiply(R.const(0.1, "float64"), y)
|
||||
y_grad_new: R.Tensor((3,), "float64") = R.add(lv14, y_grad)
|
||||
lv15: R.Tensor((3,), "float64") = R.multiply(R.const(0.8, "float64"), y_m)
|
||||
lv16: R.Tensor((3,), "float64") = R.multiply(R.const(0.2, "float64"), y_grad_new)
|
||||
y_m_new: R.Tensor((3,), "float64") = R.add(lv15, lv16)
|
||||
lv17: R.Tensor((3,), "float64") = R.multiply(R.const(0.85, "float64"), y_v)
|
||||
lv18: R.Tensor((3,), "float64") = R.multiply(y_grad_new, y_grad_new)
|
||||
lv19: R.Tensor((3,), "float64") = R.multiply(R.const(0.15, "float64"), lv18)
|
||||
y_v_new: R.Tensor((3,), "float64") = R.add(lv17, lv19)
|
||||
lv20: R.Tensor((), "float64") = R.subtract(R.const(1, "float64"), beta1_prod)
|
||||
y_m_hat: R.Tensor((3,), "float64") = R.divide(y_m_new, lv20)
|
||||
lv21: R.Tensor((), "float64") = R.subtract(R.const(1, "float64"), beta2_prod)
|
||||
y_v_hat: R.Tensor((3,), "float64") = R.divide(y_v_new, lv21)
|
||||
lv22: R.Tensor((3,), "float64") = R.sqrt(y_v_hat)
|
||||
lv23: R.Tensor((3,), "float64") = R.add(lv22, R.const(1e-07, "float64"))
|
||||
lv24: R.Tensor((3,), "float64") = R.divide(y_m_hat, lv23)
|
||||
lv25: R.Tensor((3,), "float64") = R.multiply(R.const(0.01, "float64"), lv24)
|
||||
y_new: R.Tensor((3,), "float64") = R.subtract(y, lv25)
|
||||
params_new: R.Tuple(R.Tensor((3, 3), "float64"), R.Tensor((3,), "float64")) = (
|
||||
x_new,
|
||||
y_new,
|
||||
)
|
||||
optim_states_new: R.Tuple(
|
||||
R.Tensor((), "int64"),
|
||||
R.Tensor((), "float64"),
|
||||
R.Tensor((), "float64"),
|
||||
R.Tensor((3, 3), "float64"),
|
||||
R.Tensor((3,), "float64"),
|
||||
R.Tensor((3, 3), "float64"),
|
||||
R.Tensor((3,), "float64"),
|
||||
) = (num_steps_new, beta1_prod, beta2_prod, x_m_new, y_m_new, x_v_new, y_v_new)
|
||||
R.output(params_new, optim_states_new)
|
||||
return (params_new, optim_states_new)
|
||||
|
||||
assert_structural_equal(adam, adam_expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user