wehub-resource-sync
189 lines
6.3 KiB
Python
189 lines
6.3 KiB
Python
# 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.
|
|
"""Numeric tests for relax optimizer APIs."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import IRModule, relax
|
|
from tvm.relax.training.optimizer import SGD, Adam, MomentumSGD
|
|
from tvm.runtime.vm import VirtualMachine
|
|
from tvm.script.parser import relax as R
|
|
from tvm.testing import assert_allclose
|
|
|
|
|
|
def _legalize_and_build(mod: IRModule, target, dev):
|
|
ex = tvm.compile(mod, target)
|
|
vm = VirtualMachine(ex, dev)
|
|
return vm
|
|
|
|
|
|
def _numpy_to_tvm(data):
|
|
if isinstance(data, list | tuple):
|
|
return [_numpy_to_tvm(_data) for _data in data]
|
|
return tvm.runtime.tensor(data)
|
|
|
|
|
|
def _tvm_to_numpy(data):
|
|
if isinstance(data, list | tuple | tvm_ffi.Array):
|
|
return [_tvm_to_numpy(_data) for _data in data]
|
|
return data.numpy()
|
|
|
|
|
|
def _assert_allclose_nested(data1, data2):
|
|
if isinstance(data1, list | tuple):
|
|
assert isinstance(data2, list | tuple)
|
|
assert len(data1) == len(data2)
|
|
for x, y in zip(data1, data2):
|
|
_assert_allclose_nested(x, y)
|
|
else:
|
|
assert_allclose(data1, data2)
|
|
|
|
|
|
def _assert_run_result_same(tvm_func: Callable, np_func: Callable, np_inputs: list):
|
|
result = _tvm_to_numpy(tvm_func(*[_numpy_to_tvm(i) for i in np_inputs]))
|
|
expected = np_func(*np_inputs)
|
|
_assert_allclose_nested(result, expected)
|
|
|
|
|
|
def _test_optimizer(target, dev, np_func, opt_type, *args, **kwargs):
|
|
x = relax.Var("x", R.Tensor((3, 3), "float32"))
|
|
y = relax.Var("y", R.Tensor((3,), "float32"))
|
|
opt = opt_type(*args, **kwargs).init([x, y])
|
|
mod = IRModule.from_expr(opt.get_function().with_attr("global_symbol", "main"))
|
|
tvm_func = _legalize_and_build(mod, target, dev)["main"]
|
|
|
|
param_arr = [np.random.rand(3, 3).astype(np.float32), np.random.rand(3).astype(np.float32)]
|
|
grad_arr = [np.random.rand(3, 3).astype(np.float32), np.random.rand(3).astype(np.float32)]
|
|
state_arr = _tvm_to_numpy(opt.state)
|
|
|
|
_assert_run_result_same(tvm_func, np_func, [param_arr, grad_arr, state_arr])
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"lr,weight_decay",
|
|
[
|
|
(0.01, 0),
|
|
(0.01, 0.02),
|
|
],
|
|
)
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_sgd(lr, weight_decay):
|
|
target = "llvm"
|
|
dev = tvm.device(target)
|
|
|
|
def np_func(param_tuple, grad_tuple, state_tuple):
|
|
num_steps = state_tuple[0]
|
|
param_tuple_new, state_tuple_new = [], []
|
|
state_tuple_new.append(num_steps + 1)
|
|
for i in range(len(param_tuple)):
|
|
param = param_tuple[i]
|
|
grad = grad_tuple[i]
|
|
param_tuple_new.append(param - lr * (grad + weight_decay * param))
|
|
return param_tuple_new, state_tuple_new
|
|
|
|
_test_optimizer(target, dev, np_func, SGD, lr, weight_decay)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"lr,momentum,dampening,weight_decay,nesterov",
|
|
[
|
|
(0.01, 0.9, 0, 0, False),
|
|
(0.01, 0.9, 0.85, 0.02, False),
|
|
(0.01, 0.9, 0.85, 0.02, True),
|
|
],
|
|
)
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_momentum_sgd(lr, momentum, dampening, weight_decay, nesterov):
|
|
target = "llvm"
|
|
dev = tvm.device(target)
|
|
|
|
def np_func(param_tuple, grad_tuple, state_tuple):
|
|
num_steps = state_tuple[0]
|
|
param_tuple_new, state_tuple_new = [], []
|
|
state_tuple_new.append(num_steps + 1)
|
|
|
|
for i in range(len(param_tuple)):
|
|
param = param_tuple[i]
|
|
grad = grad_tuple[i]
|
|
velocity = state_tuple[i + 1]
|
|
grad = param * weight_decay + grad
|
|
velocity = momentum * velocity + grad * (1 - dampening)
|
|
if nesterov:
|
|
param = param - (grad + momentum * velocity) * lr
|
|
else:
|
|
param = param - velocity * lr
|
|
param_tuple_new.append(param)
|
|
state_tuple_new.append(velocity)
|
|
|
|
return param_tuple_new, state_tuple_new
|
|
|
|
_test_optimizer(
|
|
target, dev, np_func, MomentumSGD, lr, momentum, dampening, weight_decay, nesterov
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"lr,betas,eps,weight_decay",
|
|
[
|
|
(0.01, (0.9, 0.999), 1e-08, 0),
|
|
(0.01, (0.8, 0.85), 1e-07, 0.1),
|
|
],
|
|
)
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_adam(lr, betas, eps, weight_decay):
|
|
target = "llvm"
|
|
dev = tvm.device(target)
|
|
|
|
def np_func(param_tuple, grad_tuple, state_tuple):
|
|
num_steps = state_tuple[0]
|
|
num_steps_new = num_steps + 1
|
|
|
|
param_tuple_new = []
|
|
state_tuple_new = [None] * len(state_tuple) # type: ignore
|
|
state_tuple_new[0] = num_steps_new
|
|
state_tuple_new[1] = state_tuple[1] * betas[0]
|
|
state_tuple_new[2] = state_tuple[2] * betas[1]
|
|
|
|
for i in range(len(param_tuple)):
|
|
param = param_tuple[i]
|
|
grad = grad_tuple[i]
|
|
m = state_tuple[i + 3]
|
|
v = state_tuple[i + 3 + len(param_tuple)]
|
|
grad = grad + weight_decay * param
|
|
m = betas[0] * m + (1 - betas[0]) * grad
|
|
v = betas[1] * v + (1 - betas[1]) * grad * grad
|
|
m_hat = m / (1 - betas[0] ** num_steps_new)
|
|
v_hat = v / (1 - betas[1] ** num_steps_new)
|
|
param = param - lr * m_hat / (np.sqrt(v_hat) + eps)
|
|
param_tuple_new.append(param)
|
|
state_tuple_new[i + 3] = m
|
|
state_tuple_new[i + 3 + len(param_tuple)] = v
|
|
|
|
return param_tuple_new, state_tuple_new
|
|
|
|
_test_optimizer(target, dev, np_func, Adam, lr, betas, eps, weight_decay)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|