A Simple ML Model Base Class

A software design pattern for machine learning code

Scikit-learn’s Approach to Base Classes

The most well known ML software package in python is scikit-learn, and it provides a set of abstract base classes in the base.py module. The scikit-learn API is a great place to learn about machine learning software engineering in general, but in this case we want to focus on it’s approach to base classes for making predictions with ML models.

  • How to load serialized model assets when instantiating a model
  • How to publish input and output data schema information

Some Solutions

Over the last few years, a few big tech companies have been developing proprietary in-house machine learning infrastructure and software. Some of these companies sell access to their ML platform and others have published details about their approach to ML infrastructure. Also, there have been a few open source projects that seek to simplify the deployment of ML models to production systems. In this section I will describe some solutions that have emerged recently for the problems described above.

AWS Sagemaker

AWS Sagemaker is a platform for training and deploying ML models within the AWS ecosystem. The platform has several ready-made ML algorithms that can be leveraged without writing a lot of code. However, a way to deploy custom ML code to the platform is provided. To deploy a prediction endpoint on top of the Sagemaker service, a Python Flask application with a “/ping” and “/invocations” endpoints must be created and deployed within a Docker container.


Facebook published a blog post about their ML systems here. The FBLearner Flow system is made up of workflows and operators. A workflow is a single unit of work with a specific set of inputs and outputs, a workflow is made up of operators which do simple operations on data. The blog post shows how to train a Decision Tree model on the iris data set. The blog post does not provide many implementation details about their internal Python packages. An interesting part of the approach taken is the fact that schema metadata is attached to every workflow created, ensuring type safety at runtime. There are not details about loading and storing model assets. Facebook’s FBFlow Python package does not use base classes that developers can inherit from to write code, but uses function annotations to attach metadata to ML model code.


Uber published a blog post about their approach to custom ML models here. Uber’s PyML package is used to deploy ML models that are not natively supported by Uber’s Michelangelo ML platform, which is described here. The PyML package does not specify how to write model training code, but does provide a base class for writing ML model prediction code. The base class is called DataFrameModel. The interface is very simple, it only has two methods: the __init__() method, and the predict() method. The model assets are required to be deserialized in the class constructor and all prediction code is in the predict method of the class.

Seldon Core

Seldon Core is an open source project for hosting ML models. It supports custom Python models, as described here. The model code is required to be in a Python class with an __init__() method and a predict() method, it follows Uber’s design closely but does not use an abstract base class to enforce the interface. Another difference is that Seldon allows the model class to return results in several different ways, and not just in Pandas dataframes. Seldon also allows the model class to return column name metadata for the model inputs, but no type metadata.

A Simple ML Model Base Class

NOTE: All of the code shown in this section can be found in this Github repository.

class MLModel(ABC):
“”” An abstract base class for ML model prediction code “””
def input_schema(self):
raise NotImplementedError()
def output_schema(self):
raise NotImplementedError()
def __init__(self):
raise NotImplementedError()
def predict(self, data):
class MLModelException(Exception):
“”” Exception type for use within MLModel derived classes “””
def __init__(self,*args,**kwargs):
Exception.__init__(self, *args, **kwargs)

Using the Base Class

This blog post deals purely with the ML code that will be used for predicting in production and not with the model training code. However, we still need to have a model to work with. Here’s a simple scikit-learn model training script:

iris = datasets.load_iris()svm_model = svm.SVC(gamma=0.001, C=100.0)
svm_model.fit(iris.data[:-1], iris.target[:-1])
dir_path = os.path.dirname(os.path.realpath(__file__))
file = open(os.path.join(dir_path, “model_files”, “svc_model.pickle”), ‘wb’)
pickle.dump(svm_model, file)
class IrisSVCModel(MLModel):
“”” A demonstration of how to use “””
input_schema = Schema({‘sepal_length’: float,
‘sepal_width’: float,
‘petal_length’: float,
‘petal_width’: float})
# the output of the model will be one of three strings
output_schema = Schema({‘species’: Or(“setosa”,
def __init__(self):
dir_path = os.path.dirname(os.path.realpath(__file__))
file = open(os.path.join(dir_path, “model_files”, “svc_model.pickle”), ‘rb’)
self._svm_model = pickle.load(file)
def predict(self, data):
# calling the super method to validate against the
# input_schema
# converting the incoming dictionary into a numpy array
# that can be accepted by the scikit-learn model
X = array([data[“sepal_length”],
data[“petal_width”]]).reshape(1, -1)
# making the prediction
y_hat = int(self._svm_model.predict(X)[0])
# converting the prediction into a string that will match
# the output schema of the model, this list will map the
# output of the scikit-learn model to the string expected by
# the output schema
targets = [‘setosa’, ‘versicolor’, ‘virginica’]
species = targets[y_hat]
return {“species”: species}
>>> model = IrisSVCModel()>>> print(json.dumps(model.input_schema.json_schema("https://example.com   /my-schema.json")))

{"type": "object", "properties": {"sepal_length": {"type": "number"}, "sepal_width": {"type": "number"},


In this post I showed a few different approaches to deploying ML model code to production systems. I also showed an implementation of a Python base class that brings together the best features of the different approaches discussed. In conclusion I will discuss some of the benefits of the approach I sketched out above.

Coder and machine learning enthusiast

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store