Skip to content

Collection of SOTA efficient computer vision models for embedded applications, with pre-trained weights and training recipes

License

Notifications You must be signed in to change notification settings

Deeplite/deeplite-torch-zoo

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

logo

πŸš€ Deeplite Torch Zoo πŸš€ is a collection of state-of-the-art efficient computer vision models for embedded applications in PyTorch.

Build Status codecov

For information on YOLOBench, click here.

The main features of this library are:

  • High-level API to create models, dataloaders, and evaluation functions
  • Single interface for SOTA classification models:
  • Single interface for SOTA YOLO detectors (compatible with Ultralytics training):

πŸ“‹ Table of content

  1. Quick start
  2. Installation
  3. Training scripts
  4. Contributing
  5. Credit

⏳ Quick start

Create a classification model

from deeplite_torch_zoo import get_model, list_models

model = get_model(
    model_name='edgevit_xs',        # model names for imagenet available via `list_models('imagenet')`
    dataset_name='imagenet',        # dataset name, since resnet18 is different for e.g. imagenet and cifar100
    pretrained=False,               # if True, will try to load a pre-trained checkpoint
)

# creating a model with 42 classes for transfer learning:

model = get_model(
    model_name='fasternet_t0',        # model names for imagenet available via `list_models('imagenet')`
    num_classes=42,                   # number of classes for transfer learning
    dataset_name='imagenet',   # take weights from checkpoint pre-trained on this dataset
    pretrained=False,                 # if True, will try to load all weights with matching tensor shapes
)

Create an object detection model

from deeplite_torch_zoo import get_model

model = get_model(
    model_name='yolo4n',        # creates a YOLOv4n model on COCO
    dataset_name='coco',        # (`n` corresponds to width factor 0.25, depth factor 0.33)
    pretrained=False,           # if True, will try to load a pre-trained checkpoint
)

# one could create a YOLO model with timm backbone,
# PAN neck and YOLOv8 decoupled anchor-free head like this:

model = get_model(
    model_name='yolo_timm_fbnetv3_d',    # creates a YOLO with FBNetV3-d backbone from timm
    dataset_name='coco',                 #
    pretrained=False,                    # if True, will try to load a pre-trained checkpoint
    custom_head='yolo8',                 # will replace default detection head
                                         # with YOLOv8 detection head
)

Create PyTorch dataloaders

from deeplite_torch_zoo import get_dataloaders

dataloaders = get_dataloaders(
    data_root='./',                      # folder with data, will be used for download
    dataset_name='imagewoof',            # datasets to if applicable,
    num_workers=8,                       # number of dataloader workers
    batch_size=64,                       # dataloader batch size (train and test)
)

# dataloaders['train'] -> train dataloader
# dataloaders['test'] -> test dataloader
#
# see below for the list of supported datasets

The list of supported datasets is available for classification and object detection.

Creating an evaluation function

from deeplite_torch_zoo import get_eval_function

eval_function = get_eval_function(
    model_name='yolo8s',
    dataset_name='voc',
)

# required arg signature is fixed for all eval functions
metrics = eval_function(model, test_dataloader)

(Experimental) Training with patched Ultralytics trainer

from deeplite_torch_zoo.trainer import Detector

model = Detector(model_name='yolo7n')        # will create a wrapper around YOLOv7n model
                                             # (YOLOv7n model with YOLOv8 detection head)
model.train(data='VOC.yaml', epochs=100)     # same arguments as Ultralytics trainer

# alternatively:

torch_model = get_model(
    model_name='yolo7n',
    dataset_name='coco',
    pretrained=False,
    custom_head='yolo8',
)
model = Detector(torch_model=torch_model)    # either `model_name` or `torch_model`
model.train(data='VOC.yaml', epochs=100)     # should be provided

πŸ›  Installation

PyPI version:

$ pip install deeplite-torch-zoo

Latest version from source:

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

πŸ’ͺ Training scripts

We provide several training scripts as an example of how deeplite-torch-zoo can be integrated into existing training pipelines:

🀝 Contributing

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

Repositories used to build Deeplite Torch Zoo

Object Detection

Segmentation

Classification

DNN building block implementations

Misc

About

Collection of SOTA efficient computer vision models for embedded applications, with pre-trained weights and training recipes

Resources

License

Code of conduct

Stars

Watchers

Forks

Packages

No packages published

Languages