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.1-py39-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 562861afe61560c1246e5cc84e31edd629ea5365f77b62ab39240d5a50ee8daa |
|
MD5 | 8689c9b521077922dd044ebf5f05bc54 |
|
BLAKE2b-256 | d233b164ada2b95722c15076b25dc7555a70455b72fb2bbc43b44c86fe2e5a3c |
Hashes for deeplite_torch_zoo-2.0.1-py38-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | e8fc28b3d7a6c1b635415407f0c35f8be3ee06071191206ccf62d1eea7219be0 |
|
MD5 | b527df97f8eda9b61be699c9cdf633f2 |
|
BLAKE2b-256 | 173c399e2510d34fa909a2480b3ce9939d787ca3afffa9b8734ebc804bcb7351 |
Hashes for deeplite_torch_zoo-2.0.1-py37-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 38dfc0d0cfbbf102c1bf0c0acbc40947ecf6b92975d21964bb38b3d575e63d57 |
|
MD5 | 90268833da26a062b4069b8d6be9ab31 |
|
BLAKE2b-256 | eec4bede6602767ad65a0e1ccf53e4cf9172c79210b6d664772536235b79e30f |