VigClsHead

class mmpretrain.models.heads.VigClsHead(num_classes, in_channels, hidden_dim=1024, act_cfg={'type': 'GELU'}, dropout=0.0, **kwargs)

The classification head for Vision GNN.

参数:

• num_classes (int) – Number of categories excluding the background category.

• in_channels (int) – Number of channels in the input feature map.

• hidden_dim (int) – The number of middle channels. Defaults to 1024.

• act_cfg (dict) – The config of activation function. Defaults to dict(type='GELU').

• dropout (float) – The dropout rate.

• loss (dict) – Config of classification loss. Defaults to dict(type='CrossEntropyLoss', loss_weight=1.0).

• init_cfg (dict, optional) – the config to control the initialization. Defaults to None.

forward(feats)

The forward process.

pre_logits(feats)

The process before the final classification head.

The input feats is a tuple of tensor, and each tensor is the feature of a stage_blocks stage. In VigClsHead, we just obtain the feature of the last stage.