Home

Awesome

EdgeAI-TorchVision

Notice 1

If you have not visited the following landing pages, please do so before attempting to use this repository.

Notice 2

<hr>

Develop Embedded Friendly Deep Neural Network Models in PyTorch PyTorch

Introduction

This is an extension of the popular GitHub repository pytorch/vision that implements torchvision - PyTorch based datasets, model architectures, and common image transformations for computer vision.

Apart from the features in underlying torchvision, we support the following features

It is important to note that we do not modify the torchvision python package itself - so off-the-shelf, pip installed torchvision python package can be used with the scripts in this repository. See the setup documentation and the setup file for details. However, we do modify the training scripts in references that uses the torchvision package. When we need to modify a model, we do that by modifying the model object/instance in the training script using our model surgery tool.

<hr>

Setup

Setup Instructions

<hr>

Models and Scripts

We have different categories of models and scripts in this repository:

Category 1: "lite" variants of original torchvision models

These are models that created replacing unsupported layers of torchvision models with supported ones - we call them "lite" models. This replacing is done by using our Model Optimization Tools, just after the model is created, in the training script.

It is important to note that we do not modify the torchvision python package itself - so off-the-shelf, pip installed torchvision python package can be used with the scripts in this repository. However, we do modify the training scripts in references that uses the torchvision package. When we need to modify a model, we do that by modifying the model object/instance in the training script using our Model Optimization Tools.

To see example usages of Model Optimization Tools, please refer to references/classification/train.py and references/segmentation/train.py

The shell scripts run_torchvision_....sh can be used to train, evaluate or export these "lite" models. Accuracy results of training these lite models are in the documentation of our Model Optimization Tools.

Category 2: Our custom 'edgeailite' models

torchvision had only classification models before 2019. So we went ahead added embedded friendly models and training scripts for tasks such as Semantic Segmentation, Depth Estimation, Multi-Task Estimation etc. These are low complexity models that are suitable for embedded SoCs.

See documentation of our edgeailite extensions to torchvision

Scripts are provided for training these models for a variety of tasks:

These models, transforms, and training scrpts are in ./references/edgeailite. The models in our model zoo that were trained using these scripts carries a keyword "edgeailite".

The shell scripts run_edgeailite_....sh can be used to train, evaluate or export these "edgeailite" models.

Original torchvision models and documentation

This repository is built on top of 0.15.x release of torchvision. We do not modify torchvision python package itself, so the user can use the original models as well. See the original torchvision documentation:

<hr>