Skip to main content
Version: 4.4

Train & Log a Custom Machine Learning Model

Log a Custom Machine Learning Model with Katonic-SDK Log package.

Import necessary packages​

import os

import mlflow.pyfunc
from katonic.log.logmodel import LogModel

Define Experiment name​

experiment_name = "custom_model"

Initiate LogModel with experiment name​

lm = LogModel(experiment_name, source_name='custom_model_logging.ipynb', features=[])

Check Metadata of the created / existing experiment​

# experiment id
exp_id = lm.id

print("experiment name : ", lm.name)
print("experiment location : ", lm.location)
print("experiment id : ", lm.id)
print("experiment status : ", lm.stage)

Artifact path where you want to log your model​

artifact_path = "custom-model-log"

Create custom model class​

class AddN(mlflow.pyfunc.PythonModel):

    def __init__(self, n):
        self.n = n

    def predict(self, context, model_input):
        return model_input.apply(lambda column: column + self.n)

        

add5_model = AddN(n=5)

Log Custom Model using Log pkg​

lm.model_logging(
    model_name="add_n",
    model_type="custom-model",
    model=add5_model,
    artifact_path=artifact_path,
    current_working_dir=f'{os.getcwd()}/custom_model_logging.ipynb'
)

Check all the logged Experiments​

You can search and get all the logged experiments with experiment ID.

df_runs = lm.search_runs(exp_id)
print("Number of runs done : ", len(df_runs))