Repository: ludwig-ai/ludwig
Pull request: feat: Adding feature type shared parameter capability for hyperopt
This change enables using the `defaults` keyword within hyperopt to set default parameters for feature groups. This will help add search spaces more concisely for datasets with a large number of features while also reducing the search space to allow for a deeper search during hyperopt. 

For e.g., a user can now add 

``` 
"defaults.text.num_filters": {"space": "choice", "categories": [64, 128, 256, 512]}
```

in their hyperopt search space so that all text features (inputs and outputs) use the sampled `cell_type` for that particular trial. 

The API for this change is as follows: `defaults.<feature_type>.<parameter>`, where:

- `feature_type`: a valid input or output feature type, for e.g., text/category/sequence etc.
- `parameter`: the parameter to search over for that particular feature type, for e.g., `embedding_size` for category features

Example Configuration:

```yaml
hyperopt:
    goal: minimize
    parameters:
        trainer.learning_rate:
            lower: 0.0001
            upper: 0.01
            space: loguniform
        defaults.text.num_filters:
            space: choice
            categories: [64, 128, 256, 512]
        defaults.category.embedding_size:
            space: choice
            categories: [64, 128, 256, 512]
        ....
    ....
```

_Note: The output `status` that gets logged every few seconds will show the parameter as `defaults.text.cell_type` in the printed table with each trial status, but underneath the surface, this is getting updated correctly. This is unfortunately not something we can control, but I think it is okay to leave it this way anyway because the table width would become too large if we were to add a column for each feature whose default parameter was updated in the trial status table that gets printed.

<img width="1563" alt="Screen Shot 2022-06-13 at 1 12 33 PM" src="https://user-images.githubusercontent.com/106701836/173437112-f8f423e7-735c-48f0-b90f-4d983cb7b1fe.png">_ 

Problem statement:
[Hyperopt] Allow default shared parameter search spaces
It would be helpful to enable using a new `defaults` keyword within hyperopt to set default parameters for feature groups. This will help add search spaces more concisely for datasets with a large number of features while also reducing the search space to allow for a deeper search during hyperopt.

For e.g., a user should be able to add

`"defaults.input_features.text.cell_type": {"space": "choice", "categories": ["lstm", "gru"]}`

in their hyperopt search space so that all text input features use the sampled cell_type for that particular trial.

The API for this change will be as follows: `defaults.<feature_group>.<feature_type>.<parameter>`, where:

- `feature_group`: either input_features or output_features to distinguish which features to apply this to
- `feature_type`: the feature type in question, for e.g., text/category/sequence etc.
- `parameter`: the parameter that one wants to search over for that particular feature type, for e.g., vocab_size for text

Example Configuration:

```yaml
hyperopt:
    goal: minimize
    parameters:
        trainer.learning_rate:
            lower: 0.0001
            upper: 0.01
            space: loguniform
        defaults.input_features.text.cell_type:
            space: choice
            categories: ["lstm", "rnn", "gru"]
        defaults.output_features.category.vocab_size:
            space: randint
            lower: 4
            upper: 8
        ....
    ....
```

Hints from issue discussion:

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/ludwig/constants.py b/ludwig/constants.py
index 23355f6cf67..e299d42c0b8 100644
--- a/ludwig/constants.py
+++ b/ludwig/constants.py
@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+INPUT_FEATURES = "input_features"
+OUTPUT_FEATURES = "output_features"
+
 BINARY = "binary"
 CATEGORY = "category"
 INT = "int"
@@ -150,6 +153,10 @@
 
 COMBINER = "combiner"
 
+ENCODER = "encoder"
+
+DEFAULTS = "defaults"
+
 BALANCE_PERCENTAGE_TOLERANCE = 0.03
 IMBALANCE_DETECTION_RATIO = 0.05
 
diff --git a/ludwig/hyperopt/execution.py b/ludwig/hyperopt/execution.py
index 89aea9861e6..b3551e7e970 100644
--- a/ludwig/hyperopt/execution.py
+++ b/ludwig/hyperopt/execution.py
@@ -16,7 +16,20 @@
 from ludwig.api import LudwigModel
 from ludwig.backend import initialize_backend, RAY
 from ludwig.callbacks import Callback
-from ludwig.constants import COLUMN, MAXIMIZE, TEST, TRAINER, TRAINING, TYPE, VALIDATION
+from ludwig.constants import (
+    COLUMN,
+    COMBINER,
+    DEFAULTS,
+    INPUT_FEATURES,
+    MAXIMIZE,
+    OUTPUT_FEATURES,
+    PREPROCESSING,
+    TEST,
+    TRAINER,
+    TRAINING,
+    TYPE,
+    VALIDATION,
+)
 from ludwig.hyperopt.results import HyperoptResults, RayTuneResults, TrialResults
 from ludwig.hyperopt.sampling import RayTuneSampler
 from ludwig.hyperopt.search_algos import get_search_algorithm
@@ -192,6 +205,7 @@ def execute(
         backend=None,
         random_seed=default_random_seed,
         debug=False,
+        shared_params_dict=None,
         **kwargs,
     ) -> HyperoptResults:
         pass
@@ -348,19 +362,22 @@ def _evaluate_best_model(
                 f"{traceback.format_exc()}"
             )
 
-    def _run_experiment(self, config, checkpoint_dir, hyperopt_dict, decode_ctx, is_using_ray_backend=False):
+    def _run_experiment(
+        self, config, checkpoint_dir, hyperopt_dict, decode_ctx, shared_params_dict, is_using_ray_backend=False
+    ):
         for gpu_id in ray.get_gpu_ids():
             # Previous trial may not have freed its memory yet, so wait to avoid OOM
             wait_for_gpu(gpu_id)
+
         # Some config values may be JSON encoded as strings, so decode them here
         config = RayTuneSampler.decode_values(config, decode_ctx)
 
         trial_id = tune.get_trial_id()
-        modified_config = substitute_parameters(copy.deepcopy(hyperopt_dict["config"]), config)
-
         trial_dir = Path(tune.get_trial_dir())
         trial_location = ray.util.get_node_ip_address()
 
+        modified_config = substitute_parameters(copy.deepcopy(hyperopt_dict["config"]), config, shared_params_dict)
+
         hyperopt_dict["config"] = modified_config
         hyperopt_dict["experiment_name "] = f'{hyperopt_dict["experiment_name"]}_{trial_id}'
         hyperopt_dict["output_directory"] = str(trial_dir)
@@ -571,6 +588,7 @@ def execute(
         random_seed=default_random_seed,
         debug=False,
         hyperopt_log_verbosity=3,
+        shared_params_dict=None,
         **kwargs,
     ) -> RayTuneResults:
         if isinstance(dataset, str) and not has_remote_protocol(dataset) and not os.path.isabs(dataset):
@@ -662,7 +680,12 @@ def execute(
 
         def run_experiment_trial(config, local_hyperopt_dict, checkpoint_dir=None):
             return self._run_experiment(
-                config, checkpoint_dir, local_hyperopt_dict, self.decode_ctx, _is_ray_backend(backend)
+                config,
+                checkpoint_dir,
+                local_hyperopt_dict,
+                self.decode_ctx,
+                shared_params_dict,
+                _is_ray_backend(backend),
             )
 
         tune_config = {}
@@ -805,9 +828,36 @@ def get_build_hyperopt_executor(executor_type):
 executor_registry = {"ray": RayTuneExecutor}
 
 
-def set_values(model_dict, name, parameters_dict):
-    if name in parameters_dict:
-        params = parameters_dict[name]
+def set_values(model_dict, feature_name, parameters_dict, feature_type=None, shared_params_dict=None):
+    """Updates the parameters of feature_name in model_dict based on hyperopt parameters sampled for each trial
+    stored in parameters_dict.
+
+    # Inputs
+
+    :param model_dict: (dict) contains the parameters for the specific input feature from config
+    :param feature_name: (str) name of the feature to update
+    :param parameters_dict: (dict) parameters from the hyperopt sampler to be used to overwrite
+         default params in model_dict
+    :param feature_type: (str, default: none) type of the feature to update
+    :param shared_params_dict: (dict, default: none) mapping of feature type to feature names that
+         can be considered for shared parameter updates based on original config
+    """
+
+    # Update shared params
+    if DEFAULTS in parameters_dict and shared_params_dict:
+        if feature_type in shared_params_dict and feature_name in shared_params_dict[feature_type]:
+            if feature_type in parameters_dict[DEFAULTS]:
+                shared_params = parameters_dict[DEFAULTS][feature_type]
+                for key, value in shared_params.items():
+                    if isinstance(value, dict):
+                        for sub_key, sub_value in value.items():
+                            model_dict[key][sub_key] = sub_value
+                    else:
+                        model_dict[key] = value
+
+    # Update or overwrite any feature specific hyperopt params
+    if feature_name in parameters_dict:
+        params = parameters_dict[feature_name]
         for key, value in params.items():
             if isinstance(value, dict):
                 for sub_key, sub_value in value.items():
@@ -831,15 +881,21 @@ def get_parameters_dict(parameters):
     return parameters_dict
 
 
-def substitute_parameters(config, parameters):
+def substitute_parameters(config, parameters, shared_params_dict):
     parameters_dict = get_parameters_dict(parameters)
-    for input_feature in config["input_features"]:
-        set_values(input_feature, input_feature[COLUMN], parameters_dict)
-    for output_feature in config["output_features"]:
+    for input_feature in config[INPUT_FEATURES]:
+        set_values(
+            input_feature,
+            input_feature[COLUMN],
+            parameters_dict,
+            feature_type=input_feature[TYPE],
+            shared_params_dict=shared_params_dict,
+        )
+    for output_feature in config[OUTPUT_FEATURES]:
         set_values(output_feature, output_feature[COLUMN], parameters_dict)
-    set_values(config["combiner"], "combiner", parameters_dict)
+    set_values(config[COMBINER], COMBINER, parameters_dict)
     set_values(config[TRAINER], TRAINER, parameters_dict)
-    set_values(config["preprocessing"], "preprocessing", parameters_dict)
+    set_values(config[PREPROCESSING], PREPROCESSING, parameters_dict)
     return config
 
 
diff --git a/ludwig/hyperopt/run.py b/ludwig/hyperopt/run.py
index 4734b58c56c..f2e13b516bb 100644
--- a/ludwig/hyperopt/run.py
+++ b/ludwig/hyperopt/run.py
@@ -9,7 +9,21 @@
 from ludwig.api import LudwigModel
 from ludwig.backend import Backend, initialize_backend, LocalBackend
 from ludwig.callbacks import Callback
-from ludwig.constants import COMBINED, EXECUTOR, HYPEROPT, LOSS, MINIMIZE, RAY, TEST, TRAINING, TYPE, VALIDATION
+from ludwig.constants import (
+    COMBINED,
+    ENCODER,
+    EXECUTOR,
+    HYPEROPT,
+    INPUT_FEATURES,
+    LOSS,
+    MINIMIZE,
+    NAME,
+    RAY,
+    TEST,
+    TRAINING,
+    TYPE,
+    VALIDATION,
+)
 from ludwig.features.feature_registries import output_type_registry
 from ludwig.hyperopt.execution import executor_registry, get_build_hyperopt_executor, RayTuneExecutor
 from ludwig.hyperopt.results import HyperoptResults
@@ -169,6 +183,10 @@ def hyperopt(
     else:
         config_dict = config
 
+    # Get mapping of features that don't have an encoder for shared parameters
+    # TODO: Consider output features as well
+    shared_params_dict = get_shared_params_dict(config_dict[INPUT_FEATURES])
+
     # merge config with defaults
     config = merge_with_defaults(config_dict)
 
@@ -332,6 +350,7 @@ def hyperopt(
         backend=backend,
         random_seed=random_seed,
         hyperopt_log_verbosity=hyperopt_log_verbosity,
+        shared_params_dict=shared_params_dict,
         **kwargs,
     )
 
@@ -373,3 +392,31 @@ def update_hyperopt_params_with_defaults(hyperopt_params):
         hyperopt_params[EXECUTOR],
         executor_defaults,
     )
+
+
+def get_shared_params_dict(features):
+    """This returns a mapping of feature type to list of features without an encoder. They may be considered for
+    potential shared parameter search spaces depending on the parameter space defined later within the hyperopt
+    config.
+
+    # Inputs
+
+    :param features: (dict) config which defines the different feature
+    related parameters of the model (either input or output)
+
+    # Returns
+
+    :return: feature_group_to_features_map (dict) that maps each feature
+    type to a set of feature names without encoders
+    """
+    feature_group_to_features_map = {}
+
+    for feature in features:
+        if not feature.get(ENCODER, 0):
+            feature_name = feature[NAME]
+            feature_type = feature[TYPE]
+            if feature_type not in feature_group_to_features_map:
+                feature_group_to_features_map[feature_type] = set()
+            feature_group_to_features_map.get(feature_type).add(feature_name)
+
+    return feature_group_to_features_map
diff --git a/tests/integration_tests/test_hyperopt.py b/tests/integration_tests/test_hyperopt.py
index 373efc207f5..11d8871b60c 100644
--- a/tests/integration_tests/test_hyperopt.py
+++ b/tests/integration_tests/test_hyperopt.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import contextlib
+import json
 import logging
 import os.path
 from typing import Dict, Optional, Tuple
@@ -21,7 +22,7 @@
 import ray
 import torch
 
-from ludwig.constants import ACCURACY, RAY, TRAINER
+from ludwig.constants import ACCURACY, INPUT_FEATURES, RAY, TEXT, TRAINER, TYPE
 from ludwig.hyperopt.execution import get_build_hyperopt_executor
 from ludwig.hyperopt.results import HyperoptResults, RayTuneResults
 from ludwig.hyperopt.run import hyperopt, update_hyperopt_params_with_defaults
@@ -314,3 +315,96 @@ def test_hyperopt_run_hyperopt(csv_filename, search_space, tmpdir, ray_cluster):
 
     if os.path.isfile(os.path.join(tmpdir, "test_hyperopt", "hyperopt_statistics.json")):
         os.remove(os.path.join(tmpdir, "test_hyperopt", "hyperopt_statistics.json"))
+
+
+@pytest.mark.distributed
+@pytest.mark.parametrize("search_space", ["random"])
+def test_hyperopt_run_hyperopt_with_shared_params(csv_filename, search_space, tmpdir, ray_cluster):
+    input_features = [
+        text_feature(name="title", cell_type="rnn", reduce_output="sum"),
+        text_feature(name="body", cell_type="rnn"),
+        category_feature(vocab_size=2, reduce_input="sum"),
+        category_feature(vocab_size=6),
+    ]
+
+    output_features = [category_feature(vocab_size=2, reduce_input="sum")]
+
+    rel_path = generate_data(input_features, output_features, csv_filename)
+
+    config = {
+        "input_features": input_features,
+        "output_features": output_features,
+        "combiner": {"type": "concat", "num_fc_layers": 2},
+        TRAINER: {"epochs": 2, "learning_rate": 0.001},
+    }
+
+    categorical_feature_name = input_features[2]["name"]
+    output_feature_name = output_features[0]["name"]
+
+    cell_types = ["lstm", "gru"]
+
+    # Create default search space for text features with various cell types
+    search_parameters = {
+        "trainer.learning_rate": {
+            "lower": 0.0001,
+            "upper": 0.01,
+            "space": "loguniform",
+        },
+        categorical_feature_name + ".vocab_size": {"space": "randint", "lower": 2, "upper": 8},
+        "defaults.text.cell_type": {"space": "choice", "categories": cell_types},
+    }
+
+    hyperopt_configs = {
+        "parameters": search_parameters,
+        "goal": "minimize",
+        "output_feature": output_feature_name,
+        "validation_metrics": "loss",
+        "executor": {"type": "ray", "num_samples": 1 if search_space == "grid" else RANDOM_SEARCH_SIZE},
+        "search_alg": {"type": "variant_generator"},
+    }
+
+    # add hyperopt parameter space to the config
+    config["hyperopt"] = hyperopt_configs
+
+    hyperopt_results = hyperopt(config, dataset=rel_path, output_directory=tmpdir, experiment_name="test_hyperopt")
+
+    # check for return results
+    assert isinstance(hyperopt_results, HyperoptResults)
+
+    hyperopt_results_df = hyperopt_results.experiment_analysis.results_df
+
+    # check if trials match random_search_size
+    assert hyperopt_results_df.shape[0] == RANDOM_SEARCH_SIZE
+
+    # Check that the trials did sample from defaults in the search space
+    for _, trial_row in hyperopt_results_df.iterrows():
+        cell_type = trial_row["config.defaults.text.cell_type"].replace('"', "")
+        assert cell_type in cell_types
+
+    # check for existence of the hyperopt statistics file
+    assert os.path.isfile(os.path.join(tmpdir, "test_hyperopt", "hyperopt_statistics.json"))
+
+    # Check that each trial's text input configs got updated correctly
+    for _, trial_row in hyperopt_results_df.iterrows():
+        trial_dir = trial_row["trial_dir"]
+        parameters_file_path = os.path.join(trial_dir, "test_hyperopt_run", "model", "model_hyperparameters.json")
+        try:
+            params_fd = open(parameters_file_path)
+            model_parameters = json.load(params_fd)
+            input_features = model_parameters[INPUT_FEATURES]
+            text_input_cell_types = set()  # Used to track that all text features have the same cell_type
+            for input_feature in input_features:
+                if input_feature[TYPE] == TEXT:
+                    cell_type = input_feature["cell_type"]
+                    # Check that cell type got updated from the sampler
+                    assert cell_type in cell_types
+                    text_input_cell_types.add(cell_type)
+            # All text features with defaults should have the same cell type for this trial
+            assert len(text_input_cell_types) == 1
+            params_fd.close()
+        # Likely unable to open trial dir so fail this test
+        except OSError:
+            assert 1 == 0
+
+    if os.path.isfile(os.path.join(tmpdir, "test_hyperopt", "hyperopt_statistics.json")):
+        os.remove(os.path.join(tmpdir, "test_hyperopt", "hyperopt_statistics.json"))


Code review comments:
diff_hunk: @@ -805,9 +828,36 @@ def get_build_hyperopt_executor(executor_type):
 executor_registry = {"ray": RayTuneExecutor}
 
 
-def set_values(model_dict, name, parameters_dict):
-    if name in parameters_dict:
-        params = parameters_dict[name]
+def set_values(model_dict, feature_name, parameters_dict, feature_type=None, shared_params_dict=None):
+    """Updates the parameters of feature_name in model_dict based on hyperopt parameters sampled for each trial
+    stored in parameters_dict.
+
+    # Inputs
+
+    :param model_dict: (dict) contains the parameters for the specific input feature from config
path: ludwig/hyperopt/execution.py
text: @user1:
Thanks for adding docs!

Instead of adding type hints for each feature type in the documentation, can you add type hints to the function definitions?

diff_hunk: @@ -373,3 +392,31 @@ def update_hyperopt_params_with_defaults(hyperopt_params):
         hyperopt_params[EXECUTOR],
         executor_defaults,
     )
+
+
+def get_shared_params_dict(features):
+    """This returns a mapping of feature type to list of features without an encoder. They may be considered for
+    potential shared parameter search spaces depending on the parameter space defined later within the hyperopt
+    config.
+
+    # Inputs
+
+    :param features: (dict) config which defines the different feature
+    related parameters of the model (either input or output)
+
+    # Returns
+
+    :return: feature_group_to_features_map (dict) that maps each feature
+    type to a set of feature names without encoders
+    """
+    feature_group_to_features_map = {}
+
+    for feature in features:
+        if not feature.get(ENCODER, 0):
+            feature_name = feature[NAME]
+            feature_type = feature[TYPE]
+            if feature_type not in feature_group_to_features_map:
+                feature_group_to_features_map[feature_type] = set()
path: ludwig/hyperopt/run.py
text: @user1:
You can use a `defaultdict(set)` instead of an empty dictionary for `feature_group_to_features_map` and then remove the check to see if a feature type exists.