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/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