wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/training/__init__.py
+27
View File
@@ -0,0 +1,27 @@
# isort: skip_file
# 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.
"""The Relax training APIs."""
from . import loss
from . import optimizer
from . import trainer
from . import utils
from .setup_trainer import SetupTrainer
from .trainer import Trainer
from .utils import AppendLoss
python/tvm/relax/training/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""FFI APIs for tvm.relax.training"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.training", __name__)
python/tvm/relax/training/loss.py
+382
View File
@@ -0,0 +1,382 @@
# 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.
# pylint: disable=redefined-builtin, invalid-name
# ruff: noqa: RUF012
"""Loss functions library for relax."""
# isort: off
from typing import Literal
# isort: on
from ..block_builder import BlockBuilder
from ..expr import Expr, Function, Type, Var
from ..op import abs, argmax, mean, multiply, reshape, subtract, sum
from ..op.nn import log_softmax, nll_loss
def _create_param_var(param: Var | Type, param_name: str) -> Var:
"""If param is a Type, create a Var with the given Type and name.
If param is a Var, create a Var with the same Type and name as the given param Var."""
if isinstance(param, Type):
param = Var(param_name, param)
if not isinstance(param, Var):
raise TypeError("The type of param should be Var or Type, but got " + type(param))
return Var(param.name_hint, param.ty)
class Loss:
r"""Base class of all loss.
Generally, loss function will take one or more **input parameters** (that is outputs of
the backbone of a model), one or more **target parameters**, and generate a scalar value
denoting the loss.
You can use `relax.transform.AppendLoss` to append the loss function to a one-dataflowblock
backbone function in a IRModule. That will generate a one-dataflowblock function accepting
instances and targets, and then returning the loss.
Most loss functions involve a reduction of losses from all instances in a batch. We use
`reduction` parameter to denote the reduction method. Possible reduction methods include
`"mean"`, `"sum"` and `"none"`.
Parameters
----------
loss_name : str
The name of the loss function. Should be provided when calling `super().__init__` in
constructor functions of subclasses.
num_backbone_outputs : int
The number of `prediction_outputs` of the backbone function, alos the number of the
backbone_prediction_outputs of the loss function. See `relax.transform.AppendLoss`.
Should be provided when calling `super().__init__` in constructor functions of subclasses.
For example, `CrossEntropyLoss` requires one backbone prediction output; `MarginRankingLoss`
requires two backbone prediction outputs.
reduction : Literal["mean", "sum", "none"]
The reduction method to apply to output. Can be "mean", "sum" or "none".
none : no reduction will be applied,
mean : the sum of the output will be divided by the batch_size,
sum : the output will be summed.
"""
_valid_reductions = ["mean", "sum", "none"]
def __init__(
self,
loss_name: str,
num_backbone_outputs: int,
reduction: Literal["mean", "sum", "none"] = "mean",
) -> None:
self._loss_name = loss_name
self._reduction = reduction
self._num_backbone_outputs = num_backbone_outputs
if self._reduction not in self._valid_reductions:
raise ValueError("Reduction can only be one of these values: ", self._valid_reductions)
@property
def num_backbone_outputs(self) -> int:
"""Get the number of number of the outputs of the backbone function."""
return self._num_backbone_outputs
def _with_reduction(self, expr: Expr) -> Expr:
"""Add a reduction to the final loss.
Parameters
----------
expr : Expr
The loss expr.
Returns
-------
ret : Expr
The reduced result.
"""
if self._reduction == "sum":
expr = sum(expr)
elif self._reduction == "mean":
expr = mean(expr)
elif self._reduction != "none":
raise ValueError("Reduction can only be one of these values: ", self._valid_reductions)
return expr
class L1Loss(Loss):
r"""Mean element-wise absolute value difference.
Parameters
----------
reduction : Literal["mean", "sum", "none"]
The reduction method to apply to output. Can be "mean", "sum" or "none".
none : no reduction will be applied,
mean : the sum of the output will be divided by the batch_size,
sum : the output will be summed.
"""
def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean") -> None:
super().__init__("l1_loss", 1, reduction)
def __call__(
self,
predictions: Var | Type,
targets: Var | Type,
) -> Function:
"""Get the relax function of L1Loss. If the parameters are
type, it will create corresponding variables.
Parameters
----------
predictions : Union[Var, Type]
The predictions of the model in the calculation of loss.
targets : Union[Var, Type]
The ground truth in the calculation of loss.
Returns
-------
The relax function of L1Loss with the loss name as its global symbol.
"""
bb = BlockBuilder()
predictions = _create_param_var(predictions, "predictions")
targets = _create_param_var(targets, "targets")
with bb.function(self._loss_name, [predictions, targets]):
with bb.dataflow():
lv = abs(subtract(predictions, targets))
loss = bb.emit_output(self._with_reduction(lv))
bb.emit_func_output(loss)
return bb.get()[self._loss_name]
class MSELoss(Loss):
r"""Measures the element-wise mean squared error.
Parameters
----------
reduction : Literal["mean", "sum", "none"]
The reduction method to apply to output. Can be "mean", "sum" or "none".
none : no reduction will be applied,
mean : the sum of the output will be divided by the batch_size,
sum : the output will be summed.
"""
def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean") -> None:
super().__init__("mse_loss", 1, reduction)
def __call__(
self,
predictions: Var | Type,
targets: Var | Type,
) -> Function:
"""Get the relax function of MSELoss. If the parameters are
type, it will create corresponding variables.
Parameters
----------
predictions : Union[Var, Type]
The predictions of the model in the calculation of loss.
targets : Union[Var, Type]
The ground truth in the calculation of loss.
Returns
-------
The relax function of MSELoss with the loss name as its global symbol.
"""
bb = BlockBuilder()
predictions = _create_param_var(predictions, "predictions")
targets = _create_param_var(targets, "targets")
with bb.function(self._loss_name, [predictions, targets]):
with bb.dataflow():
lv = subtract(predictions, targets)
lv = multiply(lv, lv)
loss = bb.emit_output(self._with_reduction(lv))
bb.emit_func_output(loss)
return bb.get()[self._loss_name]
class CrossEntropyLoss(Loss):
r"""CrossEntropyLoss. It is a combination of a log_softmax computation and a nll_loss.
Parameters
----------
reduction : Literal["mean", "sum", "none"]
The reduction method to apply to output. Can be "mean", "sum" or "none".
none : no reduction will be applied,
mean : the sum of the output will be divided by the batch_size,
sum : the output will be summed.
ignore_index : int
Specifies a target value that is ignored and does not contribute to the input gradient.
"""
ignore_index: int
def __init__(
self,
reduction: Literal["mean", "sum", "none"] = "mean",
ignore_index: int = -100,
) -> None:
super().__init__("cross_entropy_loss", 1, reduction)
self.ignore_index = ignore_index
def __call__(
self,
predictions: Var | Type,
targets: Var | Type,
weights: Var | Type | None = None,
) -> Function:
"""Get the relax function of CrossEntropyLoss. If the parameters are
type, it will create corresponding variables.
Parameters
----------
predictions : Union[Var, Type]
The predictions of the model in the calculation of loss.
targets : Union[Var, Type]
The ground truth in the calculation of loss.
weights : Optional[Union[Var, Type]]
a manual rescaling weight given to each class. It has to be a Tensor of size C.
Returns
-------
The relax function of CrossEntropyLoss with the loss name as its global symbol.
"""
bb = BlockBuilder()
predictions = _create_param_var(predictions, "predictions")
targets = _create_param_var(targets, "targets")
arg_list = [predictions, targets]
if weights:
weights = _create_param_var(weights, "weights")
arg_list.append(weights)
with bb.function(self._loss_name, arg_list):
with bb.dataflow():
logits = bb.emit(log_softmax(predictions))
loss = bb.emit_output(
nll_loss(logits, targets, weights, self._reduction, self.ignore_index)
)
bb.emit_func_output(loss)
return bb.get()[self._loss_name]
class CategoricalCrossEntropyLoss(Loss):
r"""CategoricalCrossEntropyLoss.
It is a combination of a converting one-hot target vector to a label,
a log_softmax computation and a nll_loss.
Parameters
----------
reduction : Literal["mean", "sum", "none"]
The reduction method to apply to output. Can be "mean", "sum" or "none".
none : no reduction will be applied,
mean : the sum of the output will be divided by the batch_size,
sum : the output will be summed.
ignore_index : int
Specifies a target value that is ignored and does not contribute to the input gradient.
"""
ignore_index: int
def __init__(
self,
reduction: Literal["mean", "sum", "none"] = "mean",
ignore_index: int = -100,
) -> None:
super().__init__("categorical_cross_entropy_loss", 1, reduction)
self.ignore_index = ignore_index
def __call__(
self,
predictions: Var | Type,
targets: Var | Type,
weights: Var | Type | None = None,
) -> Function:
"""Get the relax function of CategoricalCrossEntropyLoss. If the parameters are
type, it will create corresponding variables.
Parameters
----------
predictions : Union[Var, Type]
The predictions of the model in the calculation of loss.
targets : Union[Var, Type]
The ground truth in the calculation of loss.
weights : Optional[Union[Var, Type]]
a manual rescaling weight given to each class. It has to be a Tensor of size C.
Returns
-------
The relax function of CategoricalCrossEntropyLoss with the loss name as its global symbol.
"""
if "int" not in str(targets.dtype):
raise TypeError(
f"Dtype of targets expected to be int/uint. \
However, the dtype of targets is {targets.dtype}"
)
bb = BlockBuilder()
predictions = _create_param_var(predictions, "predictions")
targets = _create_param_var(targets, "targets")
arg_list = [predictions, targets]
if weights:
weights = _create_param_var(weights, "weights")
arg_list.append(weights)
# In the case of ignore_index >= 0,
# the nll_loss function is used to handle the ignore index.
# In other cases where ignore_index is not needed, just use the simpe product.
with bb.function(self._loss_name, arg_list):
with bb.dataflow():
logits = bb.emit(log_softmax(predictions))
if self.ignore_index >= 0:
targets = bb.emit(
reshape(argmax(targets, axis=1), shape=(targets.ty.shape[0],))
)
loss = bb.emit_output(
nll_loss(logits, targets, weights, self._reduction, self.ignore_index)
)
else:
lv = bb.emit(-logits * targets.astype("float32"))
if weights:
lv = bb.emit(lv * weights)
loss = bb.emit_output(self._with_reduction(lv))
bb.emit_func_output(loss)
return bb.get()[self._loss_name]
python/tvm/relax/training/optimizer.py
+722
View File
@@ -0,0 +1,722 @@
# 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]
python/tvm/relax/training/setup_trainer.py
+213
View File
@@ -0,0 +1,213 @@
# 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.
# pylint: disable=not-callable, unused-argument
"""Setup Trainer Pass."""
import tvm
from tvm.ir.module import IRModule
from tvm.tirx.expr import IntImm
from ..analysis import check_well_formed
from ..expr import Tuple
from ..training.utils import AppendLoss
from ..transform import DecomposeOpsForInference, DecomposeOpsForTraining, Gradient, LegalizeOps
from ..type import TensorType
from .loss import Loss
from .optimizer import Optimizer
@tvm.transform.module_pass(opt_level=0, name="SetupTrainer")
class SetupTrainer:
"""Transform a backbone module to a complete, legalized trainer module.
The provided backbone module should contain at least a function named `backbone`, and has two
int attributes `param_num` and `state_num`, as follows:
.. code-block:: python
@I.ir_module
class Backbone:
I.module_attrs({"param_num": 1, "state_num": 1})
@R.function
def backbone(input_instances, parameters, states):
# Predicts the result
# Should contain only one DataflowBlock
...
return backbone_result, updated_states
Here each of input_instances, parameters, states, backbone_result and updated_states can
denote a number of parameters. The length of parameters and the length of states is specified
by param_num and state_num respectively.
`states` denote the states that we need to maintain as the training process proceeds, such as
the running mean and the running var of the batch norm operator. The updated states is returned
in `updated_states`. States can be empty if there is no state that needs to be updated.
The transformed module will at least contain the functions and attributes listed below:
.. code-block:: python
@I.ir_module
class Module:
I.module_attrs({"input_num": 1, "param_num": 1, "state_num": 1, "optim_states": ...})
@R.function
def backbone(input_instances, parameters, states):
# Predicts the result. It is provided in the input module.
...
return backbone_result, updated_states
@R.function
def backbone_loss(input_instances, parameters, states, targets):
# Runs like backbone and then computes the loss between the result and targets.
...
return loss, updated_states
@R.function
def backbone_loss_adjoint(input_instances, parameters, states, targets):
# Runs like backbone_loss and then calculates the gradient of parameters.
...
return (loss, updated_states), gradient_of_params
@R.function
def optimizer(params, gradients, optim_states):
# Update parameters and optimizer states with the gradient computed
...
return (updated_params, updated_optim_states)
The transformed module contains an attribute `optim_states` as the initial optimizer states.
Then the transformed module will be legalized by `relax.transform.LegalizeOps()` to lower
relax operators into TIR functions.
Parameters
----------
loss : Loss
The loss function. It will be appended to the backbone function using
relax.transform.AppendLoss.
optimizer : Optimizer
The optimizer. It will be put as the `optimizer` function of the transformed module.
loss_args : List[TensorType]
The arguments to call the loss function.
legalize : bool
Whether to legalize the module. Default: True.
"""
BACKBONE_FUNC: str = "backbone"
BACKBONE_LOSS_FUNC: str = "backbone_loss"
ADJOINT_FUNC: str = "backbone_loss_adjoint"
OPTIMIZER_FUNC: str = "optimizer"
PARAM_NUM_ATTR_KEY: str = "param_num"
STATE_NUM_ATTR_KEY: str = "state_num"
def __init__(
self, loss: Loss, optimizer: Optimizer, loss_args: list[TensorType], legalize=True
):
self._loss = loss
self._optimizer = optimizer
self._loss_args = loss_args
self._legalize = legalize
def _check_well_formed(self, mod: IRModule):
if not check_well_formed(mod):
raise ValueError("SetupTrainer: The backbone module is not well formed.")
try:
func = mod[self.BACKBONE_FUNC]
except (KeyError, ValueError) as exc:
raise ValueError(
f"SetupTrainer: The backbone module does not contain a function named "
f"{self.BACKBONE_FUNC}"
) from exc
# Check function attrs
if self.PARAM_NUM_ATTR_KEY not in mod.attrs or not isinstance(
mod.attrs[self.PARAM_NUM_ATTR_KEY], IntImm | int
):
raise ValueError(
f"SetupTrainer: The backbone module should has an integer attribute named "
f"{self.PARAM_NUM_ATTR_KEY}"
)
if self.STATE_NUM_ATTR_KEY not in mod.attrs or not isinstance(
mod.attrs[self.STATE_NUM_ATTR_KEY], IntImm | int
):
raise ValueError(
f"SetupTrainer: The backbone module should has an integer attribute named "
f"{self.STATE_NUM_ATTR_KEY}"
)
nparam = int(mod.attrs[self.PARAM_NUM_ATTR_KEY])
nstate = int(mod.attrs[self.STATE_NUM_ATTR_KEY])
# Check parameters and return values
if len(func.params) < nparam + nstate:
raise ValueError(
"SetupTrainer: The number of parameters of the predict function should be no less "
"than the number of parameters and states"
)
if nstate > 0:
if not isinstance(func.body.body, Tuple) or len(func.body.body) <= nstate:
raise ValueError(
"SetupTrainer: When model state exists, the predict function should return a "
"tuple of length more than the number of states"
)
def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
"""Transform the backbone module into a trainer module."""
self._check_well_formed(mod)
mod = AppendLoss(
self.BACKBONE_FUNC,
self._loss(*self._loss_args), # type: ignore
self._loss.num_backbone_outputs,
self.BACKBONE_LOSS_FUNC,
)(mod)
# Decompose batch_norm operator, which behaves differently in inference and training stages
mod = DecomposeOpsForInference(self.BACKBONE_FUNC)(mod)
mod = DecomposeOpsForTraining(self.BACKBONE_LOSS_FUNC)(mod)
# Gradient pass.
param_num = int(mod.attrs[self.PARAM_NUM_ATTR_KEY])
state_num = int(mod.attrs[self.STATE_NUM_ATTR_KEY])
input_num = len(mod[self.BACKBONE_FUNC].params) - param_num - state_num
params = mod[self.BACKBONE_LOSS_FUNC].params[input_num : input_num + param_num]
mod = Gradient(self.BACKBONE_LOSS_FUNC, require_grads=params, target_index=0)(mod)
# Add optimizer function.
self._optimizer.init(params)
# Need the global symbol to match the function's name
mod[self.OPTIMIZER_FUNC] = self._optimizer.get_function().with_attr(
"global_symbol", self.OPTIMIZER_FUNC
)
# Module attrs
mod = mod.with_attrs(
{
"input_num": input_num,
"optim_state": self._optimizer.state,
}
)
if self._legalize:
mod = LegalizeOps()(mod)
return mod
python/tvm/relax/training/trainer.py
+349
View File
@@ -0,0 +1,349 @@
# 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.
# pylint: disable=invalid-name
"""Unified Trainer API for relax training."""
import numpy as np # type: ignore
import tvm
from tvm import relax
from tvm.ir.module import IRModule
from tvm.runtime._tensor import Tensor
class Trainer:
r"""Unified wrapper for relax training. It accepts the IRModule (that is the result of
SetupTrainer) and the relax VM (that contains the built result of the IRModule), and helps run
the VM. It maintains the parameters, the model states and the optimizer states internally.
Parameters
----------
train_mod : tvm.IRModule
The IRModule that will be run. Should be the result of a backbone module being transformed
by the SetupTrainer pass.
vm : tvm.relax.VirtualMachine
The relax virtual machine that contains the built result of train_mod. Considering the
complexity and flexibility of building, we require user build the train_mod outside of
trainer and pass the result vm.
device : tvm.runtime.Device
The device to place the parameters and states in.
zero_init_param_state : bool
If true, all parameters and states will be inited to zero. It requires all parameters and
states have static shape.
Examples
--------
.. code-block:: python
setup_trainer = SetupTrainer(
MSELoss(reduction="sum"),
SGD(0.001),
[pred_ty, target_ty],
)
train_mod = setup_trainer(Backbone)
ex = tvm.compile(train_mod, target)
vm = relax.VirtualMachine(ex, dev)
trainer = training.Trainer(train_mod, vm, dev, False)
trainer.xaiver_uniform_init_params()
trainer.predict(input_instances)
trainer.update([input_instances], [labels])
"""
BACKBONE_FUNC: str = "backbone"
BACKBONE_LOSS_FUNC: str = "backbone_loss"
ADJOINT_FUNC: str = "backbone_loss_adjoint"
OPTIMIZER_FUNC: str = "optimizer"
def __init__(
self,
train_mod: IRModule,
vm: relax.VirtualMachine,
device: tvm.runtime.Device,
zero_init_param_state: bool = True,
) -> None:
self.mod = train_mod.without_attr("optim_state")
self.vm = vm
self.device = device
self._optim_state = [d.copyto(device) for d in train_mod.attrs["optim_state"]]
self._input_num = int(train_mod.attrs["input_num"])
self._param_num = int(train_mod.attrs["param_num"])
self._state_num = int(train_mod.attrs["state_num"])
# are used to initialize params and states
self._param_vars = train_mod[self.ADJOINT_FUNC].params[
self._input_num : self._input_num + self._param_num
]
self._state_vars = train_mod[self.ADJOINT_FUNC].params[
(self._input_num + self._param_num) : (
self._input_num + self._param_num + self._state_num
)
]
self._params: list[Tensor | None] = [None] * self._param_num
self._param_name_to_pos: dict[str, int] = {
p.name_hint: i for i, p in enumerate(self._param_vars)
}
self._states: list[Tensor | None] = [None] * self._state_num
self._state_name_to_pos: dict[str, int] = {
s.name_hint: i for i, s in enumerate(self._state_vars)
}
if zero_init_param_state:
self.zero_init_params()
self.zero_init_states()
@staticmethod
def _get_shape_list(expr):
return [int(dim) for dim in expr.ty.shape]
def xaiver_uniform_init_params(self):
"""Xaiver uniformly initialize parameters using the method described in `Understanding the
difficulty of training deep feedforward neural networks` - Glorot, X. & Bengio, Y.
(2010).
Requires all parameters have static shapes.
"""
self._params = []
for p in self._param_vars:
shape, dtype = self._get_shape_list(p), p.ty.dtype
self._params.append(
tvm.runtime.tensor(
(np.sqrt(6.0 / np.sum(shape)) * np.random.uniform(-1.0, 1.0, shape)).astype(
dtype
),
self.device,
)
)
def zero_init_params(self):
"""Zero initialize all parameters. Requires all parameters have static shapes."""
self._params = [
tvm.runtime.tensor(np.zeros(self._get_shape_list(p), p.ty.dtype.dtype), self.device)
for p in self._param_vars
]
def zero_init_states(self):
"""Zero initialize all states. Requires all states have static shapes."""
self._states = [
tvm.runtime.tensor(np.zeros(self._get_shape_list(s), s.ty.dtype.dtype), self.device)
for s in self._state_vars
]
def load_params(
self,
params: list[np.ndarray | Tensor] | dict[str, np.ndarray | Tensor],
):
"""Load parameters from a dict or a list. Will convert parameters into tvm.runtime.Tensor
in self.device.
Parameters
----------
params : List[Union[np.ndarray, Tensor]], Dict[str, Union[np.ndarray, Tensor]]
The numerical value of the parameters.
If params is a list, its length should be param_num. The value of parameters at the
corresponding index will be updated.
If params is a dict, it should map variable name to value. The name should be the same
as the parameter name in the backbone function. The values of the corresponding
parameters will be updated.
"""
if isinstance(params, list):
if len(params) != self._param_num:
raise ValueError(
f"The length of extern parameters is {len(params)}, which does not "
f"match the number of parameters {self._param_num}"
)
self._params = [tvm.runtime.tensor(v, self.device) for v in params]
elif isinstance(params, dict):
for key, val in params.items():
if key not in self._param_name_to_pos:
raise ValueError(f"Parameter {key} is not found in the model")
self._params[self._param_name_to_pos[key]] = tvm.runtime.tensor(val, self.device)
else:
raise ValueError("The type of extern_params should be either list or dict")
def load_states(
self,
states: list[np.ndarray | Tensor] | dict[str, np.ndarray | Tensor],
):
"""Load model states from a dict or a list. Will convert states into tvm.runtime.Tensor
in self.device.
Parameters
----------
states : List[Union[np.ndarray, Tensor]], Dict[str, Union[np.ndarray, Tensor]]
The numerical value of the model states.
If states is a list, its length should be state_num. The value of states at the
corresponding index will be updated.
If params is a dict, it should map variable name to value. The name should be the same
as the state name in the backbone function. The values of the corresponding states will
be updated.
"""
if isinstance(states, list):
if len(states) != self._state_num:
raise ValueError(
f"The length of extern states is {len(states)}, which does not match "
f"the number of model states {self._state_num}"
)
self._states = [tvm.runtime.tensor(v, self.device) for v in states]
elif isinstance(states, dict):
for key, val in states.items():
if key not in self._param_name_to_pos:
raise ValueError(f"Parameter {key} is not found in the model")
self._states[self._param_name_to_pos[key]] = tvm.runtime.tensor(val, self.device)
else:
raise ValueError("The type of extern_states should be either list or dict")
def export_params(self) -> dict[str, Tensor]:
"""Export parameters to a dict (parameter name -> Tensor).
Returns
-------
exported_dict : Dict[str, Tensor]
The exported dictionary of parameters.
"""
return {key: self._params[pos] for key, pos in self._param_name_to_pos.items()}
def export_states(self) -> dict[str, Tensor]:
"""Export model states to a dict (parameter name -> Tensor).
Returns
-------
exported_dict : Dict[str, Tensor]
The exported dictionary of model states.
"""
return {key: self._states[pos] for key, pos in self._state_name_to_pos.items()}
def _check_inited(self):
"""Check that all parameters and model states are initialized."""
idx_not_inited_param = next((i for i, p in enumerate(self._params) if p is None), -1)
if idx_not_inited_param != -1:
raise RuntimeError(
f"The {idx_not_inited_param}-th parameter is not initialized before training or "
"inference."
)
idx_not_inited_state = next((i for i, s in enumerate(self._states) if s is None), -1)
if idx_not_inited_state != -1:
raise RuntimeError(
f"The {idx_not_inited_state}-th model state is not initialized before training or "
"inference."
)
def predict(self, *input_instances: np.ndarray | Tensor) -> Tensor:
"""Call the `backbone` function and return the prediction result of the backbone.
Parameters
----------
*input_instances : Union[np.ndarray, Tensor]
The values corresponding to the input_instances part of the backbone function.
Parameters and model states are not needed to provide.
Returns
-------
output : Tensor
The result of the backbone function. If the backbone contains model states, the updated
states WILL NOT be returned.
"""
self._check_inited()
if len(input_instances) != self._input_num:
raise ValueError("The length of the input does not match the backbone")
all_inputs: list[Tensor] = (
[tvm.runtime.tensor(i, self.device) for i in input_instances]
+ self._params
+ self._states
)
res = self.vm[self.BACKBONE_FUNC](*all_inputs)
# remove the states part, if they exist
if self._state_num != 0:
res = res[: -self._state_num]
if len(res) == 1:
res = res[0]
return res
def update(
self,
input_instances: np.ndarray | Tensor | list[np.ndarray | Tensor],
targets: np.ndarray | Tensor | list[np.ndarray | Tensor],
) -> Tensor:
"""Update parameters and model states. It will calculate the gradients of parameters
and update them using the `optimizer` function.
Parameters, model states and optimizer states are provided in the function, so you do not
need to provied them.
Parameters
----------
input_instances : Union[np.ndarray, Tensor, List[Union[np.ndarray, Tensor]]]
The values corresponding to the input_instances part of the backbone function.
Parameters and model states are not needed to provide.
If there are more than one input instances, you can provide a list.
targets : Union[np.ndarray, Tensor, List[Union[np.ndarray, Tensor]]]
The values corresponding to the targets part of the backbone function.
If there are more than one targets, you can provide a list.
Returns
-------
loss : Tensor
The loss stored in tvm.runtime.Tensor.
"""
self._check_inited()
if not isinstance(input_instances, list):
input_instances = [input_instances]
if not isinstance(targets, list):
targets = [targets]
if len(input_instances) != self._input_num:
raise ValueError("The length of the input does not match the backbone")
all_inputs: list[Tensor] = (
[tvm.runtime.tensor(i, self.device) for i in input_instances]
+ self._params
+ self._states
+ [tvm.runtime.tensor(i, self.device) for i in targets]
)
ret, grads = self.vm[self.ADJOINT_FUNC](*all_inputs)
# update model states
if self._state_num != 0:
self._states = list(ret[1:])
ret = ret[0]
# update params
new_params, self._optim_state = self.vm[self.OPTIMIZER_FUNC](
self._params, grads, self._optim_state
)
self._params = list(new_params)
return ret
python/tvm/relax/training/utils.py
+211
View File
@@ -0,0 +1,211 @@
# 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.
# pylint: disable=invalid-name, unused-argument
"""Utility functions for relax training."""
from collections.abc import Callable
from tvm_ffi import register_global_func
import tvm
from tvm import relax
from tvm.ir import Call
from tvm.relax.block_builder import BlockBuilder
from ..expr import Function, Var
from . import _ffi_api
def AppendLoss(
func_name: str,
loss_function: Function,
num_backbone_outputs: int = 1,
new_func_name: str | None = None,
) -> tvm.ir.transform.Pass:
"""Append the loss function to the backbone function specified by `func_name`. Generally, the
loss function is generated by instances of `relax.training.Loss`.
The backbone function and the loss function should satisfy a few restrictions:
- Both backbone and loss should contain exactly one DataflowBlock.
- Backbone should return either one Var, or a tuple of Vars
- Loss should return a scalar(0-dim Tensor) Var
They should be like:
.. code-block:: python
@R.function
def backbone(input_instances, parameters, states):
with R.dataflow():
# Predicts the result
...
return backbone_result, updated_states
@R.function
def loss(backbone_result, targets):
with R.dataflow():
# calculate the loss between backbone_result and targets
...
# loss should be a scalar Var
return loss
Here each of input_instances, parameters, states, backbone_result and updated_states can
denote a number of parameters.
`states` denote the states that we need to maintain as the training process proceeds, such as
the running mean and the running var of the batch norm operator. The updated states is returned
in `updated_states`. States can be empty if there is no state that needs to be updated.
The appended result contains only one DataflowBlock containing all bindings in backbone and
loss. It will be like:
.. code-block:: python
@R.function
def backbone_loss(input_instances, parameters, states, targets):
with R.dataflow():
# all bindings in backbone and loss
...
return loss, updated_states
Parameters
----------
func_name : str
The name of the backbone function in the IRModule.
loss_func : Function
The loss function.
num_backbone_outputs : int
Specify the number of `prediction_outputs` of the backbone function. Default: 1.
new_func_name : Optional[str]
Specify the name of the appended result. If it is not specified, the name will be
`func_name + "_loss"`.
Returns
-------
ret : Function
The result function.
Examples
--------
.. code-block:: python
@I.ir_module
class Module
@R.function
def predict(x: R.Tensor((2, 4), "float32"), y: R.Tensor((2, 4), "float32")):
with R.dataflow():
out = R.add(x, y)
R.output(out)
return out
@R.function
def loss(predictions: R.Tensor((2, 4), "float32"), labels: R.Tensor((2, 4), "float32")):
with R.dataflow():
lv = R.subtract(predictions, labels)
lv1 = R.multiply(lv, lv)
gv = R.sum(lv1)
R.output(gv)
return gv
expected = AppendLoss("predict", loss)(Module)
expected.show()
Will get
.. code-block:: python
@I.ir_module
class Module
@R.function
def predict(x: R.Tensor((2, 4), "float32"), y: R.Tensor((2, 4), "float32")):
with R.dataflow():
out = R.add(x, y)
R.output(out)
return out
@R.function
def predict_loss(x: R.Tensor((2, 4), "float32"), y: R.Tensor((2, 4), "float32"),
labels: R.Tensor((2, 4), "float32")) -> R.Tensor((), "float32"):
with R.dataflow():
out: R.Tensor((2, 4), "float32") = R.add(x, y)
lv: R.Tensor((2, 4), "float32") = R.subtract(out, labels)
lv1: R.Tensor((2, 4), "float32") = R.multiply(lv, lv)
gv: R.Tensor((), "float32") = R.sum(lv1)
R.output(gv)
return gv
Notes
-----
This util can be replaced if we have inline pass. It is equivalent to inline a tail call in
some sense.
"""
return _ffi_api.AppendLoss( # type: ignore
func_name,
loss_function,
num_backbone_outputs,
new_func_name,
)
def register_te_gradient(te_grad_name: str, te_grad_func: Callable | None = None):
"""Register a te gradient function bind with name te_grad_name. te_grad_name can be referenced
later in call_tir_with_grad nodes.
Parameters
----------
te_grad_name : str
The registered name of the te gradient function. Should be align with the te_grad_name in
call_tir_with_grad nodes.
grad_func : Callable
The te grad function.
It must be a function taking (output_grad: Tensor, arg1: Tensor, arg2: Tensor, ...)
as inputs and returning a list of Tensor created by te.compute.
Returns
-------
mod : IRModule
The mod with corresponding attributes attached.
"""
def register(func: Callable):
func_prefix = "tvm.relax.te_grad._register."
# The handler function is used to let the backend (cpp side) to emit_te.
# It's a wrapper of the te_grad_func.
# It takes the blockbuilder, the gradient var of the output and the forward call expr.
# It will return the emitted var.
def handler(
orig_var: Var, call_tir_with_grad: Call, output_grad: Var, ctx: BlockBuilder
) -> relax.Expr:
return ctx.emit_te(
func,
output_grad,
*call_tir_with_grad.args[1],
**call_tir_with_grad.attrs.te_grad_kwargs,
primfunc_name_hint=te_grad_name,
)
register_global_func(func_prefix + te_grad_name, handler)
return func
return register(te_grad_func) if te_grad_func else register