How to Register a ML model using MLflow

Rajamannar A K 86 Reputation points


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))  
   model = mlflow.pyfunc.load_model("/dbfs/FileStore/pyfunc/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")  

this is the error which I get while running the register model line


Azure Machine Learning
Azure Machine Learning
An Azure machine learning service for building and deploying models.
2,442 questions
Azure Databricks
Azure Databricks
An Apache Spark-based analytics platform optimized for Azure.
1,859 questions
0 comments No comments
{count} votes

1 answer

Sort by: Most helpful
  1. Ramr-msft 17,521 Reputation points

    @Rajamannar A K Thanks for the question. Can you please add more details about the trained model and deployment that you are trying.

    Here is link to an example of how to deploy and use object detection model in Azure ML with MLFlow option.

    the process to do MLFlow using a custom python flavor then deploy to Azure ML using the Azure version of MLFlow.
    Created the template here:

    1. Main deployment script to AML: aml_mlflow_deployment/mlflow_aml_deployment.ipynb at main · james-tn/aml_mlflow_deployment (
    2. Custom mlflow.pyfunc module: aml_mlflow_deployment/ at main · james-tn/aml_mlflow_deployment (

    There's also the option of simply using Azure DevOps to pull the model out of the MLFlow model repo in Azure Databricks, containerize it, and put it directly to ACI/AKS for API scoring outside of Azure Databricks.

    0 comments No comments