wehub-resource-sync
723 lines
27 KiB
Python
723 lines
27 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.
|
|
# ruff: noqa: F401
|
|
"""Provide abstraction for defining optimizers and a set of common optimizers."""
|
|
|
|
from decimal import Decimal
|
|
from typing import Optional, Union
|
|
|
|
import numpy as np # type: ignore
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
|
|
from ..block_builder import BlockBuilder
|
|
from ..expr import Function, TupleGetItem, Var, const
|
|
from ..expr import Tuple as RxTuple
|
|
from ..op import add, divide, multiply, sqrt, subtract
|
|
from ..type import TensorType, TupleType
|
|
|
|
|
|
# TODO(chaofan, yixin): Migrate key logics to C++
|
|
class Optimizer:
|
|
"""Relax training optimizer. This class could generate relax Functions for optimizing specified
|
|
parameters, and store the states used in the optimization process, such as momentum.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the optimizer function. This parameter is provided by subclasses.
|
|
|
|
Attributes
|
|
----------
|
|
dtype : str
|
|
The only dtype of the optimizer. It will be used as the dtype of the optimizer states,
|
|
and the dtype of necessary constants, such as the learning rate. Will be set in `init()`.
|
|
|
|
name : str
|
|
The name of the optimizer.
|
|
|
|
param_list : List[Var]
|
|
The list of variables to optimize. Will be set in `init()`.
|
|
|
|
state : tvm_ffi.Array
|
|
`state` is an runtime Array representing the state of the optimizer. Will be set in
|
|
`init()`.
|
|
|
|
The states of the optimizer can store necessary information in the optimization process at
|
|
runtime, such as the number of steps, the momentum in momentum SGD, etc.
|
|
|
|
`opt.state` should be used as the last argument of the function that is got through
|
|
`get_function()`, and its new value is returned as the last return value of that function.
|
|
|
|
See examples for more details.
|
|
|
|
Examples
|
|
--------
|
|
The usage of optimizers should resemble the following pattern. We will take SGD as an example.
|
|
For detailed examples, please see the tutorial.
|
|
|
|
.. code-block:: python
|
|
|
|
# Construct the optimizer
|
|
opt = relax.optimizer.SGD(0.1)
|
|
|
|
# Initialize the parameter list, the dtype and the optimizer state
|
|
# x is the relax Var we want to optimize
|
|
opt.init(x)
|
|
|
|
# The above two lines is equivalent to one line:
|
|
opt = relax.optimizer.SGD(0.1).init(x)
|
|
|
|
# Get the optimizer function
|
|
# mod is an IRModule constructed earlier
|
|
mod["SGD"] = opt.get_function()
|
|
|
|
# Legalize and build mod
|
|
lowered_mod = LegalizeOps()(mod)
|
|
ex = build(lowered_mod, target="llvm")
|
|
vm = VirtualMachine(ex, tvm.cpu())
|
|
|
|
# Optimization process
|
|
# param_tuple is a runtime tuple of parameters
|
|
# param_gradient is a runtime tuple of the gradient of the parameters in param_tuple,
|
|
# respectively
|
|
# param_gradient can be gained by the automatic differentiation pass. Please see
|
|
# `relax.transform.Gradient`
|
|
param_tuple, opt.state = vm["SGD"](param_tuple, param_gradient, opt.state)
|
|
"""
|
|
|
|
dtype: str
|
|
name: str
|
|
param_list: list[Var]
|
|
state: tvm_ffi.Array
|
|
|
|
def __init__(self, name: str) -> None:
|
|
self.name = name
|
|
self.param_list = None
|
|
self.state = None
|
|
self.dtype = None
|
|
|
|
def init(self, params: Var | list[Var]) -> "Optimizer":
|
|
"""Set the parameters, determine the dtype, and construct the initial state for the
|
|
optimizer.
|
|
|
|
Parameters
|
|
----------
|
|
params : Union[Var, List[Var]]
|
|
The parameter or the list of parameters to optimize.
|
|
|
|
Parameters should all be Vars of floating point Tensors, including float32, float64,
|
|
float16, etc. Currently, all parameters should have the same dtype, and that dtype
|
|
will be used as the dtype of the optimizer states.
|
|
|
|
Returns
|
|
-------
|
|
self : Optimizer
|
|
The optimizer itself.
|
|
"""
|
|
if not isinstance(params, list):
|
|
params = [params]
|
|
self._set_params_and_dtype(params)
|
|
# State should be initialized in any implementation of optimizer.
|
|
self.state = None
|
|
return self
|
|
|
|
def _set_params_and_dtype(self, params: list[Var]) -> None:
|
|
"""Check params is legal and set the param_list and dtype of the optimizer."""
|
|
params_set = set()
|
|
dtype = None
|
|
for x in params:
|
|
if not isinstance(x, Var):
|
|
raise ValueError(f"Parameter {x} is not a Var")
|
|
if not isinstance(x.ty, TensorType):
|
|
raise ValueError(
|
|
f"Optimizers only support Tensor parameters, but parameter {x.name_hint} has "
|
|
f"type {x.ty}"
|
|
)
|
|
data_type = tvm.DataType(x.ty.dtype.dtype)
|
|
if data_type.type_code not in (tvm.DataTypeCode.BFLOAT, tvm.DataTypeCode.FLOAT):
|
|
raise ValueError(
|
|
f"Optimizers only support Tensor parameters of floating point dtype, but dtype "
|
|
f"of {x.name_hint} is {x.ty.dtype}"
|
|
)
|
|
if dtype is None:
|
|
dtype = x.ty.dtype
|
|
else:
|
|
if dtype != x.ty.dtype:
|
|
raise ValueError(
|
|
f"All parameters should have the same dtype, but parameter {x.name_hint} "
|
|
f"has dtype {x.ty.dtype}, which differs from the previous dtype "
|
|
f"{dtype}"
|
|
)
|
|
if x in params_set:
|
|
raise ValueError(f"Parameter {x.name_hint} appears more than once")
|
|
params_set.add(x)
|
|
self.param_list = params
|
|
self.dtype = dtype
|
|
|
|
def _check_init(self):
|
|
"""Check that the optimizer is initialized. This method should be called at the start of
|
|
get_function().
|
|
"""
|
|
if self.param_list is None or self.state is None or self.dtype is None:
|
|
raise RuntimeError("Please call init() for the optimizer before calling get_function()")
|
|
|
|
def get_function(self) -> Function:
|
|
"""Use blockbuilder to construct an optimizer function that executes updates of the
|
|
parameters and the optimizer state.
|
|
|
|
The optimizer function will take in a tuple of parameters, a tuple of gradients of
|
|
parameters, and a tuple of optimizer states. It will return a tuple of updated parameters,
|
|
and a tuple of optimizer states.
|
|
|
|
Returns
|
|
-------
|
|
func : Function
|
|
The optimizer function.
|
|
|
|
Examples
|
|
--------
|
|
An example of the returned optimizer function. This function executes the stochastic
|
|
gradient descent method with lr = 0.1.
|
|
|
|
.. code-block:: python
|
|
|
|
@R.function
|
|
def SGD(
|
|
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")),
|
|
):
|
|
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)
|
|
"""
|
|
self._check_init()
|
|
raise NotImplementedError()
|
|
|
|
|
|
# TODO(chaofan, yixin): Support symbolic shapes
|
|
def _get_shape_as_int_list(var: Var) -> list[int]:
|
|
return [int(val) for val in var.ty.shape]
|
|
|
|
|
|
# We need to subtract on hyperparameters, but do not want to introduce floating point error.
|
|
# Floating point error would lead to a few problems, such as making assert_structural_equal not
|
|
# pass in unit tests
|
|
def _high_precision_subtract(lhs: float, rhs: float) -> float:
|
|
return float(Decimal(str(lhs)) - Decimal(str(rhs)))
|
|
|
|
|
|
class SGD(Optimizer):
|
|
"""Implements stochastic gradient descent.
|
|
|
|
The returned function of `get_function()` is equivalent to the following numpy code:
|
|
|
|
.. code-block:: python
|
|
|
|
def SGD(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
|
|
|
|
Parameters
|
|
----------
|
|
lr : float
|
|
learning rate
|
|
|
|
weight_decay : float
|
|
weight decay (L2 penalty) (default: 0)
|
|
"""
|
|
|
|
def __init__(self, lr: float, weight_decay: float = 0) -> None:
|
|
super().__init__("SGD")
|
|
self.lr = float(lr)
|
|
self.weight_decay = float(weight_decay)
|
|
|
|
def init(self, params: Var | list[Var]) -> "SGD":
|
|
"""Set the parameters, determine the dtype, and construct the initial state for the
|
|
optimizer.
|
|
|
|
The state of SGD is `(num_steps,)`.
|
|
|
|
Parameters
|
|
----------
|
|
params : Union[Var, List[Var]]
|
|
The parameter or the list of parameters to optimize.
|
|
|
|
Parameters should all be Vars of floating point Tensors, including float32, float64,
|
|
float16, etc. Currently, all parameters should have the same dtype, and that dtype
|
|
will be used as the dtype of the optimizer states.
|
|
|
|
Returns
|
|
-------
|
|
self : SGD
|
|
The SGD optimizer itself.
|
|
"""
|
|
if not isinstance(params, list):
|
|
params = [params]
|
|
self._set_params_and_dtype(params)
|
|
self.state = (
|
|
# num_steps = 0
|
|
tvm.runtime.tensor(np.zeros((), "int64")),
|
|
)
|
|
return self
|
|
|
|
def get_function(self) -> Function:
|
|
"""Use blockbuilder to construct an optimizer function that executes updates of the
|
|
parameters and the optimizer state. `init()` should be called before `get_function()`.
|
|
|
|
Returns
|
|
-------
|
|
func : Function
|
|
The optimizer function.
|
|
"""
|
|
self._check_init()
|
|
|
|
plist = self.param_list
|
|
len_param = len(plist)
|
|
dtype = self.dtype
|
|
|
|
# input variables
|
|
param_var = Var("params", TupleType([p.ty for p in plist]))
|
|
grad_var = Var("gradients", TupleType([p.ty for p in plist]))
|
|
state_var = Var("optim_states", TupleType([TensorType((), "int64")]))
|
|
|
|
# constants
|
|
lr = const(self.lr, dtype)
|
|
weight_decay = const(self.weight_decay, dtype)
|
|
one = const(1, "int64")
|
|
|
|
builder = BlockBuilder()
|
|
with builder.function(self.name, [param_var, grad_var, state_var]):
|
|
with builder.dataflow():
|
|
param_list_new, state_list_new = [], []
|
|
|
|
# handle num_steps
|
|
num_steps = builder.emit(TupleGetItem(state_var, 0), "num_steps")
|
|
num_steps_new = builder.emit(add(num_steps, one), "num_steps_new")
|
|
state_list_new.append(num_steps_new)
|
|
|
|
# computation logics
|
|
for i in range(len_param):
|
|
name = self.param_list[i].name_hint
|
|
p = builder.emit(TupleGetItem(param_var, i), name)
|
|
g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
|
|
if self.weight_decay:
|
|
g = builder.emit(add(multiply(weight_decay, p), g), name + "_grad_new")
|
|
p_new = builder.emit(subtract(p, multiply(lr, g)), name + "_new")
|
|
param_list_new.append(p_new)
|
|
|
|
# handle return values
|
|
params_new = builder.emit_output(RxTuple(param_list_new), "params_new")
|
|
optim_states_new = builder.emit_output(RxTuple(state_list_new), "optim_states_new")
|
|
builder.emit_func_output((params_new, optim_states_new))
|
|
return builder.get()[self.name]
|
|
|
|
|
|
class MomentumSGD(Optimizer):
|
|
"""Implements stochastic gradient descent with momentum. Optionally supports Nesterov
|
|
momentum.
|
|
|
|
The returned function of `get_function()` is equivalent to the following numpy code:
|
|
|
|
.. code-block:: python
|
|
|
|
def MomentumSGD(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
|
|
|
|
Parameters
|
|
----------
|
|
lr : float
|
|
learning rate
|
|
|
|
momentum : float
|
|
momentum factor (default: 0)
|
|
|
|
weight_decay : float
|
|
weight decay (L2 penalty) (default: 0)
|
|
|
|
dampening : float
|
|
dampening for momentum (default: 0)
|
|
|
|
nesterov : bool
|
|
enables Nesterov momentum (default: False)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
lr: float,
|
|
momentum: float,
|
|
dampening: float = 0,
|
|
weight_decay: float = 0,
|
|
nesterov: bool = False,
|
|
) -> None:
|
|
super().__init__("MomentumSGD")
|
|
self.lr = float(lr)
|
|
self.momentum = float(momentum)
|
|
self.weight_decay = float(weight_decay)
|
|
self.dampening = float(dampening)
|
|
self.nesterov = nesterov
|
|
|
|
def init(self, params: Var | list[Var]) -> "MomentumSGD":
|
|
"""Set the parameters, determine the dtype, and construct the initial state for the
|
|
optimizer.
|
|
|
|
The state of MomentumSGD is
|
|
`(num_steps, velocity_of_param_0, ..., velocity_of_param_n-1)`.
|
|
|
|
Parameters
|
|
----------
|
|
params : Union[Var, List[Var]]
|
|
The parameter or the list of parameters to optimize.
|
|
|
|
Parameters should all be Vars of floating point Tensors, including float32, float64,
|
|
float16, etc. Currently, all parameters should have the same dtype, and that dtype
|
|
will be used as the dtype of the optimizer states.
|
|
|
|
Returns
|
|
-------
|
|
self : MomentumSGD
|
|
The MomentumSGD optimizer itself.
|
|
"""
|
|
if not isinstance(params, list):
|
|
params = [params]
|
|
self._set_params_and_dtype(params)
|
|
self.state = (
|
|
# num_steps = 0
|
|
tvm.runtime.tensor(np.zeros((), "int64")),
|
|
# v_{param} is initialized to all zeros
|
|
*(
|
|
tvm.runtime.tensor(np.zeros(_get_shape_as_int_list(p), p.ty.dtype.dtype))
|
|
for p in self.param_list
|
|
),
|
|
)
|
|
return self
|
|
|
|
def get_function(self) -> Function:
|
|
"""Use blockbuilder to construct an optimizer function that executes updates of the
|
|
parameters and the optimizer state. `init()` should be called before `get_function()`.
|
|
|
|
Returns
|
|
-------
|
|
func : Function
|
|
The optimizer function.
|
|
"""
|
|
self._check_init()
|
|
plist = self.param_list
|
|
len_param = len(plist)
|
|
dtype = self.dtype
|
|
|
|
# input variables
|
|
param_var = Var("params", TupleType([p.ty for p in plist]))
|
|
grad_var = Var("gradients", TupleType([p.ty for p in plist]))
|
|
state_var = Var(
|
|
"optim_states",
|
|
TupleType([TensorType((), "int64"), *(p.ty for p in plist)]),
|
|
)
|
|
|
|
# constants
|
|
lr = const(self.lr, dtype)
|
|
momentum = const(self.momentum, dtype)
|
|
weight_decay = const(self.weight_decay, dtype)
|
|
dampening_inv = const(_high_precision_subtract(1, self.dampening), dtype)
|
|
one = const(1, "int64")
|
|
|
|
builder = BlockBuilder()
|
|
with builder.function(self.name, [param_var, grad_var, state_var]):
|
|
with builder.dataflow():
|
|
param_list_new, state_list_new = [], []
|
|
|
|
# handle num_steps
|
|
num_steps = builder.emit(TupleGetItem(state_var, 0), "num_steps")
|
|
num_steps_new = builder.emit(add(num_steps, one), "num_steps_new")
|
|
state_list_new.append(num_steps_new)
|
|
|
|
# computation logics
|
|
for i in range(len_param):
|
|
name = self.param_list[i].name_hint
|
|
p = builder.emit(TupleGetItem(param_var, i), name)
|
|
g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
|
|
v = builder.emit(TupleGetItem(state_var, i + 1), name + "_v")
|
|
if self.weight_decay:
|
|
g = builder.emit(add(multiply(weight_decay, p), g), name + "_grad_new")
|
|
damp_g = multiply(dampening_inv, g) if self.dampening else g
|
|
v_new = builder.emit(add(multiply(momentum, v), damp_g), name + "_v_new")
|
|
g_new = (
|
|
builder.emit(add(g, multiply(momentum, v_new)), name + "_g_nest")
|
|
if self.nesterov
|
|
else v_new
|
|
)
|
|
p_new = builder.emit(subtract(p, multiply(lr, g_new)), name + "_new")
|
|
param_list_new.append(p_new)
|
|
state_list_new.append(v_new)
|
|
|
|
# handle return values
|
|
params_new = builder.emit_output(RxTuple(param_list_new), "params_new")
|
|
optim_states_new = builder.emit_output(RxTuple(state_list_new), "optim_states_new")
|
|
builder.emit_func_output((params_new, optim_states_new))
|
|
return builder.get()[self.name]
|
|
|
|
|
|
class Adam(Optimizer):
|
|
"""Implements Adam optimization algorithm.
|
|
|
|
The returned function of `get_function()` is equivalent to the following numpy code:
|
|
|
|
.. code-block:: python
|
|
|
|
def Adam(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)
|
|
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
|
|
|
|
Parameters
|
|
----------
|
|
lr : float
|
|
learning rate
|
|
|
|
betas : Tuple[float, float]
|
|
coefficients used for computing running averages of gradient and its square
|
|
(default: (0.9, 0.999))
|
|
|
|
eps : float
|
|
term added to the denominator to improve numerical stability (default: 1e-8)
|
|
|
|
weight_decay : float
|
|
weight decay (L2 penalty) (default: 0)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
lr: float,
|
|
betas: tuple[float, float] = (0.9, 0.999),
|
|
eps: float = 1e-08,
|
|
weight_decay: float = 0,
|
|
) -> None:
|
|
super().__init__("Adam")
|
|
self.lr = float(lr)
|
|
self.beta1 = float(betas[0])
|
|
self.beta2 = float(betas[1])
|
|
self.eps = float(eps)
|
|
self.weight_decay = float(weight_decay)
|
|
|
|
def init(self, params: Var | list[Var]) -> "Adam":
|
|
"""Set the parameters, determine the dtype, and construct the initial state for the
|
|
optimizer.
|
|
|
|
The state of Adam is
|
|
|
|
.. code-block:: python
|
|
|
|
(
|
|
num_steps,
|
|
beta_0_prod, # beta0 ** num_steps
|
|
beta_1_prod, # beta1 ** num_steps
|
|
first_momentum_of_param_0, ..., first_momentum_of_param_n-1,
|
|
second_momentum_of_param_0, ..., second_momentum_of_param_n-1
|
|
)
|
|
|
|
Parameters
|
|
----------
|
|
params : Union[Var, List[Var]]
|
|
The parameter or the list of parameters to optimize.
|
|
|
|
Parameters should all be Vars of floating point Tensors, including float32, float64,
|
|
float16, etc. Currently, all parameters should have the same dtype, and that dtype
|
|
will be used as the dtype of the optimizer states.
|
|
|
|
Returns
|
|
-------
|
|
self : Adam
|
|
The Adam optimizer itself.
|
|
"""
|
|
if not isinstance(params, list):
|
|
params = [params]
|
|
self._set_params_and_dtype(params)
|
|
self.state = (
|
|
# num_steps, beta_0_prod, beta_1_prod
|
|
tvm.runtime.tensor(np.zeros((), "int64")),
|
|
tvm.runtime.tensor(np.ones((), self.dtype)),
|
|
tvm.runtime.tensor(np.ones((), self.dtype)),
|
|
# first_momentum
|
|
*(
|
|
tvm.runtime.tensor(np.zeros(_get_shape_as_int_list(p), p.ty.dtype.dtype))
|
|
for p in self.param_list
|
|
),
|
|
# second_momentum
|
|
*(
|
|
tvm.runtime.tensor(np.zeros(_get_shape_as_int_list(p), p.ty.dtype.dtype))
|
|
for p in self.param_list
|
|
),
|
|
)
|
|
return self
|
|
|
|
def get_function(self) -> Function:
|
|
"""Use blockbuilder to construct an optimizer function that executes updates of the
|
|
parameters and the optimizer state. `init()` should be called before `get_function()`.
|
|
|
|
Returns
|
|
-------
|
|
func : Function
|
|
The optimizer function.
|
|
"""
|
|
self._check_init()
|
|
plist = self.param_list
|
|
len_param = len(plist)
|
|
dtype = self.dtype
|
|
|
|
# input variables
|
|
param_var = Var("params", TupleType([p.ty for p in plist]))
|
|
grad_var = Var("gradients", TupleType([p.ty for p in plist]))
|
|
state_var = Var(
|
|
"optim_states",
|
|
TupleType(
|
|
[
|
|
TensorType((), "int64"),
|
|
TensorType((), dtype),
|
|
TensorType((), dtype),
|
|
*(p.ty for p in plist),
|
|
*(p.ty for p in plist),
|
|
]
|
|
),
|
|
)
|
|
|
|
# constants
|
|
lr = const(self.lr, dtype)
|
|
beta1 = const(self.beta1, dtype)
|
|
beta2 = const(self.beta2, dtype)
|
|
beta1_inv = const(_high_precision_subtract(1, self.beta1), dtype)
|
|
beta2_inv = const(_high_precision_subtract(1, self.beta2), dtype)
|
|
eps = const(self.eps, dtype)
|
|
weight_decay = const(self.weight_decay, dtype)
|
|
one_int = const(1, "int64")
|
|
one_float = const(1, dtype)
|
|
|
|
builder = BlockBuilder()
|
|
with builder.function(self.name, [param_var, grad_var, state_var]):
|
|
with builder.dataflow():
|
|
param_list_new = []
|
|
state_list_new = [None] * (len_param * 2 + 3) # type: List[Optional[Var]]
|
|
|
|
# handle num_steps
|
|
num_steps = builder.emit(TupleGetItem(state_var, 0), "num_steps")
|
|
num_steps_new = builder.emit(add(num_steps, one_int), "num_steps_new")
|
|
state_list_new[0] = num_steps_new
|
|
beta1_prod = builder.emit(multiply(TupleGetItem(state_var, 1), beta1), "beta1_prod")
|
|
beta2_prod = builder.emit(multiply(TupleGetItem(state_var, 2), beta2), "beta2_prod")
|
|
state_list_new[1] = beta1_prod
|
|
state_list_new[2] = beta2_prod
|
|
|
|
# computation logics
|
|
for i in range(len_param):
|
|
name = self.param_list[i].name_hint
|
|
p = builder.emit(TupleGetItem(param_var, i), name)
|
|
g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
|
|
m = builder.emit(TupleGetItem(state_var, i + 3), name + "_m")
|
|
v = builder.emit(TupleGetItem(state_var, i + 3 + len_param), name + "_v")
|
|
if self.weight_decay:
|
|
g = builder.emit(add(multiply(weight_decay, p), g), name + "_grad_new")
|
|
m_new = builder.emit(
|
|
add(multiply(beta1, m), multiply(beta1_inv, g)), name + "_m_new"
|
|
)
|
|
v_new = builder.emit(
|
|
add(multiply(beta2, v), multiply(beta2_inv, multiply(g, g))),
|
|
name + "_v_new",
|
|
)
|
|
m_hat = builder.emit(
|
|
divide(m_new, subtract(one_float, state_list_new[1])), name + "_m_hat"
|
|
)
|
|
v_hat = builder.emit(
|
|
divide(v_new, subtract(one_float, state_list_new[2])), name + "_v_hat"
|
|
)
|
|
p_new = builder.emit(
|
|
subtract(p, multiply(lr, divide(m_hat, add(sqrt(v_hat), eps)))),
|
|
name + "_new",
|
|
)
|
|
param_list_new.append(p_new)
|
|
state_list_new[i + 3] = m_new
|
|
state_list_new[i + 3 + len_param] = v_new
|
|
|
|
# handle return values
|
|
params_new = builder.emit_output(RxTuple(param_list_new), "params_new")
|
|
optim_states_new = builder.emit_output(RxTuple(state_list_new), "optim_states_new")
|
|
builder.emit_func_output((params_new, optim_states_new))
|
|
return builder.get()[self.name]
|