wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
26446540fadaf1a5bdd5dfc29e594b312da09221
apache--tvm/python/tvm/relax/training/optimizer.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

723 lines
27 KiB
Python
Raw Blame History

# 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]
Reference in New Issue View Git Blame Copy Permalink