Skip to main content

Reuse Model Object

Question

How can I reuse model code in training and prediction flows?

Solution

A common pattern when using Metaflow is to move complex business logic outside of the flow. This makes the logic callable from multiple flows and more easily tested independent of the flow.

1Make Class Used in Multiple Flows

Imagine you have the following model class:

model.py
class Model():

def init_model(self, model_type = None, params:dict = {}):
return model_type(**params)

def train(self, model, features, labels):
return model.fit(features, labels)

def score(self, model, features, true_labels):
preds = model.predict(features)
return {
"accuracy": sum(true_labels==preds)/len(true_labels)
}

Now you can use multiple inheritance with this object when we instantiate our FlowSpec class.

2Use Model Class in Training Flow

This flow demonstrates how the Model class functions can be inherited by the flow. The flow shows how to:

  • Instantiate and splits a dataset from scikit-learn.
  • Initialize a model using the previously defined class.
  • Train the model.
  • Score the model on a validation set and prints the result.
train_model_flow.py
from metaflow import step, FlowSpec
from model import Model

class TrainingFlow(FlowSpec, Model):

@step
def start(self):
from sklearn import datasets
from sklearn.model_selection import train_test_split
self.iris = datasets.load_iris()
X, y = self.iris['data'], self.iris['target']
self.labels = self.iris['target_names']
split = train_test_split(X, y, test_size=0.2)
self.X_train, self.X_test = split[0], split[1]
self.y_train, self.y_test = split[2], split[3]
self.next(self.make_model)

@step
def make_model(self):
from sklearn.ensemble import RandomForestClassifier
self.params = {"max_depth": 8}
self.model = self.init_model(
model_type = RandomForestClassifier,
params = self.params
)
self.next(self.train_model)

@step
def train_model(self):
self.model = self.train(self.model, self.X_train, self.y_train)
self.next(self.end)

@step
def end(self):
scores = self.score(self.model, self.X_test, self.y_test)
print('Accuracy: ', scores['accuracy'])

if __name__ == "__main__":
TrainingFlow()
python train_model_flow.py run
     Workflow starting (run-id 864):
[864/start/4336 (pid 17451)] Task is starting.
[864/start/4336 (pid 17451)] Task finished successfully.
[864/make_model/4337 (pid 17461)] Task is starting.
[864/make_model/4337 (pid 17461)] Task finished successfully.
[864/train_model/4338 (pid 17466)] Task is starting.
[864/train_model/4338 (pid 17466)] Task finished successfully.
[864/end/4339 (pid 17475)] Task is starting.
[864/end/4339 (pid 17475)] Accuracy: 0.9
[864/end/4339 (pid 17475)] Task finished successfully.
Done!

3Use Model Class in Scoring Flow

Now you can use multiple inheritance again to instantiate a different flow.

This flow shows how to:

  • Create a test dataset to score.
  • Instantiate a model using the trained model object from TrainFlow.
  • Use the common Model class function to score the model on the test dataset.
scoring_model_flow.py
from metaflow import step, FlowSpec
from model import Model

class ScoringFlow(FlowSpec, Model):

sibling_flow = 'TrainingFlow'

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

@step
def score_trained_model(self):
from metaflow import Flow
run = Flow(self.sibling_flow).latest_successful_run
self.model = run['end'].task.data.model
self.scores = self.score(self.model, self.X, self.y)
self.next(self.end)

@step
def end(self):
print('Accuracy: ', self.scores['accuracy'])

if __name__ == "__main__":
ScoringFlow()
python scoring_model_flow.py run
     Workflow starting (run-id 865):
[865/start/4341 (pid 17487)] Task is starting.
[865/start/4341 (pid 17487)] Task finished successfully.
[865/score_trained_model/4342 (pid 17499)] Task is starting.
[865/score_trained_model/4342 (pid 17499)] Task finished successfully.
[865/end/4343 (pid 17511)] Task is starting.
[865/end/4343 (pid 17511)] Accuracy: 0.98
[865/end/4343 (pid 17511)] Task finished successfully.
Done!

Further Reading

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

Join Slack