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")

[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.