Method get_function

python/tvm/relax/training/optimizer.py:304–353

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)

        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", TupleStructInfo([p.struct_info for p in plist]))
        grad_var = Var("gradients", TupleStructInfo([p.struct_info for p in plist]))
        state_var = Var("optim_states", TupleStructInfo([TensorStructInfo((), "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):