MultiTaskHead

class mmpretrain.models.heads.MultiTaskHead(task_heads, init_cfg=None, **kwargs)

Multi task head.

Parameters:

• task_heads (dict) – Sub heads to use, the key will be use to rename the loss components.

• common_cfg (dict) – The common settings for all heads. Defaults to an empty dict.

• init_cfg (dict, optional) – The extra initialization settings. Defaults to None.

forward(feats)

The forward process.

loss(feats, data_samples, **kwargs)

Calculate losses from the classification score.

Parameters:

• feats (tuple[Tensor]) – The features extracted from the backbone.

• data_samples (List[MultiTaskDataSample]) – The annotation data of every samples.

• **kwargs – Other keyword arguments to forward the loss module.

Returns:

a dictionary of loss components, each task loss

key will be prefixed by the task_name like “task1_loss”

Return type:

dict[str, Tensor]

predict(feats, data_samples=None)

Inference without augmentation.

Parameters:

• feats (tuple[Tensor]) – The features extracted from the backbone.

• data_samples (List[MultiTaskDataSample], optional) – The annotation data of every samples. If not None, set pred_label of the input data samples. Defaults to None.

Returns:

A list of data samples which contains the predicted results.

Return type:

List[MultiTaskDataSample]