Repository: microsoft/nni
File: nni/contrib/compression/utils/evaluator.py (Python)

Code before:
        self.model = model
        self._ori_model_attr.update({
            'training_step': model.training_step,
            'configure_optimizers': model.configure_optimizers,
            'configure_callbacks': model.configure_callbacks
        })
        self._param_names_map = param_names_map
        self._patch_configure_optimizers()

    def unbind_model(self):
        if self.model:
            self.revert_loss()
            self.revert_optimizer_step()
            self.remove_all_hooks()
            self._revert_configure_optimizers()
            self._param_names_map = None
            self._ori_model_attr.clear()
            self.model = None
        else:
            _logger.warning('Did not bind any model, no need to unbind model.')

    def patch_optim_param_group(self, module_name_param_dict: Dict[str, List[Tensor]]):
        assert isinstance(self.model, pl.LightningModule)
        assert module_name_param_dict is not None

        optimizers_lr_schedulers: Any = self.model.configure_optimizers()

        if self._opt_returned_dicts:
            def new_configure_optimizers(_):  # type: ignore
                optimizers = [opt_lrs_dict['optimizer'] for opt_lrs_dict in optimizers_lr_schedulers]
                # add param group
                self._optimizer_add_param_group(self.model, module_name_param_dict, optimizers) # type: ignore

                return optimizers_lr_schedulers

        elif self._lr_scheduler_helpers:
            def new_configure_optimizers(_):  # type: ignore
                optimizers, lr_schedulers = optimizers_lr_schedulers
                # add param_group
                self._optimizer_add_param_group(self.model, module_name_param_dict, optimizers) # type: ignore

                return optimizers, lr_schedulers

        else:
            def new_configure_optimizers(_):
                # add param_group
                self._optimizer_add_param_group(self.model, module_name_param_dict, optimizers_lr_schedulers) # type: ignore

                return optimizers_lr_schedulers

        self.model.configure_optimizers = types.MethodType(new_configure_optimizers, self.model)

Review comment (performance):
please init the optimizers when call `configure_optimizers`, that is:
```
old_configure_optimizers = self.model.configure_optimizers

def new_configure_optimizers(_):
    optimizers_lr_schedulers = old_configure_optimizers()
    ...
```

Code after:
            raise TypeError(err_msg)

        self._initialization_complete = True

    def bind_model(self, model: pl.LightningModule, param_names_map: Dict[str, str] | None = None):
        err_msg = 'Evaluator initialization is not complete, please call `_init_optimizer_helpers` before bind model.'
        assert self._initialization_complete is True, err_msg
        assert isinstance(model, pl.LightningModule)
        if self.model is not None:
            _logger.warning('Already bound a model, will unbind it before bind a new model.')
            self.unbind_model()

        self.model = model
        self._ori_model_attr.update({
            'training_step': model.training_step,
            'configure_optimizers': model.configure_optimizers,
            'configure_callbacks': model.configure_callbacks
        })
        self._param_names_map = param_names_map
        self._patch_configure_optimizers()

    def unbind_model(self):
        if self.model:
            self.revert_loss()
            self.revert_optimizer_step()
            self.remove_all_hooks()
            self._revert_configure_optimizers()
            self._param_names_map = None
            self._ori_model_attr.clear()
            self.model = None
        else:
            _logger.warning('Did not bind any model, no need to unbind model.')

    def patch_optim_param_group(self, module_name_param_dict: Dict[str, List[Tensor]]):
        assert isinstance(self.model, pl.LightningModule)
        assert module_name_param_dict is not None

        old_configure_optimizers = self.model.configure_optimizers

        if self._opt_returned_dicts:
            def new_configure_optimizers(_):  # type: ignore
                optimizers_lr_schedulers: Any = old_configure_optimizers() # type: ignore
                optimizers = [opt_lrs_dict['optimizer'] for opt_lrs_dict in optimizers_lr_schedulers]
                # add param group
                self._optimizer_add_param_group(self.model, module_name_param_dict, optimizers) # type: ignore

                return optimizers_lr_schedulers

        elif self._lr_scheduler_helpers:
            def new_configure_optimizers(_):  # type: ignore
                optimizers_lr_schedulers: Any = old_configure_optimizers() # type: ignore