Holocron: a Deep Learning toolbox for PyTorch¶
Holocron is meant to bridge the gap between PyTorch and latest research papers. It brings training components that are not available yet in PyTorch with a similar interface.
This project is meant for:
speed: architectures in this repo are picked for both pure performances and minimal latency
research: train your models easily to SOTA standards
Installation¶
Create and activate a virtual environment and then install Holocron:
Check out the installation guide for more options
Quick start¶
Load a model and use it just like any PyTorch model:
from PIL import Image
from torchvision.transforms.v2 import Compose, ConvertImageDtype, Normalize, PILToTensor, Resize
from torchvision.transforms.v2.functional import InterpolationMode
from holocron.models.classification import repvgg_a0
# Load your model
model = repvgg_a0(pretrained=True).eval()
# Read your image
img = Image.open(path_to_an_image).convert("RGB")
# Preprocessing
config = model.default_cfg
transform = Compose([
Resize(config['input_shape'][1:], interpolation=InterpolationMode.BILINEAR),
PILToTensor(),
ConvertImageDtype(torch.float32),
Normalize(config['mean'], config['std'])
])
input_tensor = transform(img).unsqueeze(0)
# Inference
with torch.inference_mode():
output = model(input_tensor)
print(config['classes'][output.squeeze(0).argmax().item()], output.squeeze(0).softmax(dim=0).max())
Model zoo¶
Image classification¶
- TridentNet from "Scale-Aware Trident Networks for Object Detection"
- SKNet from "Selective Kernel Networks"
- PyConvResNet from "Pyramidal Convolution: Rethinking Convolutional Neural Networks for Visual Recognition"
- ReXNet from "ReXNet: Diminishing Representational Bottleneck on Convolutional Neural Network"
- RepVGG from "RepVGG: Making VGG-style ConvNets Great Again"
Semantic segmentation¶
- U-Net from "U-Net: Convolutional Networks for Biomedical Image Segmentation"
- U-Net++ from "UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation"
- UNet3+ from "UNet 3+: A Full-Scale Connected UNet For Medical Image Segmentation"