Method transform_module

python/tvm/relax/training/setup_trainer.py:173–213

Transform the backbone module into a trainer module.

(self, mod: IRModule, ctx: tvm.transform.PassContext)

                )

    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