Shortcuts

mmpretrain.models.heads.swav_head 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import torch
from mmengine.model import BaseModule

from mmpretrain.registry import MODELS


[文档]@MODELS.register_module() class SwAVHead(BaseModule): """Head for SwAV Pre-training. Args: loss (dict): Config dict for module of loss functions. """ def __init__(self, loss: dict) -> None: super().__init__() self.loss_module = MODELS.build(loss)
[文档] def loss(self, pred: torch.Tensor) -> torch.Tensor: """Generate loss. Args: pred (torch.Tensor): NxC input features. Returns: torch.Tensor: The SwAV loss. """ loss = self.loss_module(pred) return loss
Read the Docs v: dev
Versions
latest
stable
mmcls-1.x
mmcls-0.x
dev
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.