Deeplite Torch Zoo

The deeplite-torch-zoo package is a collection of popular (pretrained) CNN model architectures and benchmark datasets for PyTorch. The models are grouped under different datasets and different task types such as classification, object detection, and semantic segmentation. The primary aim of deeplite-torch-zoo is to booststrap applications by starting with the most suitable pretrained models for a given task. In addition, the pretrained models from deeplite-torch-zoo can be used as a good starting point for optimizing model architectures using our neutrino_engine

• Installation

  • Install using pip
  • Install from source
  • Install in Dev mode

• How to Use

  • Loading Datasets
    • Classification Datasets
    • Object Detection Datasets
  • Loading Models
    • Classification Models
    • Object Detection Models

• Available Models

• Available Datasets

• Training on Custom Data

• Benchmark Results

• Contribute a Model/Dataset to the Zoo

Installation

Install using pip (release version)

Use following command to install the package from our internal PyPI repository.

    $ pip install --upgrade pip
    $ pip install deeplite-torch-zoo

Install from source (development version)

    $ git clone https://github.com/Deeplite/deeplite-torch-zoo.git
    $ pip install .

Install in dev mode

    $ git clone https://github.com/Deeplite/deeplite-torch-zoo.git
    $ pip install -e .
    $ pip install -r requirements-test.txt

To test the installation, one can run the basic tests using pytest command in the root folder.

How to Use

The deeplite-torch-zoo is collection of benchmark computer vision datasets and pretrained models. There are four primary wrapper functions to load datasets, models and evaluation functions: get_data_splits_by_name, get_model_by_name, get_eval_function and create_model which can be imported as

from deeplite_torch_zoo import get_data_splits_by_name
from deeplite_torch_zoo import get_model_by_name
from deeplite_torch_zoo import get_eval_function
from deeplite_torch_zoo import create_model

Loading Datasets

The loaded datasets are available as a dictionary of the following format: {'train': train_dataloder, 'test': test_dataloader}. The train_dataloder and test_dataloader are objects of type torch.utils.data.DataLoader.

Classification Datasets

    data_splits = get_data_splits_by_name(
        data_root="./", dataset_name="cifar100", model_name="resnet18", batch_size=128
    )

The list of all available classification datasets can be found here. Please note that it is always necessary to pass the model name upon the creation of dataloader because the dataset class logic might depend on the model type.

Object Detection Datasets

The following sample code loads the PASCAL VOC dataset. train contains the data loader for the trainval data splits of the VOC2007 and/or VOC2012. If both datasets are provided it concatenates both VOC2007 and VOC2012 train sets. Otherwise, it returns the train set for the provided dataset. 'test' contains dataloader (always with batch_size=1) for the test split of VOC2007. You also need to provide the model name to instantiate the dataloaders.

data_splits = get_data_splits_by_name(
        data_root=PATH_TO_VOCdevkit,
        dataset_name="voc",
        model_name="yolo3",
        batch_size=BATCH_SIZE,
    )

The list of all available object detection datasets can be found here.

NOTE: As it can be observed the data_loaders are provided based on the corresponding model (model_name). Different object detection models consider inputs/outputs in different formats, and thus the our data_splits are formatted according to the needs of the model (e.g. for SSD or YOLO detection models).

Loading and Creating Models

Models are generally provided with weights pretrained on specific datasets. One would load a model X pretrained on a dataset Y to get the appropriate weights for the task Y. The get_model_by_name could used for this purpose. There is also an option to create a new model with an arbitrary number of categories for the downstream tasl and load the weights from another dataset for transfer learning (e.g. to load COCO weights to train a model on the VOC dataset). The create_model method should be generally used for that. Note that get_model_by_name always returns a fully-trained model for the specified task, this method thus does not allow specifying a custom number of classes.

Classification Models

To get a pretrained classification model one could use

    model = get_model_by_name(
        model_name="resnet18",
        dataset_name="cifar100",
        pretrained=True, # or False, if pretrained weights are not required
        progress=False, # or True, if a progressbar is required
        device="cpu", # or "cuda"
    )

To create a new model with ImageNet weights and a custom number of classes one could use

    model = create_model(
        model_name="resnet18",
        pretraining_dataset="imagenet",
        num_classes=42,
        pretrained=True, # or False, if pretrained weights are not required
        progress=False, # or True, if a progressbar is required
        device="cpu", # or "cuda"
    )

This method would load the ImageNet-pretrained weights to all the modules of the model where one could match the shape of the weight tensors (i.e. all the layers except the last fully-connected one in the above case).

The list of all available classification models can be found here.

Object Detection Models

    model = get_model_by_name(
        model_name="yolo4s",
        dataset_name="voc",
        pretrained=True, # or False, if pretrained weights are not required
        progress=False, # or True, if a progressbar is required
    )

Likewise, to create a object detection model with an arbitrary number of classes

    model = get_model_by_name(
        model_name="yolo4s",
        num_classes=5,
        dataset_name="coco",
        pretrained=True, # or False, if pretrained weights are not required
        progress=False, # or True, if a progressbar is required
    )

The list of all available Object Detection models can be found here.

Creating an evaluation function

To create an evaluation fuction for the given model and dataset one could call get_eval_function passing the model_name and dataset_name arguments:

    eval_fn = get_eval_function(
        model_name="resnet50",
        dataset_name="imagenet",
    )

The returned evaluation function is a Python callable that takes two arguments: a PyTorch model object and a PyTorch dataloader object (logically corresponding to the test split dataloader) and returns a dictionary with metric names as keys and their corresponding values.

Available Models

There is an useful utility function list_models which can be imported as

from deeplite_torch_zoo import list_models

This utility will help in listing available pretrained models or datasets.

For instance list_models("yolo5") will provide the list of available pretrained models that contain yolo5 in their model names. Similar results e.g. can be obtained using list_models("yo"). Filtering models by the corresponding task type is also possible by passing the string of the task type with the task_type_filter argument (the following task types are available: classification, object_detection, semantic_segmentation).

    +------------------+------------------------------------+
    | Available models |          Source datasets           |
    +==================+====================================+
    | yolo5_6l         | voc                                |
    +------------------+------------------------------------+
    | yolo5_6m         | coco, voc                          |
    +------------------+------------------------------------+
    | yolo5_6m_relu    | person_detection, voc              |
    +------------------+------------------------------------+
    | yolo5_6ma        | coco                               |
    +------------------+------------------------------------+
    | yolo5_6n         | coco, person_detection, voc, voc07 |
    +------------------+------------------------------------+
    | yolo5_6n_hswish  | coco                               |
    +------------------+------------------------------------+
    | yolo5_6n_relu    | coco, person_detection, voc        |
    +------------------+------------------------------------+
    | yolo5_6s         | coco, person_detection, voc, voc07 |
    +------------------+------------------------------------+
    | yolo5_6s_relu    | person_detection, voc              |
    +------------------+------------------------------------+
    | yolo5_6sa        | coco, person_detection             |
    +------------------+------------------------------------+
    | yolo5_6x         | voc                                |
    +------------------+------------------------------------+

Train on Custom Dataset

One could refer to the example training scripts to see how the zoo could be integrated into differen training pipelines. For more details please see

• Training a Classification Model
• Training an Object Detection Model

Benchmark Results

Please refer to our documentation for the detailed performance metrics of the pretrained models available in the deeplite-torch-zoo. After downloading a model, please evaluate the model using deeplite-profiler to verify the performance metric values. However, one may see different numbers for the execution time as the target hardware and/or the load on the system may impact it.

Contribute a Model/Dataset to the Zoo

NOTE: If you looking for an SDK documentation, please head over here.

We always welcome community contributions to expand the scope of deeplite-torch-zoo and also to have additional new models and datasets. Please refer to the documentation for the detailed steps on how to add a model and dataset. In general, we follow the fork-and-pull Git workflow.

1. Fork the repo on GitHub
2. Clone the project to your own machine
3. Commit changes to your own branch
4. Push your work back up to your fork
5. Submit a Pull request so that we can review your changes

NOTE: Be sure to merge the latest from "upstream" before making a pull request!

Credit

Object Detection

• The implementation of mb-ssd models: qfgaohao/pytorch-ssd
• The implementation of resnet-ssd: Nvidia-SSD
• The implementation of yolov5: ultralytics/yolov5
• The implementation of flexible-yolov5: Bobo-y/flexible-yolov5
• The implementation of yolov7: WongKinYiu/yolov7

Segmentation

• The implementation of deeplab: pytorch-deeplab-xception
• The implementation of unet_scse: nyoki-mtl/pytorch-segmentation
• The implementation of fcn: wkentaro/pytorch-fcn
• The implementation of Unet: milesial/Pytorch-UNet

Classification

• The implementation of models on CIFAR100 dataset: kuangliu/pytorch-cifar
• The implementation of Mobilenetv1 model on VWW dataset: qfgaohao/pytorch-ssd
• The implementation of Mobilenetv3 model on VWW dataset: d-li14/mobilenetv3.pytorch

Misc

• mAP metric calculation code: bes-dev/mean_average_precision