Necks

Classification

class torchok.models.necks.classification.hrnet.HRNetClassificationNeck(in_channels: Union[List[int], Tuple[int, ...]])

Bases: BaseModel

HRNet neck for classification task.

__init__(in_channels: Union[List[int], Tuple[int, ...]])

Init HRNetClassificationNeck.

Parameters

in_channels – Input channels.

forward(x: List[Tensor]) → Tensor

Forward method.

training: bool

Segmentation

class torchok.models.necks.segmentation.hrnet.HRNetSegmentationNeck(in_channels: Union[List[int], Tuple[int, ...]])

Bases: BaseModel

HRNet neck for segmentation task.

__init__(in_channels: Union[List[int], Tuple[int, ...]])

Init HRNetSegmentationNeck.

Parameters

in_channels – Input channels.

forward(features: List[Tensor]) → List[Tensor]

Forward method.

training: bool

TorchOK Unet.

Adapted from: https://github.com/qubvel/segmentation_models.pytorch/tree/master/segmentation_models_pytorch/decoders/unet Licensed under MIT license [see LICENSE for details]

class torchok.models.necks.segmentation.unet.UnetNeck(in_channels: List[int], decoder_channels: List[int] = (512, 256, 128, 64, 64), use_batchnorm: bool = True, use_attention: bool = False, center: bool = True)

Bases: BaseModel

Unet is a fully convolutional neural network for image semantic segmentation. Paper: https://arxiv.org/pdf/1505.04597.

__init__(in_channels: List[int], decoder_channels: List[int] = (512, 256, 128, 64, 64), use_batchnorm: bool = True, use_attention: bool = False, center: bool = True)

Init Unet. The number of stages used in decoder inferred from the decoder_channels, larger depth - more features are generated. e.g. for depth=3 encoder will generate list of features with following spatial shapes [(H,W), (H/2, W/2), (H/4, W/4), (H/8, W/8)], so in general the deepest feature tensor will have spatial resolution (H/(2^depth), W/(2^depth)]

Parameters

  • in_channels – List of Numbers of Conv2D layer filters from backbone.

  • decoder_channels – List of numbers of Conv2D layer filters in decoder blocks.

  • use_batchnorm – If True, BatchNormalisation applied between every Conv2D and activation layers.

  • use_attention – If True will use SCSEModule.

  • center – If True will use ‘’CenterBlock’’.

Raises

ValueError – If the number of blocks is not equal to the length of the decoder_channels or encoder_channels - 1.

forward(features: List[Tensor]) → List[Tensor]

Forward method.

training: bool

Detection

TorchOk supports many mmdetection necks as long as custom necks. mmdetection necks can be accessed from the registry with their original names (see the link above and follow the corresponding class).

Custom detection necks