From c3d3f3ca773c90db63729e26bf5743a31fbc18b2 Mon Sep 17 00:00:00 2001 From: Brian Raf Date: Fri, 26 Apr 2024 14:50:26 +0000 Subject: [PATCH] Adding cli option for optuna search --- model_analyzer/config/input/config_command_profile.py | 5 +++-- tests/test_config.py | 9 ++++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/model_analyzer/config/input/config_command_profile.py b/model_analyzer/config/input/config_command_profile.py index fc13cdb08..4e505fc16 100755 --- a/model_analyzer/config/input/config_command_profile.py +++ b/model_analyzer/config/input/config_command_profile.py @@ -920,14 +920,15 @@ def _add_run_search_configs(self): ConfigField( "run_config_search_mode", flags=["--run-config-search-mode"], - choices=["brute", "quick"], + choices=["brute", "quick", "optuna"], field_type=ConfigPrimitive(str), default_value=DEFAULT_RUN_CONFIG_SEARCH_MODE, description="The search mode for Model Analyzer to find and evaluate" " model configurations. 'brute' will brute force all combinations of" " configuration options. 'quick' will attempt to find a near-optimal" " configuration as fast as possible, but isn't guaranteed to find the" - " best.", + " best. 'optuna' is a more generalized version of quick search allowing " + " the user to quick search over any parameter.", ) ) self._add_config( diff --git a/tests/test_config.py b/tests/test_config.py index 190075dea..3b42b8b05 100755 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1979,7 +1979,7 @@ def test_copy(self): def test_multi_model_search_mode(self): """ - Test that multi-model is only run in quick + Test that multi-model is only run in quick/optuna """ args = [ "model-analyzer", @@ -2011,6 +2011,13 @@ def test_multi_model_search_mode(self): self._evaluate_config(new_args, yaml_content, subcommand="profile") + # Optuna should pass + new_args = list(args) + new_args.append("--run-config-search-mode") + new_args.append("optuna") + + self._evaluate_config(new_args, yaml_content, subcommand="profile") + def test_quick_search_mode(self): """ Test that only legal options are specified in quick search