LatentPredictHead

class mmpretrain.models.heads.LatentPredictHead(loss, predictor, init_cfg=None)

Head for latent feature prediction.

This head builds a predictor, which can be any registered neck component. For example, BYOL and SimSiam call this head and build NonLinearNeck. It also implements similarity loss between two forward features.

Parameters:

• loss (dict) – Config dict for the loss.

• predictor (dict) – Config dict for the predictor.

• init_cfg (dict or List[dict], optional) – Initialization config dict. Defaults to None.

loss(input, target)

Forward head.

Parameters:

• input (torch.Tensor) – NxC input features.

• target (torch.Tensor) – NxC target features.

Returns:

The latent predict loss.

Return type:

torch.Tensor