Skip to main content

Use XGBoost with Metaflow

Question

How can I build and fit an XGBoost model in a Metaflow flow?

Solution

There are two common ways to fit XGBoost models and you can use both with Metaflow. In XGBoost documentation they are referred to as the learning API and the scikit-learn API. This example uses the learning API but you can build flows with either.

1Run Flow

The flow shows how to:

  1. Load training data.
  2. Instantiate the XGBoost model.
  3. Train the model with cross-validation.
xgb_learning_api.py
from metaflow import FlowSpec, step, Parameter

class XGBFlow(FlowSpec):

@step
def start(self):
from sklearn import datasets
self.iris = datasets.load_iris()
self.X = self.iris['data']
self.y = self.iris['target']
self.next(self.train_model)

@step
def train_model(self):
import xgboost as xgb
dtrain = xgb.DMatrix(self.X, self.y)
self.results = xgb.cv(
params = {'num_class':3,
'objective':'multi:softmax',
'eval_metric':"mlogloss"},
dtrain=dtrain,
verbose_eval=False
)
self.next(self.end)

@step
def end(self):
print("Flow is done.")

if __name__ == "__main__":
XGBFlow()
python xgb_learning_api.py run
Workflow starting (run-id 1654221281882630):
[1654221281882630/start/1 (pid 71160)] Task is starting.
[1654221281882630/start/1 (pid 71160)] Task finished successfully.
[1654221281882630/train_model/2 (pid 71199)] Task is starting.
[1654221281882630/train_model/2 (pid 71199)] Task finished successfully.
[1654221281882630/end/3 (pid 71262)] Task is starting.
[1654221281882630/end/3 (pid 71262)] Flow is done.
[1654221281882630/end/3 (pid 71262)] Task finished successfully.
Done!

2Access Artifacts Outside of Flow

The following can be run in a Python script or notebook to access the contents of the dataframe that was stored as a flow artifact with self.results:

from metaflow import Flow
run = Flow('XGBFlow').latest_run
run.data.results.head()
train-mlogloss-mean train-mlogloss-std test-mlogloss-mean test-mlogloss-std
0 0.741877 0.001425 0.750814 0.002562
1 0.533298 0.003306 0.550585 0.001667
2 0.394987 0.002554 0.421669 0.002304
3 0.300281 0.002392 0.337402 0.003478
4 0.231565 0.001567 0.280347 0.004483

Further Reading

For help or feedback, please join Metaflow Slack. To suggest an article, you may open an issue on GitHub.

Join Slack