MLflow PyTorch Flavor

Introduction

PyTorch is an open-source machine learning library developed by Facebook’s AI Research lab. It provides a flexible and intuitive framework for deep learning and is particularly favored for its dynamic computation graph (eager mode), which provides a more pythonic development flow compared to static graph frameworks (graph mode). PyTorch is efficient for large-scale data processing and neural network training. Due to its ease of use and robust community support, PyTorch has become a popular choice among researchers and developers in the AI field.

MLflow has built-in support (we call it MLflow PyTorch flavor) for PyTorch workflow, at a high level in MLflow we provide a set of APIs for:

  • Simplified Experiment Tracking: Log parameters, metrics, and models during model training.

  • Experiments Management: Store your PyTorch experiments in MLflow server, and you can view and share them from MLflow UI.

  • Effortless Deployment: Deploy PyTorch models with simple API calls, catering to a variety of production environments.

Developer Guide of PyTorch with MLflow

To learn more about the nuances of the pytorch flavor in MLflow, please read the developer guide. It will walk you through the following topics:

View the Developer Guide
  • Logging PyTorch Experiments with MLflow: How to log PyTorch experiments to MLflow, including training metrics, model parameters, and training hyperparamers.

  • Log Your PyTorch Models with MLflow: How to log your PyTorch models with MLflow and how to load them back for inference.