Repository: microsoft/nni
File: nni/retiarii/experiment/pytorch.py (Python)

Code before:
class RetiariiExperiment(Experiment):
    def __init__(self, base_model: nn.Module, trainer: Union[Evaluator, BaseOneShotTrainer],
                 applied_mutators: List[Mutator] = None, strategy: BaseStrategy = None):
        # TODO: The current design of init interface of Retiarii experiment needs to be reviewed.
        self.config: RetiariiExeConfig = None
        self.port: Optional[int] = None

        self.base_model = base_model
        self.trainer = trainer
        self.applied_mutators = applied_mutators
        self.strategy = strategy

        self._dispatcher = RetiariiAdvisor()
        self._dispatcher_thread: Optional[Thread] = None
        self._proc: Optional[Popen] = None
        self._pipe: Optional[Pipe] = None

    def _start_strategy(self):
        base_model_ir, self.applied_mutators = preprocess_model(
            self.base_model, self.trainer, self.applied_mutators,
            full_ir=self.config.execution_engine not in ['py', 'benchmark'],
            dummy_input=self.config.dummy_input
        )

        _logger.info('Start strategy...')
        search_space = BaseStrategy.get_model_space(base_model_ir, self.applied_mutators)
        self.update_search_space(search_space)
        self.strategy.run(base_model_ir, self.applied_mutators)
        _logger.info('Strategy exit')
        # TODO: find out a proper way to show no more trial message on WebUI
        # self._dispatcher.mark_experiment_as_ending()

    def start(self, port: int = 8080, debug: bool = False) -> None:
        """
        Start the experiment in background.
        This method will raise exception on failure.
        If it returns, the experiment should have been successfully started.
        Parameters
        ----------
        port
            The port of web UI.
        debug
            Whether to start in debug mode.
        """
        atexit.register(self.stop)

        # we will probably need a execution engine factory to make this clean and elegant
        if self.config.execution_engine == 'base':
            from ..execution.base import BaseExecutionEngine
            engine = BaseExecutionEngine()
        elif self.config.execution_engine == 'cgo':

Review comment (refactor):
Should call `dry_run_for_formated_search_space` here. Not sure whether every search space can be successfully formatted though.

Code after:

class RetiariiExperiment(Experiment):
    def __init__(self, base_model: nn.Module, trainer: Union[Evaluator, BaseOneShotTrainer],
                 applied_mutators: List[Mutator] = None, strategy: BaseStrategy = None):
        # TODO: The current design of init interface of Retiarii experiment needs to be reviewed.
        self.config: RetiariiExeConfig = None
        self.port: Optional[int] = None

        self.base_model = base_model
        self.trainer = trainer
        self.applied_mutators = applied_mutators
        self.strategy = strategy

        self._dispatcher = RetiariiAdvisor()
        self._dispatcher_thread: Optional[Thread] = None
        self._proc: Optional[Popen] = None
        self._pipe: Optional[Pipe] = None

    def _start_strategy(self):
        base_model_ir, self.applied_mutators = preprocess_model(
            self.base_model, self.trainer, self.applied_mutators,
            full_ir=self.config.execution_engine not in ['py', 'benchmark'],
            dummy_input=self.config.dummy_input
        )

        _logger.info('Start strategy...')
        search_space = dry_run_for_formatted_search_space(base_model_ir, self.applied_mutators)
        self.update_search_space(search_space)
        self.strategy.run(base_model_ir, self.applied_mutators)
        _logger.info('Strategy exit')
        # TODO: find out a proper way to show no more trial message on WebUI
        # self._dispatcher.mark_experiment_as_ending()

    def start(self, port: int = 8080, debug: bool = False) -> None:
        """
        Start the experiment in background.
        This method will raise exception on failure.
        If it returns, the experiment should have been successfully started.
        Parameters
        ----------
        port
            The port of web UI.
        debug
            Whether to start in debug mode.
        """
        atexit.register(self.stop)

        # we will probably need a execution engine factory to make this clean and elegant
        if self.config.execution_engine == 'base':
            from ..execution.base import BaseExecutionEngine
            engine = BaseExecutionEngine()