EdgeAI Torch Model Optimization Toolkit

March 25, 2026 ยท View on GitHub

Table of contents

Introduction

Tools to help development of Embedded Friendly Deep Neural Network Models in Pytorch. We call these model optimization tools because they help in making the models more efficient. This contains xmodelopt which is our extension of torch.ao model architecture optimization tools.

Note

The Quantization and Model Surgery tools are available in all of (v1, v2, v3) and Pruning is available in both (v2 and v3), but their interfaces and functionality are slightly different. We recommend to use the latest version (torch.fx based v3) tools whenever possible, but there are situations where the legacy tools may be useful - for example, for QAT models to be used with older TIDL versions or for older target devices such TDA4VM.

Why use our toolkit?

  • Our toolkit provides the APIs for quantization, surgery, pruning, and sparsity along with other tools, for users to seamlessly introduce them in their own training code.
  • These are wrapped in easy to use interfaces so that even someone with basic knowledge of Pytorch can use these. The user can add a single line of code to introduce each of these functionality as shown in the user guides.

Getting Started

Installation

Package installation

Install the package for usage

pip3 install "git+https://github.com/TexasInstruments/edgeai-modeloptimization.git#subdirectory=torchmodelopt"

Source Installation

Install this repository as a local editable Python package (for development)

cd edgeai-modeloptimization/torchmodelopt
./setup.sh

Features

Quantization:

Quantization documentation

Sparsity:

Model Pruning documentation

Experimental

Model Surgery:

Model Surgery documentation

Pruning

Model Pruning documentation

Other Utilities:

  • This package also contains edgeai_torchmodelopt.xnn which is our extension of torch.nn neural network model utilities. This is for advanced users and contains several undocumented features.

FAQ

Question 1: I am getting error "RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss." while training.

Solution 1: Try setting 'find_unused_parameters=True' while wrapping the model in torch.nn.parallel.DistributedDataParallel.

Question 2: Can I use different parts of the toolkit together?

Solution 2: Surgery currently works with both sparsity and quantization individually, but sparsity together with quantization is cyrrently not supported. It will be made available soon.

Question 3: I am seeing a huge drop in accuracy while training when using multi-gpu configuration.

Solution 3: Some configurations are seeing that, however, after the training, the model will be able to give the desired accuracy, otherwise single-gpu configuration could be used.

Question 4: I am getting RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment.

Solution 4: copy.deepcopy(model) is not supported while training and the training code might have it. It needs to be diabled and same model needs to be used. Make sure to keep the the model on same device though to enable training.

Contributors

  • @parakh08 [Parakh Agarwal]
  • @mathmanu [Manu Mathew]
  • @rakib23r [Rekib Zaman]
  • @kpatel [Kunal Patel]