Repository: aimhubio/aim

Issue:
Xgboost integration
Add Xgboost integration


Discussion:

Resolving patch:
diff --git a/aim/sdk/adapters/xgboost.py b/aim/sdk/adapters/xgboost.py
new file mode 100644
--- /dev/null
+++ b/aim/sdk/adapters/xgboost.py
@@ -0,0 +1,86 @@
+from typing import Optional
+
+from aim.sdk.session.session import Session
+from aim.engine.configs import DEFAULT_SYSTEM_TRACKING_INT
+from aim.sdk.session.configs import DEFAULT_FLUSH_FREQUENCY
+
+
+class AimCallback(object):
+    __xgboost_callback_cls = None
+
+    @staticmethod
+    def __new__(cls, *args, **kwargs):
+        xgboost_callback_inst = cls.__get_callback_cls()
+        return xgboost_callback_inst(*args, **kwargs)
+
+    @classmethod
+    def __get_callback_cls(cls):
+        if cls.__xgboost_callback_cls is not None:
+            return cls.__xgboost_callback_cls
+        
+        from xgboost.callback import TrainingCallback, CallbackContainer
+
+        class _XgboostCallback(TrainingCallback):
+            def __init__(self, repo: Optional[str] = None,
+                         experiment: Optional[str] = None,
+                         system_tracking_interval: Optional[int]
+                         = DEFAULT_SYSTEM_TRACKING_INT,
+                         flush_frequency: Optional[int]
+                         = DEFAULT_FLUSH_FREQUENCY):
+                super().__init__()
+                self.repo = repo
+                self.experiment = experiment
+                self.flush_frequency = flush_frequency
+                self.system_tracking_interval = system_tracking_interval
+                self.initialized = False
+                self.aim_session = None
+
+            def before_training(self, model):
+                self.aim_session = Session(
+                    repo=self.repo,
+                    experiment=self.experiment,
+                    flush_frequency=self.flush_frequency,
+                    system_tracking_interval=self.system_tracking_interval
+                )
+                self.initialized = True
+                return model
+
+            def after_iteration(self, model, epoch: int,
+                                evals_log: CallbackContainer.EvalsLog) -> bool:
+                if not evals_log:
+                    return False
+
+                for data, metric in evals_log.items():
+                    for metric_name, log in metric.items():
+                        stdv: Optional[float] = None
+                        if isinstance(log[-1], tuple):
+                            score = log[-1][0]
+                            stdv = log[-1][1]
+                        else:
+                            score = log[-1]
+
+                        self.aim_session.track(score,
+                                               name=metric_name,
+                                               stdv=False)
+                        if stdv is not None:
+                            self.aim_session.track(score,
+                                                   name=metric_name,
+                                                   stdv=True)
+
+                return False
+
+            def after_training(self, model):
+                if self.initialized and self.aim_session.active:
+                    self.aim_session.close()
+                return model
+
+        cls.__xgboost_callback_cls = _XgboostCallback
+        return cls.__xgboost_callback_cls
+    
+    def __init__(self,
+                 repo: Optional[str] = None,
+                 experiment: Optional[str] = None,
+                 system_tracking_interval: Optional[int]
+                 = DEFAULT_SYSTEM_TRACKING_INT,
+                 flush_frequency: Optional[int] = DEFAULT_FLUSH_FREQUENCY):
+        pass
diff --git a/aim/xgboost.py b/aim/xgboost.py
new file mode 100644
--- /dev/null
+++ b/aim/xgboost.py
@@ -0,0 +1,2 @@
+# Alias to SDK XGBoost interface
+from aim.sdk.adapters.xgboost import AimCallback
diff --git a/examples/xgboost_track.py b/examples/xgboost_track.py
new file mode 100644
--- /dev/null
+++ b/examples/xgboost_track.py
@@ -0,0 +1,54 @@
+
+# You should download and extract the data beforehand. Simply by doing this: 
+# wget https://archive.ics.uci.edu/ml/machine-learning-databases/dermatology/dermatology.data
+
+from __future__ import division
+
+import numpy as np
+import xgboost as xgb
+from aim.xgboost import AimCallback
+
+# label need to be 0 to num_class -1
+data = np.loadtxt('./dermatology.data', delimiter=',',
+        converters={33: lambda x:int(x == '?'), 34: lambda x:int(x) - 1})
+sz = data.shape
+
+train = data[:int(sz[0] * 0.7), :]
+test = data[int(sz[0] * 0.7):, :]
+
+train_X = train[:, :33]
+train_Y = train[:, 34]
+
+test_X = test[:, :33]
+test_Y = test[:, 34]
+print(len(train_X))
+
+xg_train = xgb.DMatrix(train_X, label=train_Y)
+xg_test = xgb.DMatrix(test_X, label=test_Y)
+# setup parameters for xgboost
+param = {}
+# use softmax multi-class classification
+param['objective'] = 'multi:softmax'
+# scale weight of positive examples
+param['eta'] = 0.1
+param['max_depth'] = 6
+param['nthread'] = 4
+param['num_class'] = 6
+
+watchlist = [(xg_train, 'train'), (xg_test, 'test')]
+num_round = 50
+bst = xgb.train(param, xg_train, num_round, watchlist)
+# get prediction
+pred = bst.predict(xg_test)
+error_rate = np.sum(pred != test_Y) / test_Y.shape[0]
+print('Test error using softmax = {}'.format(error_rate))
+
+# do the same thing again, but output probabilities
+param['objective'] = 'multi:softprob'
+bst = xgb.train(param, xg_train, num_round, watchlist, callbacks=[AimCallback(repo='.', experiment='xgboost_test')])
+# Note: this convention has been changed since xgboost-unity
+# get prediction, this is in 1D array, need reshape to (ndata, nclass)
+pred_prob = bst.predict(xg_test).reshape(test_Y.shape[0], 6)
+pred_label = np.argmax(pred_prob, axis=1)
+error_rate = np.sum(pred_label != test_Y) / test_Y.shape[0]
+print('Test error using softprob = {}'.format(error_rate))


Test patch: