[Avg. reading time: 8 minutes]
Model Flavors
Remember MLflow features (Experiments - Runs - Models - Versions)
Rerun the model again.
git clone https://github.com/gchandra10/uni_multi_model
Popular MLflow Model Flavors
| Flavor | Used For | Typical Libraries / Frameworks |
|---|---|---|
sklearn | Traditional ML models (regression, classification, clustering) | Scikit-Learn, statsmodels |
xgboost | Gradient boosting trees | XGBoost |
lightgbm | High-performance gradient boosting | LightGBM |
catboost | Categorical-feature-friendly boosting | CatBoost |
pytorch | Deep learning and neural networks | PyTorch |
tensorflow / keras | Deep learning models | TensorFlow, Keras |
onnx | Portable models for cross-framework inference | ONNX Runtime |
fastai | Transfer learning and DL pipelines | FastAI |
statsmodels | Statistical / econometric models | statsmodels |
prophet | Time series forecasting | Facebook Prophet |
gluon | Deep learning (MXNet backend) | Apache MXNet |
sparkml | Distributed ML pipelines | Apache Spark MLlib |
pyfunc | Universal interface — wraps all other flavors | MLflow internal meta-flavor |
PyFunc makes ML models cross-platform — one consistent way to load and predict, regardless of how they were built.
-
Just like apps can be built separately for iOS or Android, models in MLflow can be saved in different native formats (like Scikit-Learn, PyTorch, XGBoost, etc.).
-
A cross-platform app works everywhere, and that’s what PyFunc is for ML models: a universal wrapper that runs any model with the same interface.
-
This lets teams serve and reuse models easily, without worrying about which library originally trained them.
For Example:
| Library | Save API | Predict Method |
|---|---|---|
| Scikit-Learn | joblib.dump() | model.predict() |
| TensorFlow | model.save() | model(x) |
| PyTorch | torch.save() | model.forward(x) |
| XGBoost | model.save_model() | model.predict(xgb.DMatrix(x)) |
You can use pyfunc for all the flavors
import mlflow.pyfunc
mlflow.pyfunc.save_model()
---
---
---
model = mlflow.pyfunc.load_model("models:/<name>/<stage>")
model.predict(pd.DataFrame(...))
Advantages
- One simple API for inference. Works the same whether the model was trained in Scikit-Learn, XGBoost, PyTorch, or TensorFlow.
- Reduces code differences between data-science teams using different libraries.
- PyFunc packages the model + environment (conda/requirements) together.
- Guarantees that the model runs identically on local machines, servers, or cloud.
- Ideal for CI/CD pipelines and container builds.
- Can be loaded from: Run path: runs:/<run_id>/model Registry stage: models:/name/Production
- Works seamlessly with MLflow Serving, FastAPI, Docker, or SageMaker deploys.
- Enables easy A/B comparisons between models trained in different frameworks.
- You can subclass mlflow.pyfunc.PythonModel to: Add preprocessing or feature engineering. Postprocess predictions. Integrate external systems (feature store, logging, metrics).
Limitations
- Framework-specific features are lost.
- Input is pandas centric.
- In some cases, can be slower than native runtimes. (Torch/Tensor flow)
https://github.com/gchandra10/uni_multi_model/blob/main/03_load_test_model.py