Heads

Classification

class torchok.models.heads.classification.classification_head.ClassificationHead(in_channels: int, num_classes: int, drop_rate: float = 0.0, bias: bool = True)

Bases: LinearHead

Classification head for basic input features.

__init__(in_channels: int, num_classes: int, drop_rate: float = 0.0, bias: bool = True)

Init ClassificationHead.

Shape of the input features is (, in_features) and shape of the output is (, num_classes), where * means any number of dimensions. At output the logits are returned, i.e. no softmax is applied. When number of classes is equal to 1, the classification task is considered as being a binary classification, so the channels dimension is squeezed.

Parameters

  • in_channels – number of input features

  • num_classes – number of classes

  • drop_rate – dropout rate (applied before linear layer)

  • bias – whether to use bias in the linear layer

forward(x: Tensor, target: Optional[Tensor] = None) → Tensor

Forward single input x.

Parameters

x – Input tensor.

training: bool
class torchok.models.heads.classification.arcface_head.ArcFaceHead(in_channels: int, num_classes: int, scale: Optional[float] = None, margin: Optional[float] = None, easy_margin: bool = False, dynamic_margin: bool = False, num_warmup_steps: Optional[int] = None, min_margin: Optional[float] = None)

Bases: BaseModel

Implement of arc margin distance. Classification head.

ArcFace paper: https://arxiv.org/pdf/1801.07698.pdf Code: https://github.com/ronghuaiyang/arcface-pytorch/blob/master/models/metrics.py

__init__(in_channels: int, num_classes: int, scale: Optional[float] = None, margin: Optional[float] = None, easy_margin: bool = False, dynamic_margin: bool = False, num_warmup_steps: Optional[int] = None, min_margin: Optional[float] = None)

Init ArcFaceHead class.

Parameters

  • in_channels – Size of each input sample.

  • num_classes – number of classes.

  • scale – Feature scale.

  • margin – Angular margin.

  • easy_margin – Easy margin.

  • dynamic_margin – If True margin will increase from min_margin to margin, step by step(num_warmup_steps times).

  • num_warmup_steps – Number steps with dynamic margin.

  • min_margin – Initial margin in dynamic_margin mode.

Raises

ValueError – if num_warmup_steps or min_margin is None, when dynamic_margin is True.

forward(input: Tensor, target: Optional[Tensor] = None) → Tensor

Forward method.

Parameters

  • input – Input tensor.

  • target – Target tensor. It may be None, if training mode off (inference stage).

Raises

ValueError – If training mode on and target is None.

training: bool

Segmentation

class torchok.models.heads.segmentation.base.SegmentationHead(in_channels: int, num_classes: int, do_interpolate: bool = True)

Bases: BaseModel

Base Segmentation head.

__init__(in_channels: int, num_classes: int, do_interpolate: bool = True)

Init SegmentationHead.

Parameters

  • in_channels – Size of each input sample.

  • num_classes – A number of classes for output (output shape - (batch, classes, h, w)).

  • do_interpolate – If True will interpolate features after forward pass.

forward(x: List[Tensor]) → Tensor

Forward method.

training: bool

This HRNet implementation is modified from the following repository: https://github.com/HRNet/HRNet-Semantic-Segmentation

class torchok.models.heads.segmentation.ocr.OCRSegmentationHead(in_channels: int, num_classes: int, do_interpolate: bool = True, ocr_mid_channels=128, ocr_key_channels=64)

Bases: BaseModel

Implementation of HRNet segmentation head with Object-Contextual Representations for Semantic Segmentation from https://github.com/HRNet/HRNet-Semantic-Segmentation/tree/HRNet-OCR.

__init__(in_channels: int, num_classes: int, do_interpolate: bool = True, ocr_mid_channels=128, ocr_key_channels=64)

Init OCRSegmentationHead. :param in_channels: Number of channels from segmentation neck. :param num_classes: Number of segmentation classes. :param ocr_mid_channels: Number of intermediate feature channels. :param ocr_key_channels: Number of channels in the dimension after the key/query transform.

forward(feats: Tensor) → Union[Tensor, Tuple[Tensor, Tensor]]
training: bool

Representation

class torchok.models.heads.representation.linear_head.LinearHead(in_channels, out_channels, drop_rate=0.0, bias=True, normalize=False)

Bases: BaseModel

Linear Head

__init__(in_channels, out_channels, drop_rate=0.0, bias=True, normalize=False)

Init LinearHead. :param in_channels: Input channels. :param out_channels: Output channels. :param drop_rate: Drop rate. :param bias: Bias. :param normalize: Normalize.

forward(x: Tensor, targets: Optional[Tensor] = None) → Tensor
training: bool

Detection