Introducing Digest AI: A Model Ingestion and Analysis Tool

Jan 03, 2025

Author: Victoria Godsoe, Manager AI Development

AMD, in collaboration with the ONNX Community and the Linux Foundation, has released an open-source tool to aid developers in their exploration and understanding of machine learning models. Digest AI is a powerful model analysis tool designed to help you extract valuable insights from your machine learning models. With Digest AI, you can easily organize, report on, and optimize your models, making your workflow more efficient and effective. In this blog, we will walk you through the key features and functionalities of Digest AI.

Getting Started with Digest AI

With just one click, Digest AI enables you to extract insights from your models and works on both Windows and Linux platforms. To get started, you can build Digest AI from the source code available on the GitHub repository at github.com/onnx/digest. Once you have completed the installation, run digest from your command line to launch the application. In the GitHub repository, there are also instructions on how to build an executable, which can then launch the application.

Opening and Analyzing a Model

Once you have Digest AI up and running, the next step is to open a model. There are three ways you can open a local model:

1. In the left sidebar, you can open a model by clicking on the “Open a local model” button:

Figure 1 - DigestAI
Figure 1: Open Local Model

2. You can use the keyboard shortcut CTRL+O, which will open a file explorer where you can select the model you would like to open.

3. Or you can simply drag and drop a model file directly onto the interface.

When you select a model file, Digest AI performs several checks to ensure the model's integrity, prepares it for analysis, and extracts the necessary data. After the analysis is complete, you will see a detailed summary of your model, including:

  • Basic Details: Parameters and FLOPs (Floating Point Operations), which is the algorithmic number of floating-point operations required to run a single set of inputs through the model.
  • Advanced Insights: Operation histogram, parameter and FLOPs intensity per operation type, and a model similarity graph comparing your model to other popular models.
  • Model Inputs and Outputs: Detailed information about tensor names, shapes, data types, and sizes.

There is also a Node List button that opens a window displaying all nodes in the model in a table format, which you can then sort on the various columns.

Figure 2: Node List
Figure 2: Node List

Example of the Node List View:

Figure 3 - DigestAI
Figure 3: Example of the Node List View

Saving and Reporting

Digest AI allows you to save the model analysis using the “Save Report” button found in the left menu. 

Figure 4 - DigestAI
Figure 4: Save Report

When you click the save button, you can then specify a folder to save the reports to. The following is an example of the files saved:

  • Heatmap: a PNG image that shows the similarity of the model to other public models.
  • Node CSV Report: contains all the model’s nodes, parameters, FLOPs, attributes, and inputs and outputs.
  • Report: a text document summarizing the model analysis.
  • Histogram: a PNG image of the operation count.
  • Node Type Counts CSV Report: contains all the model’s nodes and their counts.

Multi-Model Analysis

Digest AI also includes support for multi-model analysis. By clicking the “Multi-Model Analysis” button, you can select a folder containing multiple models. 

Figure 5 - DigestAI
Figure 5: Multi-Model Analysis

This will load a new view where the left pane will list unique models, while the right pane will indicate duplicates.

Figure 6 - DigestAI
Figure 6: Multi-Model Selection Window

You can right-click any model to open its summary view. 

Figure 7 - DigestAI
Figure 7: Model Summary View

Or click “Open Analysis” to analyze all models and view a global summary.

Figure 8 - DigestAI
Figure 8: Multi-Model Analysis Window

The Save Reports button allows you to save all model summaries and the global model summary as CSV files.

Beta Hugging Face Module

Digest AI also includes a Hugging Face module in beta status that enables users to download ONNX models directly from the Hugging Face hub without leaving Digest AI. You can search for a model, select it, and click Open. Digest AI will download the model, analyze it, and take you to the summary page.

Handling Dynamic Input Shapes

Some models may have dynamic input shapes, which can affect certain calculations like FLOPs. If you encounter this, Digest AI will display a warning message. To freeze a model with dynamic shapes, scroll down to the “Input Tensor” information section and click the blue snowflake icon next to the table.

Figure 9 - DigestAI
Figure 9: Freeze Dynamic Shapes

This will open a utility where you can specify static dimensions for the inputs, and you can then save the modified model with static shapes.

Figure 10 - DigestAI
Figure 10: Freeze Input Shapes Window

Digest API

Digest AI includes an API for ingesting, modifying, and analyzing machine learning models. This includes the ability to perform the following:

  • Analyzes a single model or multiple models in a directory.
  • Generates reports on model summary, node lists, node type counts, and node shape counts.

Code Block:

    # Single model
python analysis.py /path/to/model.onnx /path/to/output/directory

# Multiple models
python analysis.py /path/to/model/directory /path/to/output/directory

Conclusion

Digest AI is a powerful and user-friendly tool that can significantly enhance your machine learning development. By providing detailed model analysis, multi-model support, and integration with the Hugging Face hub, Digest AI helps you optimize your models and streamline your workflow. Ready to experience Digest AI? Visit github.com/onnx/digest and try it out today! For any feedback or questions, feel free to reach out to us at digestai@amd.com. We’d love to hear from you!

Share:

Article By

Profile Image of Victoria Godsoe

Victoria Godsoe

AI Developer Enablement Manager