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='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:
-
- 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.5-py39-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 64044d86e1ec2dc44e060eae3b9a96e66f7624c6abb5988617e5f01c7480bb8a |
|
MD5 | 159ebec778729d6e73b3ac2a9001ce6b |
|
BLAKE2b-256 | 9e89e81192c6bf2a1acb99f5c9c424373c3b41cdba52b1bacf5934ea0d09fba2 |
Hashes for deeplite_torch_zoo-2.0.5-py38-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | a0c21da57225b40216e35d07785bdcbb15deccd1aa416bc5ee32dff2e4377fd8 |
|
MD5 | fa0fb7390470372c3cd1009ec498f25c |
|
BLAKE2b-256 | 3f24d82d1f9949ace159876ec3fbf8f2cd60a01a8a1735abc71c2a71669fcfd4 |
Hashes for deeplite_torch_zoo-2.0.5-py37-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 9d890267914bf8df69f447df7db7f9591e7631009c012b1df8bd12906dc1ab3a |
|
MD5 | 06e4d600c4a02dad670bf2a1640dc645 |
|
BLAKE2b-256 | 599361215465fc9796864e24e85a6c5621d2c3a266e75fb8d763046b1254bc98 |