Add LightGBM support

trains/binding/frameworks/lightgbm_bind.py

@@ -0,0 +1,130 @@
+import sys
+
+import six
+from pathlib2 import Path
+
+from ..frameworks.base_bind import PatchBaseModelIO
+from ..frameworks import _patched_call, WeightsFileHandler, _Empty
+from ..import_bind import PostImportHookPatching
+from ...config import running_remotely
+from ...model import Framework
+
+
+class PatchLIGHTgbmModelIO(PatchBaseModelIO):
+    __main_task = None
+    __patched = None
+
+    @staticmethod
+    def update_current_task(task, **kwargs):
+        PatchLIGHTgbmModelIO.__main_task = task
+        PatchLIGHTgbmModelIO._patch_model_io()
+        PostImportHookPatching.add_on_import('lightgbm', PatchLIGHTgbmModelIO._patch_model_io)
+
+    @staticmethod
+    def _patch_model_io():
+        if PatchLIGHTgbmModelIO.__patched:
+            return
+
+        if 'lightgbm' not in sys.modules:
+            return
+        PatchLIGHTgbmModelIO.__patched = True
+        # noinspection PyBroadException
+        try:
+            import lightgbm as lgb  # noqa
+
+            lgb.Booster.save_model = _patched_call(lgb.Booster.save_model, PatchLIGHTgbmModelIO._save)
+            lgb.train = _patched_call(lgb.train, PatchLIGHTgbmModelIO._train)
+            lgb.Booster = _patched_call(lgb.Booster, PatchLIGHTgbmModelIO._load)
+        except ImportError:
+            pass
+        except Exception:
+            pass
+
+    @staticmethod
+    def _save(original_fn, obj, f, *args, **kwargs):
+        ret = original_fn(obj, f, *args, **kwargs)
+        if not PatchLIGHTgbmModelIO.__main_task:
+            return ret
+
+        if isinstance(f, six.string_types):
+            filename = f
+        elif hasattr(f, 'name'):
+            filename = f.name
+            # noinspection PyBroadException
+            try:
+                f.flush()
+            except Exception:
+                pass
+        else:
+            filename = None
+
+        # give the model a descriptive name based on the file name
+        # noinspection PyBroadException
+        try:
+            model_name = Path(filename).stem
+        except Exception:
+            model_name = None
+        WeightsFileHandler.create_output_model(obj, filename, Framework.lightgbm, PatchLIGHTgbmModelIO.__main_task,
+                                               singlefile=True, model_name=model_name)
+        return ret
+
+    @staticmethod
+    def _load(original_fn, model_file, *args, **kwargs):
+        if isinstance(model_file, six.string_types):
+            filename = model_file
+        elif hasattr(model_file, 'name'):
+            filename = model_file.name
+        elif len(args) == 1 and isinstance(args[0], six.string_types):
+            filename = args[0]
+        else:
+            filename = None
+
+        if not PatchLIGHTgbmModelIO.__main_task:
+            return original_fn(model_file, *args, **kwargs)
+
+        # register input model
+        empty = _Empty()
+        # Hack: disabled
+        if False and running_remotely():
+            filename = WeightsFileHandler.restore_weights_file(empty, filename, Framework.xgboost,
+                                                               PatchLIGHTgbmModelIO.__main_task)
+            model = original_fn(model_file=filename or model_file, *args, **kwargs)
+        else:
+            # try to load model before registering, in case we fail
+            model = original_fn(model_file=model_file, *args, **kwargs)
+            WeightsFileHandler.restore_weights_file(empty, filename, Framework.lightgbm,
+                                                    PatchLIGHTgbmModelIO.__main_task)
+
+        if empty.trains_in_model:
+            # noinspection PyBroadException
+            try:
+                model.trains_in_model = empty.trains_in_model
+            except Exception:
+                pass
+        return model
+
+    @staticmethod
+    def _train(original_fn, *args, **kwargs):
+        def trains_lightgbm_callback():
+            def callback(env):
+                # logging the results to scalars section
+                # noinspection PyBroadException
+                try:
+                    logger = PatchLIGHTgbmModelIO.__main_task.get_logger()
+                    iteration = env.iteration
+                    for data_title, data_series, value, _ in env.evaluation_result_list:
+                        logger.report_scalar(title=data_title, series=data_series, value="{:.6f}".format(value),
+                                             iteration=iteration)
+                except Exception:
+                    pass
+            return callback
+        params, train_set = args
+        kwargs.setdefault("callbacks", []).append(trains_lightgbm_callback())
+        ret = original_fn(params, train_set, **kwargs)
+        if not PatchLIGHTgbmModelIO.__main_task:
+            return ret
+        for k, v in params.items():
+            if isinstance(v, set):
+                params[k] = list(v)
+        PatchLIGHTgbmModelIO.__main_task.connect(params)
+        return ret