feat: Added model type GBM (LightGBM tree learner), as an alternative to ECD

examples/lightgbm/config.yaml

@@ -0,0 +1,35 @@
+model_type: gbm
+input_features:
+  - name: age
+    type: number
+  - name: workclass
+    type: category
+  - name: fnlwgt
+    type: number
+  - name: education
+    type: category
+  - name: education-num
+    type: number
+  - name: marital-status
+    type: category
+  - name: occupation
+    type: category
+  - name: relationship
+    type: category
+  - name: race
+    type: category
+  - name: sex
+    type: category
+  - name: capital-gain
+    type: number
+  - name: capital-loss
+    type: number
+  - name: hours-per-week
+    type: number
+  - name: native-country
+    type: category
+output_features:
+  - name: income
+    type: category
+training:
+  learning_rate: 0.05

examples/lightgbm/config_higgs.yaml

@@ -0,0 +1,70 @@
+model_type: gbm
+input_features:
+  - name: lepton_pT
+    type: number
+  - name: lepton_eta
+    type: number
+  - name: lepton_phi
+    type: number
+  - name: missing_energy_magnitude
+    type: number
+  - name: missing_energy_phi
+    type: number
+  - name: jet_1_pt
+    type: number
+  - name: jet_1_eta
+    type: number
+  - name: jet_1_phi
+    type: number
+  - name: jet_1_b-tag
+    type: number
+  - name: jet_2_pt
+    type: number
+  - name: jet_2_eta
+    type: number
+  - name: jet_2_phi
+    type: number
+  - name: jet_2_b-tag
+    type: number
+  - name: jet_3_pt
+    type: number
+  - name: jet_3_eta
+    type: number
+  - name: jet_3_phi
+    type: number
+  - name: jet_3_b-tag
+    type: number
+  - name: jet_4_pt
+    type: number
+  - name: jet_4_eta
+    type: number
+  - name: jet_4_phi
+    type: number
+  - name: jet_4_b-tag
+    type: number
+  - name: m_jj
+    type: number
+  - name: m_jjj
+    type: number
+  - name: m_lv
+    type: number
+  - name: m_jlv
+    type: number
+  - name: m_bb
+    type: number
+  - name: m_wbb
+    type: number
+  - name: m_wwbb
+    type: number
+output_features:
+  - name: label
+    type: binary
+trainer:
+  learning_rate: 0.1
+  num_boost_round: 50
+  num_leaves: 255
+  feature_fraction: 0.9
+  bagging_fraction: 0.8
+  bagging_freq: 5
+  min_data_in_leaf: 1
+  min_sum_hessian_in_leaf: 100

examples/lightgbm/train.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+
+import logging
+import os
+import shutil
+
+from ludwig.api import LudwigModel
+from ludwig.backend import initialize_backend
+from ludwig.data.cache.types import CacheableDataframe
+from ludwig.datasets import higgs  # adult_census_income
+
+shutil.rmtree("./results", ignore_errors=True)
+
+backend_config = {
+    "type": "ray",
+    "processor": {
+        "parallelism": 6,
+        "type": "dask",
+    },
+    "trainer": {
+        "num_actors": 3,
+        "cpus_per_actor": 2,
+    },
+}
+backend = initialize_backend(backend_config)
+model = LudwigModel(config="./config_higgs.yaml", logging_level=logging.INFO, backend=backend)
+
+# df = adult_census_income.load(split=False)
+df = higgs.load(split=False, add_validation_set=True)
+df = CacheableDataframe(df=df, name="cache_higgs", checksum="9YeB0J_fiQ9Dh_lL84NdZg==")
+
+(
+    train_stats,  # dictionary containing training statistics
+    preprocessed_data,  # tuple Ludwig Dataset objects of pre-processed training data
+    output_directory,  # location of training results stored on disk
+) = model.train(dataset=df)
+
+print("contents of output directory:", output_directory)
+for item in os.listdir(output_directory):
+    print("\t", item)

ludwig/api.py

@@ -43,6 +43,7 @@
     EVAL_BATCH_SIZE,
     FULL,
     LEARNING_RATE,
+    MODEL_TYPE,
     PREPROCESSING,
     TEST,
     TRAINER,
@@ -61,18 +62,19 @@
     set_disable_progressbar,
     TRAIN_SET_METADATA_FILE_NAME,
 )
-from ludwig.models.ecd import ECD
+from ludwig.models.abstractmodel import AbstractModel
 from ludwig.models.inference import InferenceModule
 from ludwig.models.predictor import (
     calculate_overall_stats,
     print_evaluation_stats,
     save_evaluation_stats,
     save_prediction_outputs,
 )
-from ludwig.models.trainer import Trainer
+from ludwig.models.registry import model_type_registry
 from ludwig.modules.metric_modules import get_best_function
 from ludwig.schema import validate_config
 from ludwig.schema.utils import load_config_with_kwargs
+from ludwig.trainers.trainer import Trainer
 from ludwig.utils import metric_utils
 from ludwig.utils.data_utils import (
     figure_data_format,
@@ -84,7 +86,7 @@
 )
 from ludwig.utils.defaults import default_random_seed, merge_with_defaults
 from ludwig.utils.fs_utils import makedirs, open_file, path_exists, upload_output_directory
-from ludwig.utils.misc_utils import get_file_names, get_output_directory
+from ludwig.utils.misc_utils import get_file_names, get_from_registry, get_output_directory
 from ludwig.utils.print_utils import print_boxed
 
 logger = logging.getLogger(__name__)
@@ -1470,8 +1472,8 @@ def _check_initialization(self):
             raise ValueError("Model has not been trained or loaded")
 
     @staticmethod
-    def create_model(config: dict, random_seed: int = default_random_seed) -> ECD:
-        """Instantiates Encoder-Combiner-Decoder (ECD) object.
+    def create_model(config: dict, random_seed: int = default_random_seed) -> AbstractModel:
+        """Instantiates AbstractModel object.
 
         # Inputs
         :param config: (dict) Ludwig config
@@ -1480,15 +1482,10 @@ def create_model(config: dict, random_seed: int = default_random_seed) -> ECD:
             splits and any other random function.
 
         # Return
-        :return: (ludwig.models.ECD) Instance of the Ludwig model object.
+        :return: (ludwig.models.AbstractModel) Instance of the Ludwig model object.
         """
-        # todo: support loading other model types based on config
-        return ECD(
-            input_features_def=config["input_features"],
-            combiner_def=config["combiner"],
-            output_features_def=config["output_features"],
-            random_seed=random_seed,
-        )
+        model_type = get_from_registry(config[MODEL_TYPE], model_type_registry)
+        return model_type(**config, random_seed=random_seed)
 
     @staticmethod
     def set_logging_level(logging_level: int) -> None:

ludwig/backend/base.py

@@ -22,7 +22,10 @@
 from ludwig.data.dataframe.pandas import PANDAS
 from ludwig.data.dataset.base import DatasetManager
 from ludwig.data.dataset.pandas import PandasDatasetManager
-from ludwig.models.ecd import ECD
+from ludwig.models.abstractmodel import AbstractModel
+from ludwig.schema.trainer import TrainerConfig
+from ludwig.trainers.registry import trainers_registry
+from ludwig.utils.misc_utils import get_from_registry
 from ludwig.utils.torch_utils import initialize_pytorch
 
 
@@ -54,7 +57,7 @@ def initialize_pytorch(self, *args, **kwargs):
 
     @contextmanager
     @abstractmethod
-    def create_trainer(self, **kwargs):
+    def create_trainer(self, **kwargs) -> "BaseTrainer":  # noqa: F821
         raise NotImplementedError()
 
     @abstractmethod
@@ -106,12 +109,17 @@ class LocalTrainingMixin:
     def initialize_pytorch(self, *args, **kwargs):
         initialize_pytorch(*args, **kwargs)
 
-    def create_trainer(self, **kwargs):
-        from ludwig.models.trainer import Trainer
+    def create_trainer(self, **kwargs) -> "BaseTrainer":  # noqa: F821
+        config: TrainerConfig = kwargs["config"]
+        model: AbstractModel = kwargs["model"]
 
-        return Trainer(**kwargs)
+        trainers_for_model = get_from_registry(model.type(), trainers_registry)
 
-    def create_predictor(self, model: ECD, **kwargs):
+        trainer_cls = get_from_registry(config.type, trainers_for_model)
+
+        return trainer_cls(**kwargs)
+
+    def create_predictor(self, model: AbstractModel, **kwargs):
         from ludwig.models.predictor import Predictor
 
         return Predictor(model, **kwargs)

ludwig/backend/horovod.py

@@ -17,10 +17,11 @@
 import time
 
 from ludwig.backend.base import Backend, LocalPreprocessingMixin
+from ludwig.constants import MODEL_GBM, MODEL_TYPE
 from ludwig.data.dataset.pandas import PandasDatasetManager
 from ludwig.models.ecd import ECD
 from ludwig.models.predictor import Predictor
-from ludwig.models.trainer import Trainer
+from ludwig.trainers.trainer import Trainer
 from ludwig.utils.horovod_utils import initialize_horovod
 from ludwig.utils.torch_utils import initialize_pytorch
 
@@ -36,7 +37,9 @@ def initialize(self):
     def initialize_pytorch(self, *args, **kwargs):
         initialize_pytorch(*args, horovod=self._horovod, **kwargs)
 
-    def create_trainer(self, **kwargs):
+    def create_trainer(self, **kwargs) -> "BaseTrainer":  # noqa: F821
+        if kwargs.get(MODEL_TYPE, "") == MODEL_GBM:
+            raise ValueError("Horovod backend does not support GBM models.")
         return Trainer(horovod=self._horovod, **kwargs)
 
     def create_predictor(self, model: ECD, **kwargs):

ludwig/backend/ray.py

@@ -32,12 +32,16 @@
 from ray.util.dask import ray_dask_get
 
 from ludwig.backend.base import Backend, RemoteTrainingMixin
-from ludwig.constants import NAME, PREPROCESSING, PROC_COLUMN
+from ludwig.constants import MODEL_ECD, MODEL_GBM, NAME, PREPROCESSING, PROC_COLUMN
 from ludwig.data.dataset.ray import RayDataset, RayDatasetManager, RayDatasetShard
+from ludwig.models.abstractmodel import AbstractModel
 from ludwig.models.ecd import ECD
 from ludwig.models.predictor import BasePredictor, get_output_columns, Predictor, RemotePredictor
-from ludwig.models.trainer import BaseTrainer, RemoteTrainer, TrainerConfig
+from ludwig.schema.trainer import TrainerConfig
+from ludwig.trainers.registry import ray_trainers_registry, register_ray_trainer
+from ludwig.trainers.trainer import BaseTrainer, RemoteTrainer
 from ludwig.utils.horovod_utils import initialize_horovod
+from ludwig.utils.misc_utils import get_from_registry
 from ludwig.utils.torch_utils import initialize_pytorch
 
 _ray112 = LooseVersion(ray.__version__) >= LooseVersion("1.12")
@@ -271,6 +275,7 @@ def tune_learning_rate_fn(
         hvd.shutdown()
 
 
+@register_ray_trainer("ray_trainer_v2", MODEL_ECD, default=True)
 class RayTrainerV2(BaseTrainer):
     def __init__(self, model, trainer_kwargs, data_loader_kwargs, executable_kwargs):
         self.model = model.cpu()
@@ -454,6 +459,7 @@ def __init__(self, **kwargs):
         super().__init__(horovod=horovod, **kwargs)
 
 
+@register_ray_trainer("ray_legacy_trainer", MODEL_ECD)
 class RayLegacyTrainer(BaseTrainer):
     def __init__(self, horovod_kwargs, executable_kwargs):
         # TODO ray: make this more configurable by allowing YAML overrides of timeout_s, etc.
@@ -717,17 +723,26 @@ def initialize_pytorch(self, **kwargs):
         initialize_pytorch(gpus=-1)
         self._pytorch_kwargs = kwargs
 
-    def create_trainer(self, model: ECD, **kwargs):
+    def create_trainer(self, model: AbstractModel, **kwargs) -> "BaseTrainer":  # noqa: F821
         executable_kwargs = {**kwargs, **self._pytorch_kwargs}
         if not self._use_legacy:
+            trainers_for_model = get_from_registry(model.type(), ray_trainers_registry)
+
+            config: TrainerConfig = kwargs["config"]
+            trainer_cls = get_from_registry(config.type, trainers_for_model)
+
             # Deep copy to workaround https://github.com/ray-project/ray/issues/24139
-            return RayTrainerV2(
-                model,
-                copy.deepcopy(self._horovod_kwargs),
-                self._data_loader_kwargs,
-                executable_kwargs,
-            )
+            all_kwargs = {
+                "model": model,
+                "trainer_kwargs": copy.deepcopy(self._horovod_kwargs),
+                "data_loader_kwargs": self._data_loader_kwargs,
+                "executable_kwargs": executable_kwargs,
+            }
+            return trainer_cls(**all_kwargs)
         else:
+            if model.name == MODEL_GBM:
+                raise RuntimeError("Legacy trainer not supported for GBM models.")
+
             # TODO: deprecated 0.5
             return RayLegacyTrainer(self._horovod_kwargs, executable_kwargs)
 

ludwig/constants.py

@@ -167,3 +167,7 @@
 )
 
 DEFAULT_AUDIO_TENSOR_LENGTH = 70000
+
+MODEL_TYPE = "model_type"
+MODEL_ECD = "ecd"
+MODEL_GBM = "gbm"

ludwig/data/dataset/pandas.py

@@ -22,7 +22,7 @@
 from ludwig.data.dataset.base import Dataset, DatasetManager
 from ludwig.data.sampler import DistributedSampler
 from ludwig.utils import data_utils
-from ludwig.utils.data_utils import DATA_TRAIN_HDF5_FP, to_numpy_dataset
+from ludwig.utils.data_utils import DATA_TRAIN_HDF5_FP, from_numpy_dataset, to_numpy_dataset
 from ludwig.utils.fs_utils import download_h5
 from ludwig.utils.misc_utils import get_proc_features
 
@@ -34,6 +34,9 @@ def __init__(self, dataset, features, data_hdf5_fp):
         self.size = len(dataset)
         self.dataset = to_numpy_dataset(dataset)
 
+    def to_df(self, features):
+        return from_numpy_dataset({feature.feature_name: self.dataset[feature.proc_column] for feature in features})
+
     def get(self, proc_column, idx=None):
         if idx is None:
             idx = range(self.size)