Source code for nni.nas.experiment.config.experiment

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import logging
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, TYPE_CHECKING, Union, List
from typing_extensions import Literal

from nni.experiment.config import utils, ExperimentConfig

from .engine import ExecutionEngineConfig, SequentialEngineConfig
from .format import ModelFormatConfig

    from nni.nas.evaluator import Evaluator
    from import BaseModelSpace
    from nni.nas.strategy import Strategy

__all__ = ['NasExperimentConfig']

RESERVED = '_reserved_'

_logger = logging.getLogger(__name__)

[docs] @dataclass(init=False) class NasExperimentConfig(ExperimentConfig): """Config for NAS experiment. Other than training service fields which also exists in a HPO experiment, additional fields provided by NAS include execution engine and model format. Execution engine is used to specify how to (e.g., distributedly or sequentially) run a trial, and model format specifies the format of the converted model space used throughout the NAS experiment. It can be constructed via 3 approaches. 1. Create a default config and then modify some fields (recommended). The default config should be good enough for most cases. Users only need to update some fields like concurrency. See details in :meth:`default`. Example:: config = NasExperimentConfig.default(model_space, evaluator, strategy) config.trial_concurrency = 4 2. Create an object by providing several required fields, and then set other fields. Though marked as optional in function signature, it's recommended to set all three fields. config = NasExperimentConfig('ts', 'graph', 'local') config.experiment_name = 'hello' config.execution_engine.dummy_input = [1, 3, 224, 224] 3. Create an empty object and set all fields manually. Put the fields into kwargs should also work:: config = NasExperimentConfig() config.execution_engine = TrainingServiceEngineConfig() config.model_format = SimplifiedModelFormatConfig() config.training_service = LocalConfig(use_active_gpu=True) # equivalent to config = NasExperimentConfig( execution_engine=TrainingServiceEngineConfig(), model_format=SimplifiedModelFormatConfig(), training_service=LocalConfig(use_active_gpu=True) ) """ # TODO: refactor this class to inherit from a new common base class with HPO config experiment_type: Literal['nas'] = 'nas' search_space: Any = RESERVED trial_code_directory: utils.PathLike = '.' trial_command: str = '' # New config field for NAS execution_engine: ExecutionEngineConfig model_format: ModelFormatConfig # Internal: to support customized fields in trial command # Useful when customized python / environment variables are needed _trial_command_params: Optional[Dict[str, Any]] = None def __init__(self, execution_engine: Union[str, ExecutionEngineConfig, None] = None, model_format: Union[str, ModelFormatConfig, None] = None, training_service_platform: Union[str, List[str], None] = None, **kwargs): # `execution_engine` and `model_format` are two shortcuts for easy configuration. # We merge them into `kwargs` and let the parent class handle them. if isinstance(execution_engine, str): kwargs.update(execution_engine=ExecutionEngineConfig(name=execution_engine)) elif isinstance(execution_engine, ExecutionEngineConfig): kwargs.update(execution_engine=execution_engine) elif execution_engine is not None: raise TypeError('execution_engine must be a string or an ExecutionEngineConfig object.') if isinstance(model_format, str): kwargs.update(model_format=ModelFormatConfig(name=model_format)) elif isinstance(model_format, ModelFormatConfig): kwargs.update(model_format=model_format) elif model_format is not None: raise TypeError('model_format must be a string or a ModelFormatConfig object.') super().__init__(training_service_platform=training_service_platform, **kwargs)
[docs] @classmethod def default(cls, model_space: 'BaseModelSpace', evaluator: 'Evaluator', strategy: 'Strategy') -> 'NasExperimentConfig': """Instantiate a default config. Infer from current setting of model space, evaluator and strategy. If the strategy is found to be a one-shot strategy, the execution engine will be set to "sequential" and model format will be set to "raw" to preserve the weights and the model object. If the strategy is found to be a multi-trial strategy, training service engine will be used by default, and the training service will be set to "local" if not provided. Model format will be set to "simplified" for performance and memory efficiency. """'Config is not provided. Will try to infer.') trial_concurrency = 1 # no effect if not going parallel training_service = 'local' # no effect if not using training service execution_engine = None model_format = None try: from nni.nas.oneshot.pytorch.strategy import OneShotStrategy, is_supernet from nni.nas.nn.pytorch import ModelSpace if isinstance(strategy, OneShotStrategy):'Strategy is found to be a one-shot strategy. ' 'Setting execution engine to "sequential" and format to "raw".') execution_engine = 'sequential' model_format = 'raw' if isinstance(model_space, ModelSpace) and is_supernet(model_space):'Model space is found to be a one-shot supernet. ' 'Setting execution engine to "sequential" and format to "raw" to preserve the weights.') execution_engine = 'sequential' model_format = 'raw' except ImportError: _logger.warning('Import of one-shot strategy failed. Assuming the strategy is not one-shot.') if execution_engine is None:'Using execution engine based on training service. Trial concurrency is set to 1.') execution_engine = 'ts' if model_format is None:'Using simplified model format.') model_format = 'simplified' if execution_engine == 'ts':'Using local training service.') import torch if torch.cuda.is_available() and torch.cuda.device_count() > 0: _logger.warning( 'GPU found but will not be used. Please set `experiment.config.trial_gpu_number` to ' 'the number of GPUs you want to use for each trial.' ) config = NasExperimentConfig( training_service_platform=training_service, execution_engine=execution_engine, model_format=model_format, trial_concurrency=trial_concurrency ) return config
def _canonicalize(self, parents): if self.search_space != RESERVED and self.search_space != {}: raise ValueError('`search_space` field can not be customized in NAS experiment.') self.search_space = {} if not Path(self.trial_code_directory).samefile(Path.cwd()): raise ValueError('`trial_code_directory` field can not be customized in NAS experiment.') trial_command_suffix = '-m nni.nas.execution.training_service trial' trial_command_template = '{envs} {python} ' + trial_command_suffix if self.trial_command and not self.trial_command.endswith(trial_command_suffix): raise ValueError('`trial_command` field can not be customized in NAS experiment.') _trial_command_params = { # Default variables 'envs': '', # TODO: maybe use sys.executable rendered in trial side (e.g., trial_runner) 'python': sys.executable, # This should override the parameters above. **(self._trial_command_params or {}) } self.trial_command = trial_command_template.format(**_trial_command_params).strip() if isinstance(self.execution_engine, SequentialEngineConfig): if not utils.is_missing(self.trial_concurrency) and self.trial_concurrency != 1: raise ValueError('`trial_concurrency` must be 1 for sequential execution engine.') self.trial_concurrency = 1 if not utils.is_missing(self.training_service): _logger.warning('`training_service` will be ignored for sequential execution engine.') self.training_service = utils.training_service_config_factory('local') super()._canonicalize([self] + parents)