The deeplite-torch-zoo package is a collection of popular pretrained deep learning models and their datasets for PyTorch framework.
Project description
🚀 Deeplite Torch Zoo 🚀 is a collection of state-of-the-art efficient computer vision models for embedded applications in PyTorch.
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:
- timm models,
- pytorchcv models,
- other SOTA efficient models (EdgeViT, FasterNet, GhostNetV2, MobileOne)
- Single interface for SOTA YOLO detectors (compatible with Ultralytics training):
- YOLOv3, v4, v5, v6-3.0, v7, v8
- YOLO with timm backbones
- other experimental configs
📋 Table of content
⏳ 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='v8', # 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
🛠 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:
-
- support for Knowledge Distillation
- training recipes provides (A1, A2, A3, USI, etc.)
🤝 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.
- Fork the repo on GitHub
- Clone the project to your own machine
- Commit changes to your own branch
- Push your work back up to your fork
- 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
- YOLOv3 implementation: ultralytics/yolov3
- YOLOv5 implementation: ultralytics/yolov5
- flexible-yolov5 implementation: Bobo-y/flexible-yolov5
- YOLOv8 implementation: ultralytics/ultralytics
- YOLOv7 implementation: WongKinYiu/yolov7
- YOLOX implementation: iscyy/yoloair
- westerndigitalcorporation/YOLOv3-in-PyTorch
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
DNN building block implementations
- d-li14/mobilenetv2.pytorch
- d-li14/efficientnetv2.pytorch
- apple/ml-mobileone
- osmr/imgclsmob
- huggingface/pytorch-image-models
- moskomule/senet.pytorch
- DingXiaoH/RepLKNet-pytorch
- huawei-noah/Efficient-AI-Backbones
Misc
- torchvision dataset implementations: pytorch/vision
- MLP implementation: aaron-xichen/pytorch-playground
- AutoAugment implementation: DeepVoltaire/AutoAugment
- Cutout implementation: uoguelph-mlrg/Cutout
- Robustness measurement image distortions: hendrycks/robustness
- Registry implementation: openvinotoolkit/openvino/tools/pot
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distributions
Hashes for deeplite_torch_zoo-2.0.2-py39-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | c03b52ecca6ae7eaf1256542402cb974ca691acecb2eb06d06ae29f99930da86 |
|
MD5 | 09b3cfab4e8f24fb628806bc742c7f0c |
|
BLAKE2b-256 | c3c5ec70f65778bdac278c8f5ac8e95d900b8908e76b1e5198df95fbc34e9e0b |
Hashes for deeplite_torch_zoo-2.0.2-py38-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 4d6d9e3842125661b67aae74099a498b89c649879d20b775fcfbdecd9e458634 |
|
MD5 | 06b05ec3fcadc645e5819621f0d6704a |
|
BLAKE2b-256 | 6101c5275e9268d2a6e00c801fd258a6e81bba58a250f776a28b2f9d270c3ac8 |
Hashes for deeplite_torch_zoo-2.0.2-py37-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 4eedd0cace3744258de815987cb932a05dd4151e3e38ced1252a33db93d8b66a |
|
MD5 | ea1f4e7a397fea77d15d09dec6299299 |
|
BLAKE2b-256 | 7e39473d154ff97239ca6038678c618e1ad11d96bb6ecb2a935e6017f25dd5bb |