holocron.models

The models subpackage contains definitions of models for addressing different tasks, including: image classification, pixelwise semantic segmentation, object detection, instance segmentation, person keypoint detection and video classification.

The following datasets are available:

Classification

Classification models expect a 4D image tensor as an input (N x C x H x W) and returns a 2D output (N x K). The output represents the classification scores for each output classes.

import holocron.models as models
darknet19 = models.darknet19(num_classes=10)

Res2Net

class holocron.models.Res2Net(block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=26, scale=4, replace_stride_with_dilation=None, norm_layer=None)[source]

Implements a Res2Net model as described in https://arxiv.org/pdf/1904.01169.pdf

Parameters:
  • block (torch.nn.Module) – class constructor to be used for residual blocks

  • layers (list<python:int>) – layout of layers

  • num_classes (int) – number of output classes

  • zero_init_residual (bool) – whether the residual connections should be initialized at zero

  • groups (int) – number of convolution groups

  • width_per_group (int) – number of channels per group

  • scale (int) – scaling ratio within blocks

  • replace_stride_with_dilation (list<bool>) – whether stride should be traded for dilation

  • norm_layer (torch.nn.Module) – norm layer to be used

holocron.models.res2net(depth, num_classes, width_per_group=26, scale=4, pretrained=False, progress=True, **kwargs)[source]

Instantiate a Res2Net model

Parameters:
  • depth (int) – depth of the model

  • num_classes (int) – number of output classes

  • scale (int) – number of branches for cascade convolutions

  • pretrained (bool) – whether the model should load pretrained weights (ImageNet training)

  • progress (bool) – whether a progress bar should be displayed while downloading pretrained weights

  • **kwargs – optional arguments of torchvision.models.resnet.ResNet

Returns:

loaded Pytorch model

Return type:

model (torch.nn.Module)

Res2NeXt

holocron.models.res2next(depth, num_classes, width_per_group=4, scale=4, pretrained=False, progress=True, **kwargs)[source]

Instantiate a Res2NeXt model

Parameters:
  • depth (int) – depth of the model

  • num_classes (int) – number of output classes

  • scale (int) – number of branches for cascade convolutions

  • pretrained (bool) – whether the model should load pretrained weights (ImageNet training)

  • progress (bool) – whether a progress bar should be displayed while downloading pretrained weights

  • **kwargs – optional arguments of torchvision.models.resnet.ResNet

Returns:

loaded Pytorch model

Return type:

model (torch.nn.Module)

Darknet

holocron.models.darknet24(pretrained=False, progress=True, **kwargs)[source]

Darknet-24 from “You Only Look Once: Unified, Real-Time Object Detection”

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

Returns:

classification model

Return type:

torch.nn.Module

holocron.models.darknet19(pretrained=False, progress=True, **kwargs)[source]

Darknet-19 from “YOLO9000: Better, Faster, Stronger”

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

Returns:

classification model

Return type:

torch.nn.Module

holocron.models.darknet53(pretrained=False, progress=True, **kwargs)[source]

Darknet-53 from “YOLOv3: An Incremental Improvement”

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

Returns:

classification model

Return type:

torch.nn.Module

Object Detection

Object detection models expect a 4D image tensor as an input (N x C x H x W) and returns a list of dictionaries. Each dictionary has 3 keys: box coordinates, classification probability, classification label.

import holocron.models as models
yolov2 = models.yolov2(num_classes=10)

YOLO

holocron.models.yolov1(pretrained=False, progress=True, **kwargs)[source]

YOLO model from “You Only Look Once: Unified, Real-Time Object Detection”

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

Returns:

detection module

Return type:

torch.nn.Module

holocron.models.yolov2(pretrained=False, progress=True, **kwargs)[source]

YOLOv2 model from “YOLO9000: Better, Faster, Stronger”

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

Returns:

detection module

Return type:

torch.nn.Module