nni

An open source AutoML toolkit for automate machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyper-parameter tuning.

13,973

1,806

13,973

418

View on GitHub

Top Related Projects

optuna

10,484

A hyperparameter optimization framework

ray

32,953

Ray is a unified framework for scaling AI and Python applications. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.

hyperopt

7,191

Distributed Asynchronous Hyperparameter Optimization in Python

mlflow

18,287

Open source platform for the machine learning lifecycle

scikit-optimize

2,739

Sequential model-based optimization with a `scipy.optimize` interface

Quick Overview

NNI (Neural Network Intelligence) is an open-source AutoML toolkit developed by Microsoft for automating machine learning lifecycle, including feature engineering, neural architecture search, model compression and hyperparameter tuning. It aims to help developers efficiently design and tune machine learning models, supporting various environments including local machines, remote servers, and cloud services.

Pros

Comprehensive AutoML toolkit covering multiple aspects of ML lifecycle
Supports a wide range of ML frameworks (PyTorch, TensorFlow, Keras, etc.)
Extensible architecture allowing easy integration of custom algorithms
Provides both command-line and web UI for experiment management

Cons

Steep learning curve for beginners due to its extensive features
Documentation can be overwhelming and sometimes lacks detailed examples
Some advanced features may require significant computational resources
Limited support for certain specialized ML tasks or domains

Code Examples

Defining a search space for hyperparameter tuning:

from nni.experiment import Experiment

search_space = {
    'learning_rate': {'_type': 'loguniform', '_value': [0.0001, 0.1]},
    'batch_size': {'_type': 'choice', '_value': [16, 32, 64, 128]},
    'num_layers': {'_type': 'choice', '_value': [1, 2, 3, 4]}
}

experiment = Experiment('local')
experiment.config.search_space = search_space

Creating a simple trial function for model training:

import nni

def trial(params):
    # Your model training code here
    model = create_model(params['num_layers'])
    loss = train_model(model, params['learning_rate'], params['batch_size'])
    nni.report_final_result(loss)

experiment.config.trial_command = 'python trial.py'
experiment.config.trial_code_directory = '.'

Configuring and running an experiment:

experiment.config.max_trial_number = 100
experiment.config.trial_concurrency = 2
experiment.config.tuner.name = 'TPE'
experiment.config.tuner.class_args['optimize_mode'] = 'minimize'

experiment.run(8080)

Getting Started

To get started with NNI:

Install NNI:

pip install nni

Create a search space and trial function as shown in the code examples above.
Configure and run an experiment:

from nni.experiment import Experiment

experiment = Experiment('local')
experiment.config.search_space = search_space
experiment.config.trial_command = 'python trial.py'
experiment.config.trial_code_directory = '.'
experiment.config.max_trial_number = 100
experiment.config.tuner.name = 'TPE'
experiment.run(8080)

Access the NNI web UI at http://localhost:8080 to monitor and manage your experiment.

Competitor Comparisons

optuna

10,484

A hyperparameter optimization framework

Pros of Optuna

More lightweight and easier to integrate into existing projects
Supports a wider range of optimization algorithms, including Bayesian optimization
Better visualization tools for hyperparameter importance and optimization history

Cons of Optuna

Less comprehensive feature set compared to NNI's end-to-end ML lifecycle management
Lacks built-in support for distributed training and model compression

Code Comparison

Optuna:

import optuna

def objective(trial):
    x = trial.suggest_float('x', -10, 10)
    return (x - 2) ** 2

study = optuna.create_study()
study.optimize(objective, n_trials=100)

NNI:

import nni

@nni.trial
def trial(params):
    x = params['x']
    return {'default': (x - 2) ** 2}

if __name__ == '__main__':
    nni.run(search_space={'x': {'_type': 'uniform', '_value': [-10, 10]}})

Both frameworks offer simple APIs for defining and running hyperparameter optimization experiments. Optuna's approach is more pythonic and flexible, while NNI provides a more structured framework with additional features for managing the entire machine learning lifecycle.

ray

32,953

Ray is a unified framework for scaling AI and Python applications. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.

Pros of Ray

More comprehensive distributed computing framework, supporting a wider range of applications beyond just hyperparameter tuning
Better scalability for large-scale distributed computing tasks
More active community and frequent updates

Cons of Ray

Steeper learning curve due to its broader scope and more complex API
Potentially overkill for simpler machine learning tasks or smaller-scale projects
Less focused on automated machine learning and neural architecture search compared to NNI

Code Comparison

Ray example:

import ray

@ray.remote
def f(x):
    return x * x

futures = [f.remote(i) for i in range(4)]
print(ray.get(futures))

NNI example:

import nni

@nni.parameter_scope('example')
def f(x):
    return x * x

nni.report_final_result(f(nni.get_next_parameter()))

Both frameworks offer distributed computing capabilities, but Ray provides a more general-purpose solution, while NNI focuses on hyperparameter optimization and neural architecture search. Ray's code example demonstrates its distributed task execution, while NNI's example shows its integration with hyperparameter tuning experiments.

hyperopt

7,191

Distributed Asynchronous Hyperparameter Optimization in Python

Pros of hyperopt

Simpler and more lightweight, making it easier to integrate into existing projects
More flexible in defining search spaces, allowing for custom probability distributions
Better suited for small to medium-scale optimization tasks

Cons of hyperopt

Less comprehensive documentation and fewer examples compared to NNI
Limited built-in visualization tools for analyzing optimization results
Lacks some advanced features like early stopping and multi-objective optimization

Code Comparison

hyperopt:

from hyperopt import fmin, tpe, hp

def objective(params):
    x, y = params['x'], params['y']
    return x**2 + y**2

space = {'x': hp.uniform('x', -10, 10),
         'y': hp.uniform('y', -10, 10)}

best = fmin(objective, space, algo=tpe.suggest, max_evals=100)

NNI:

import nni

@nni.trial
def objective(params):
    x, y = params['x'], params['y']
    return x**2 + y**2

if __name__ == '__main__':
    nni.run(objective)

Both libraries offer powerful hyperparameter optimization capabilities, but hyperopt is more focused on this specific task, while NNI provides a broader range of features for neural architecture search and model compression. The choice between them depends on the specific requirements of your project and the scale of optimization needed.

Ax

2,340

Adaptive Experimentation Platform

Pros of Ax

More focused on Bayesian optimization and multi-armed bandits
Tighter integration with PyTorch ecosystem
Better support for multi-objective optimization

Cons of Ax

Less comprehensive feature set compared to NNI
Steeper learning curve for beginners
Limited support for distributed tuning

Code Comparison

NNI example:

import nni

@nni.trace
def objective(x, y):
    return x**2 + y**2

if __name__ == '__main__':
    nni.run(objective)

Ax example:

from ax import optimize

def objective(x, y):
    return x**2 + y**2

best_parameters, values, experiment, model = optimize(
    parameters=[
        {"name": "x", "type": "range", "bounds": [-10, 10]},
        {"name": "y", "type": "range", "bounds": [-10, 10]},
    ],
    evaluation_function=objective,
    objective_name="quadratic",
)

Both repositories offer powerful optimization tools, but they cater to slightly different use cases. NNI provides a more comprehensive suite of features for neural architecture search and hyperparameter optimization, while Ax focuses on Bayesian optimization and multi-armed bandits. The choice between the two depends on the specific requirements of your project and your familiarity with the respective ecosystems.

mlflow

18,287

Open source platform for the machine learning lifecycle

Pros of MLflow

More comprehensive tracking and management of the entire ML lifecycle
Better integration with popular ML frameworks and cloud platforms
Larger community and ecosystem, with more plugins and integrations available

Cons of MLflow

Steeper learning curve for beginners due to its broader feature set
Can be overkill for simple projects or small-scale experiments
Less focus on automated hyperparameter tuning compared to NNI

Code Comparison

MLflow experiment tracking:

import mlflow

with mlflow.start_run():
    mlflow.log_param("param1", value1)
    mlflow.log_metric("metric1", score)
    mlflow.sklearn.log_model(model, "model")

NNI experiment tracking:

import nni

@nni.report_final_result
def run_experiment(params):
    model = train_model(params)
    accuracy = evaluate_model(model)
    nni.report_intermediate_result(accuracy)
    return accuracy

if __name__ == '__main__':
    nni.run(run_experiment)

Both MLflow and NNI offer powerful tools for machine learning experimentation and management. MLflow provides a more comprehensive solution for the entire ML lifecycle, while NNI focuses more on automated hyperparameter tuning and neural architecture search. The choice between the two depends on specific project requirements and team preferences.

scikit-optimize

2,739

Sequential model-based optimization with a `scipy.optimize` interface

Pros of scikit-optimize

Simpler and more lightweight, focusing specifically on Bayesian optimization
Easier to integrate into existing Python projects, especially those using scikit-learn
More extensive documentation and examples for various optimization scenarios

Cons of scikit-optimize

Limited to Bayesian optimization techniques, lacking the broader range of algorithms offered by NNI
Does not provide a comprehensive AutoML platform or neural architecture search capabilities
Less support for distributed and parallel optimization tasks

Code Comparison

scikit-optimize:

from skopt import gp_minimize

def objective(x):
    return (x[0] - 3)**2 + (x[1] + 2)**2

result = gp_minimize(objective, [(-5.0, 5.0), (-5.0, 5.0)])

NNI:

import nni

@nni.trial
def objective(params):
    x, y = params['x'], params['y']
    return (x - 3)**2 + (y + 2)**2

if __name__ == '__main__':
    nni.run(objective)

Both libraries aim to simplify optimization tasks, but NNI offers a more comprehensive suite of tools for AutoML and hyperparameter tuning, while scikit-optimize focuses on providing a streamlined experience for Bayesian optimization within the scikit-learn ecosystem.

Convert designs to code with AI

Introducing Visual Copilot: A new AI model to turn Figma designs to high quality code using your components.

Try Visual Copilot

README

NNI automates feature engineering, neural architecture search, hyperparameter tuning, and model compression for deep learning. Find the latest features, API, examples and tutorials in our official documentation (ç®ä½ä¸æçç¹è¿é).

What's NEW!

New release: v3.0 preview is available - released on May-5-2022
New demo available: Youtube entry | Bilibili å¥å£ - last updated on June-22-2022
New research paper: SparTA: Deep-Learning Model Sparsity via Tensor-with-Sparsity-Attribute - published in OSDI 2022
New research paper: Privacy-preserving Online AutoML for Domain-Specific Face Detection - published in CVPR 2022
Newly upgraded documentation: Doc upgraded

Installation

See the NNI installation guide to install from pip, or build from source.

To install the current release:

$ pip install nni

To update NNI to the latest version, add --upgrade flag to the above commands.

NNI capabilities in a glance

	Hyperparameter Tuning	Neural Architecture Search	Model Compression
Algorithms	Exhaustive search Grid Search Random Heuristic search Anneal Evolution Hyperband PBT Bayesian optimization BOHB DNGO GP Metis SMAC TPE	Multi-trial Grid Search Policy Based RL Random Regularized Evolution TPE One-shot DARTS ENAS FBNet ProxylessNAS SPOS	Pruning Level L1 Norm Taylor FO Weight Movement AGP Auto Compress More... Quantization Naive QAT LSQ Observer DoReFa BNN
	Supported Frameworks	Training Services	Tutorials
Supports	PyTorch TensorFlow Scikit-learn XGBoost LightGBM MXNet Caffe2 More...	Local machine Remote SSH servers Azure Machine Learning (AML) Kubernetes Based OpenAPI Kubeflow FrameworkController AdaptDL PAI DLC Hybrid training services	HPO PyTorch TensorFlow NAS Hello NAS NAS Benchmarks Compression Pruning Pruning Speedup Quantization Quantization Speedup

Resources

Contribution guidelines

If you want to contribute to NNI, be sure to review the contribution guidelines, which includes instructions of submitting feedbacks, best coding practices, and code of conduct.

We use GitHub issues to track tracking requests and bugs. Please use NNI Discussion for general questions and new ideas. For questions of specific use cases, please go to Stack Overflow.

Participating discussions via the following IM groups is also welcomed.

Gitter		WeChat
	OR

Over the past few years, NNI has received thousands of feedbacks on GitHub issues, and pull requests from hundreds of contributors. We appreciate all contributions from community to make NNI thrive.

Test status

Essentials

Type	Status
Fast test
Full test - HPO
Full test - NAS
Full test - compression

Training services

Type	Status
Local - linux
Local - windows
Remote - linux to linux
Remote - windows to windows
OpenPAI
Frameworkcontroller
Kubeflow
Hybrid
AzureML

Related Projects

Targeting at openness and advancing state-of-art technology, Microsoft Research (MSR) had also released few other open source projects.

OpenPAI : an open source platform that provides complete AI model training and resource management capabilities, it is easy to extend and supports on-premise, cloud and hybrid environments in various scale.
FrameworkController : an open source general-purpose Kubernetes Pod Controller that orchestrate all kinds of applications on Kubernetes by a single controller.
MMdnn : A comprehensive, cross-framework solution to convert, visualize and diagnose deep neural network models. The "MM" in MMdnn stands for model management and "dnn" is an acronym for deep neural network.
SPTAG : Space Partition Tree And Graph (SPTAG) is an open source library for large scale vector approximate nearest neighbor search scenario.
nn-Meter : An accurate inference latency predictor for DNN models on diverse edge devices.

We encourage researchers and students leverage these projects to accelerate the AI development and research.

License

The entire codebase is under MIT license.

Top Related Projects

Convert designs to code with AI

Introducing Visual Copilot: A new AI model to turn Figma designs to high quality code using your components.

Try Visual Copilot