# 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()