mmpretrain.models.heads.cls_head 源代码
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmengine.model import BaseModule
from mmpretrain.evaluation.metrics import Accuracy
from mmpretrain.registry import MODELS
from mmpretrain.structures import DataSample
[文档]@MODELS.register_module()
class ClsHead(BaseModule):
"""Classification head.
Args:
loss (dict): Config of classification loss. Defaults to
``dict(type='CrossEntropyLoss', loss_weight=1.0)``.
topk (int | Tuple[int]): Top-k accuracy. Defaults to ``(1, )``.
cal_acc (bool): Whether to calculate accuracy during training.
If you use batch augmentations like Mixup and CutMix during
training, it is pointless to calculate accuracy.
Defaults to False.
init_cfg (dict, optional): the config to control the initialization.
Defaults to None.
"""
def __init__(self,
loss: dict = dict(type='CrossEntropyLoss', loss_weight=1.0),
topk: Union[int, Tuple[int]] = (1, ),
cal_acc: bool = False,
init_cfg: Optional[dict] = None):
super(ClsHead, self).__init__(init_cfg=init_cfg)
self.topk = topk
if not isinstance(loss, nn.Module):
loss = MODELS.build(loss)
self.loss_module = loss
self.cal_acc = cal_acc
[文档] def pre_logits(self, feats: Tuple[torch.Tensor]) -> torch.Tensor:
"""The process before the final classification head.
The input ``feats`` is a tuple of tensor, and each tensor is the
feature of a backbone stage. In ``ClsHead``, we just obtain the feature
of the last stage.
"""
# The ClsHead doesn't have other module, just return after unpacking.
return feats[-1]
[文档] def forward(self, feats: Tuple[torch.Tensor]) -> torch.Tensor:
"""The forward process."""
pre_logits = self.pre_logits(feats)
# The ClsHead doesn't have the final classification head,
# just return the unpacked inputs.
return pre_logits
[文档] def loss(self, feats: Tuple[torch.Tensor], data_samples: List[DataSample],
**kwargs) -> dict:
"""Calculate losses from the classification score.
Args:
feats (tuple[Tensor]): The features extracted from the backbone.
Multiple stage inputs are acceptable but only the last stage
will be used to classify. The shape of every item should be
``(num_samples, num_classes)``.
data_samples (List[DataSample]): The annotation data of
every samples.
**kwargs: Other keyword arguments to forward the loss module.
Returns:
dict[str, Tensor]: a dictionary of loss components
"""
# The part can be traced by torch.fx
cls_score = self(feats)
# The part can not be traced by torch.fx
losses = self._get_loss(cls_score, data_samples, **kwargs)
return losses
def _get_loss(self, cls_score: torch.Tensor,
data_samples: List[DataSample], **kwargs):
"""Unpack data samples and compute loss."""
# Unpack data samples and pack targets
if 'gt_score' in data_samples[0]:
# Batch augmentation may convert labels to one-hot format scores.
target = torch.stack([i.gt_score for i in data_samples])
else:
target = torch.cat([i.gt_label for i in data_samples])
# compute loss
losses = dict()
loss = self.loss_module(
cls_score, target, avg_factor=cls_score.size(0), **kwargs)
losses['loss'] = loss
# compute accuracy
if self.cal_acc:
assert target.ndim == 1, 'If you enable batch augmentation ' \
'like mixup during training, `cal_acc` is pointless.'
acc = Accuracy.calculate(cls_score, target, topk=self.topk)
losses.update(
{f'accuracy_top-{k}': a
for k, a in zip(self.topk, acc)})
return losses
[文档] def predict(
self,
feats: Tuple[torch.Tensor],
data_samples: Optional[List[Optional[DataSample]]] = None
) -> List[DataSample]:
"""Inference without augmentation.
Args:
feats (tuple[Tensor]): The features extracted from the backbone.
Multiple stage inputs are acceptable but only the last stage
will be used to classify. The shape of every item should be
``(num_samples, num_classes)``.
data_samples (List[DataSample | None], optional): The annotation
data of every samples. If not None, set ``pred_label`` of
the input data samples. Defaults to None.
Returns:
List[DataSample]: A list of data samples which contains the
predicted results.
"""
# The part can be traced by torch.fx
cls_score = self(feats)
# The part can not be traced by torch.fx
predictions = self._get_predictions(cls_score, data_samples)
return predictions
def _get_predictions(self, cls_score, data_samples):
"""Post-process the output of head.
Including softmax and set ``pred_label`` of data samples.
"""
pred_scores = F.softmax(cls_score, dim=1)
pred_labels = pred_scores.argmax(dim=1, keepdim=True).detach()
out_data_samples = []
if data_samples is None:
data_samples = [None for _ in range(pred_scores.size(0))]
for data_sample, score, label in zip(data_samples, pred_scores,
pred_labels):
if data_sample is None:
data_sample = DataSample()
data_sample.set_pred_score(score).set_pred_label(label)
out_data_samples.append(data_sample)
return out_data_samples