How to Register a ML model using MLflow

[Itachi_Naruto]

Hi,

I have a PyTorch model which I have pushed into the dbfs now I want to serve the model using MLflow. I saw that the model needs to be in python_function model.

To do that I did the following methods

1. load the model from dbfs using torch load option

2. Then save the model in python_function model using the pyfunc.save_model function

3. After this when I register the model I get a decode error

I'm not training any model in the Databricks.

import mlflow
import mlflow.pyfunc
from torch import load as torch_load
 
py_model = torch_load( "/dbfs/FileStore/ml/ner_model" , map_location = torch_device(ner_gpu_device))
mlflow.pytorch.save_model(py_model,path="/dbfs/FileStore/pyfunc/ner_model")
 
model = mlflow.pyfunc.load_model("/dbfs/FileStore/pyfunc/ner_model")
mlflow.register_model(model,"ner_model")
 
# loading the python_function model to register
model = mlflow.pyfunc.load_model("/dbfs/FileStore/pyfunc/ner_model")
model_version = mlflow.register_model(model,"ner_model")

[Kaniz_Fatma]

Hi @Itachi_Naruto! My name is Kaniz, and I'm the technical moderator here. Great to meet you, and thanks for your question! Let's see if your peers in the community have an answer to your question first. Or else I will get back to you soon. Thanks.

[Anonymous]

I think you want to use mflow to load the model not pytorch. There is a function in mlflow to load pytorch models https://www.mlflow.org/docs/latest/python_api/mlflow.pytorch.html#mlflow.pytorch.load_model

Then once it's loaded, you can log it and register it as normal.