wehub-resource-sync
214 lines
8.2 KiB
Python
214 lines
8.2 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# 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
|