

Keras

Training

"""
Example taken from https://github.com/mlflow/mlflow/blob/master/examples/keras/train.py
"""
from __future__ import print_function

import numpy as np
import keras
from keras.datasets import reuters
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation
from keras.preprocessing.text import Tokenizer

# The following import and function call are the only additions to code required
# to automatically log metrics and parameters to MLflow.
import mlflow.keras

mlflow.set_tracking_uri("http://localhost:5000")

mlflow.keras.autolog()

max_words = 1000
batch_size = 32
epochs = 5

print("Loading data...")
(x_train, y_train), (x_test, y_test) = reuters.load_data(
    num_words=max_words, test_split=0.2
)

print(len(x_train), "train sequences")
print(len(x_test), "test sequences")

num_classes = np.max(y_train) + 1
print(num_classes, "classes")

print("Vectorizing sequence data...")
tokenizer = Tokenizer(num_words=max_words)
x_train = tokenizer.sequences_to_matrix(x_train, mode="binary")
x_test = tokenizer.sequences_to_matrix(x_test, mode="binary")
print("x_train shape:", x_train.shape)
print("x_test shape:", x_test.shape)

print(
    "Convert class vector to binary class matrix "
    "(for use with categorical_crossentropy)"
)
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
print("y_train shape:", y_train.shape)
print("y_test shape:", y_test.shape)

print("Building model...")

experiment_name = "test_keras"
if mlflow.get_experiment_by_name(experiment_name) is None:
    mlflow.create_experiment(experiment_name)

with mlflow.start_run(experiment_id=3) as run:
    model = Sequential()
    model.add(Dense(512, input_shape=(max_words,)))
    model.add(Activation("relu"))
    model.add(Dropout(0.5))
    model.add(Dense(num_classes))
    model.add(Activation("softmax"))

    model.compile(
        loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]
    )

    history = model.fit(
        x_train,
        y_train,
        batch_size=batch_size,
        epochs=epochs,
        verbose=1,
        validation_split=0.1,
    )
    score = model.evaluate(x_test, y_test, batch_size=batch_size, verbose=1)
    print("Test score:", score[0])
    print("Test accuracy:", score[1])

    mlflow.keras.log_model(model, "model", registered_model_name="keras_model")

To run it :

python3 -m examples.training.keras

Serving

from typing import List

import pandas as pd

from keras.preprocessing.text import Tokenizer

from serveml.api import ApiBuilder
from serveml.inputs import BasicInput
from serveml.loader import load_mlflow_model
from serveml.predictions import GenericPrediction

# load model
model = load_mlflow_model(
    # MlFlow model path
    "models:/keras_model/1",
    # MlFlow Tracking URI
    "http://localhost:5000",
)


# Implement deserializer for input data
class ReutersNewswireTopic(BasicInput):
    sequence: List[int]


# Implement prediction because this is a bit custom
class CustomKerasApplication(GenericPrediction):
    def _transform_input(self, input: ReutersNewswireTopic):
        """
        Transforms <ReutersNewswireTopic> object to <numpy.array>
        """
        max_words = 1000
        tokenizer = Tokenizer(num_words=max_words)
        x_train = tokenizer.sequences_to_matrix(
            [input.sequence], mode="binary"
        )
        return pd.DataFrame(x_train)


# implement application
app = ApiBuilder(
    CustomKerasApplication(model), ReutersNewswireTopic
).build_api()

To run it :

uvicorn examples.serving.keras:app --host 0.0.0.0