models.generator.Head

Contents

models.generator.Head#

class Head(cfg: str | ListGen, num_classes: int, in_shape: List[int], init_weights: bool = True)[source]#

Bases: Module

Head model holder

Applies a head model to multiple maps and merges them. For each input map, its own HeadGen model is generated. Predictions obtained for different feature maps are combined.

Parameters:

  • cfg (str | ListGen) – Lists of layer generators.

  • num_classes (int) – Number of classes.

  • in_shape (List[int]) – Input data format. Expecting to receive a list of models.neck.NeckGen.out_shape

  • init_weights (bool, optional) – If true apply weight initialization function. Defaults to True.

Methods

forward

Direct network pass

Attributes

training

forward(X: List[Tensor]) → tuple[Tensor, Tensor, Tensor][source]#

Direct network pass

Parameters:

X (List[torch.Tensor]) – Feature map list. One map shape [ts, batch, channel, h, w].

Returns:

Predictions made by a neural network. Contains three tensors:

  1. anchors: Shape [anchor, 4]

  2. cls_preds: Shape [ts, batch, anchor, num_classes + 1]

  3. bbox_preds: Shape [ts, batch, anchor, 4]

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Contents