Using MLflow’s Model Registry

Next Item Prediction

We use the MLflow Model Registry to register three different artifacts that are used in the AI Workflow:

  • Transformers4Rec (PyTorch) model

  • NVTabular workflow

  • Triton Inference Server ensemble, which is a combination of the two above

Using MLflow allows us to refer to the models by name rather than storage location. This AI Workflow does not use advanced functionality such as model lifecycle. Running A/B tests or deploying multiple next-item prediction models are out of scope for this AI Workflow, but should you venture down that path, MLflow will help keep your models organized.


We also show an example of how to log model parameters and metrics to as MLflow Experiments.


For running MLflow locally during development, refer to the Developer Guide on how to run the MLflow server.

The Deployment Guide will walk through how to run a production version of MLflow on Kubernetes.

See MLflow’s website for full documentation about how to use MLflow.

© Copyright 2022-2023, NVIDIA. Last updated on May 23, 2023.