apache--tvm/python/tvm/relax/training/trainer.py

# 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