|
|
|
import copy |
|
from typing import Dict, List, Optional, Tuple, Union |
|
|
|
import torch |
|
import torch.nn as nn |
|
from torch import Tensor |
|
|
|
from mmdet.models.utils import (filter_gt_instances, rename_loss_dict, |
|
reweight_loss_dict) |
|
from mmdet.registry import MODELS |
|
from mmdet.structures import SampleList |
|
from mmdet.structures.bbox import bbox_project |
|
from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig |
|
from .base import BaseDetector |
|
|
|
|
|
@MODELS.register_module() |
|
class SemiBaseDetector(BaseDetector): |
|
"""Base class for semi-supervised detectors. |
|
|
|
Semi-supervised detectors typically consisting of a teacher model |
|
updated by exponential moving average and a student model updated |
|
by gradient descent. |
|
|
|
Args: |
|
detector (:obj:`ConfigDict` or dict): The detector config. |
|
semi_train_cfg (:obj:`ConfigDict` or dict, optional): |
|
The semi-supervised training config. |
|
semi_test_cfg (:obj:`ConfigDict` or dict, optional): |
|
The semi-supervised testing config. |
|
data_preprocessor (:obj:`ConfigDict` or dict, optional): Config of |
|
:class:`DetDataPreprocessor` to process the input data. |
|
Defaults to None. |
|
init_cfg (:obj:`ConfigDict` or list[:obj:`ConfigDict`] or dict or |
|
list[dict], optional): Initialization config dict. |
|
Defaults to None. |
|
""" |
|
|
|
def __init__(self, |
|
detector: ConfigType, |
|
semi_train_cfg: OptConfigType = None, |
|
semi_test_cfg: OptConfigType = None, |
|
data_preprocessor: OptConfigType = None, |
|
init_cfg: OptMultiConfig = None) -> None: |
|
super().__init__( |
|
data_preprocessor=data_preprocessor, init_cfg=init_cfg) |
|
self.student = MODELS.build(detector) |
|
self.teacher = MODELS.build(detector) |
|
self.semi_train_cfg = semi_train_cfg |
|
self.semi_test_cfg = semi_test_cfg |
|
if self.semi_train_cfg.get('freeze_teacher', True) is True: |
|
self.freeze(self.teacher) |
|
|
|
@staticmethod |
|
def freeze(model: nn.Module): |
|
"""Freeze the model.""" |
|
model.eval() |
|
for param in model.parameters(): |
|
param.requires_grad = False |
|
|
|
def loss(self, multi_batch_inputs: Dict[str, Tensor], |
|
multi_batch_data_samples: Dict[str, SampleList]) -> dict: |
|
"""Calculate losses from multi-branch inputs and data samples. |
|
|
|
Args: |
|
multi_batch_inputs (Dict[str, Tensor]): The dict of multi-branch |
|
input images, each value with shape (N, C, H, W). |
|
Each value should usually be mean centered and std scaled. |
|
multi_batch_data_samples (Dict[str, List[:obj:`DetDataSample`]]): |
|
The dict of multi-branch data samples. |
|
|
|
Returns: |
|
dict: A dictionary of loss components |
|
""" |
|
losses = dict() |
|
losses.update(**self.loss_by_gt_instances( |
|
multi_batch_inputs['sup'], multi_batch_data_samples['sup'])) |
|
|
|
origin_pseudo_data_samples, batch_info = self.get_pseudo_instances( |
|
multi_batch_inputs['unsup_teacher'], |
|
multi_batch_data_samples['unsup_teacher']) |
|
multi_batch_data_samples[ |
|
'unsup_student'] = self.project_pseudo_instances( |
|
origin_pseudo_data_samples, |
|
multi_batch_data_samples['unsup_student']) |
|
losses.update(**self.loss_by_pseudo_instances( |
|
multi_batch_inputs['unsup_student'], |
|
multi_batch_data_samples['unsup_student'], batch_info)) |
|
return losses |
|
|
|
def loss_by_gt_instances(self, batch_inputs: Tensor, |
|
batch_data_samples: SampleList) -> dict: |
|
"""Calculate losses from a batch of inputs and ground-truth data |
|
samples. |
|
|
|
Args: |
|
batch_inputs (Tensor): Input images of shape (N, C, H, W). |
|
These should usually be mean centered and std scaled. |
|
batch_data_samples (List[:obj:`DetDataSample`]): The batch |
|
data samples. It usually includes information such |
|
as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`. |
|
|
|
Returns: |
|
dict: A dictionary of loss components |
|
""" |
|
|
|
losses = self.student.loss(batch_inputs, batch_data_samples) |
|
sup_weight = self.semi_train_cfg.get('sup_weight', 1.) |
|
return rename_loss_dict('sup_', reweight_loss_dict(losses, sup_weight)) |
|
|
|
def loss_by_pseudo_instances(self, |
|
batch_inputs: Tensor, |
|
batch_data_samples: SampleList, |
|
batch_info: Optional[dict] = None) -> dict: |
|
"""Calculate losses from a batch of inputs and pseudo data samples. |
|
|
|
Args: |
|
batch_inputs (Tensor): Input images of shape (N, C, H, W). |
|
These should usually be mean centered and std scaled. |
|
batch_data_samples (List[:obj:`DetDataSample`]): The batch |
|
data samples. It usually includes information such |
|
as `gt_instance` or `gt_panoptic_seg` or `gt_sem_seg`, |
|
which are `pseudo_instance` or `pseudo_panoptic_seg` |
|
or `pseudo_sem_seg` in fact. |
|
batch_info (dict): Batch information of teacher model |
|
forward propagation process. Defaults to None. |
|
|
|
Returns: |
|
dict: A dictionary of loss components |
|
""" |
|
batch_data_samples = filter_gt_instances( |
|
batch_data_samples, score_thr=self.semi_train_cfg.cls_pseudo_thr) |
|
losses = self.student.loss(batch_inputs, batch_data_samples) |
|
pseudo_instances_num = sum([ |
|
len(data_samples.gt_instances) |
|
for data_samples in batch_data_samples |
|
]) |
|
unsup_weight = self.semi_train_cfg.get( |
|
'unsup_weight', 1.) if pseudo_instances_num > 0 else 0. |
|
return rename_loss_dict('unsup_', |
|
reweight_loss_dict(losses, unsup_weight)) |
|
|
|
@torch.no_grad() |
|
def get_pseudo_instances( |
|
self, batch_inputs: Tensor, batch_data_samples: SampleList |
|
) -> Tuple[SampleList, Optional[dict]]: |
|
"""Get pseudo instances from teacher model.""" |
|
self.teacher.eval() |
|
results_list = self.teacher.predict( |
|
batch_inputs, batch_data_samples, rescale=False) |
|
batch_info = {} |
|
for data_samples, results in zip(batch_data_samples, results_list): |
|
data_samples.gt_instances = results.pred_instances |
|
data_samples.gt_instances.bboxes = bbox_project( |
|
data_samples.gt_instances.bboxes, |
|
torch.from_numpy(data_samples.homography_matrix).inverse().to( |
|
self.data_preprocessor.device), data_samples.ori_shape) |
|
return batch_data_samples, batch_info |
|
|
|
def project_pseudo_instances(self, batch_pseudo_instances: SampleList, |
|
batch_data_samples: SampleList) -> SampleList: |
|
"""Project pseudo instances.""" |
|
for pseudo_instances, data_samples in zip(batch_pseudo_instances, |
|
batch_data_samples): |
|
data_samples.gt_instances = copy.deepcopy( |
|
pseudo_instances.gt_instances) |
|
data_samples.gt_instances.bboxes = bbox_project( |
|
data_samples.gt_instances.bboxes, |
|
torch.tensor(data_samples.homography_matrix).to( |
|
self.data_preprocessor.device), data_samples.img_shape) |
|
wh_thr = self.semi_train_cfg.get('min_pseudo_bbox_wh', (1e-2, 1e-2)) |
|
return filter_gt_instances(batch_data_samples, wh_thr=wh_thr) |
|
|
|
def predict(self, batch_inputs: Tensor, |
|
batch_data_samples: SampleList) -> SampleList: |
|
"""Predict results from a batch of inputs and data samples with post- |
|
processing. |
|
|
|
Args: |
|
batch_inputs (Tensor): Inputs with shape (N, C, H, W). |
|
batch_data_samples (List[:obj:`DetDataSample`]): The Data |
|
Samples. It usually includes information such as |
|
`gt_instance`, `gt_panoptic_seg` and `gt_sem_seg`. |
|
rescale (bool): Whether to rescale the results. |
|
Defaults to True. |
|
|
|
Returns: |
|
list[:obj:`DetDataSample`]: Return the detection results of the |
|
input images. The returns value is DetDataSample, |
|
which usually contain 'pred_instances'. And the |
|
``pred_instances`` usually contains following keys. |
|
|
|
- scores (Tensor): Classification scores, has a shape |
|
(num_instance, ) |
|
- labels (Tensor): Labels of bboxes, has a shape |
|
(num_instances, ). |
|
- bboxes (Tensor): Has a shape (num_instances, 4), |
|
the last dimension 4 arrange as (x1, y1, x2, y2). |
|
- masks (Tensor): Has a shape (num_instances, H, W). |
|
""" |
|
if self.semi_test_cfg.get('predict_on', 'teacher') == 'teacher': |
|
return self.teacher( |
|
batch_inputs, batch_data_samples, mode='predict') |
|
else: |
|
return self.student( |
|
batch_inputs, batch_data_samples, mode='predict') |
|
|
|
def _forward(self, batch_inputs: Tensor, |
|
batch_data_samples: SampleList) -> SampleList: |
|
"""Network forward process. Usually includes backbone, neck and head |
|
forward without any post-processing. |
|
|
|
Args: |
|
batch_inputs (Tensor): Inputs with shape (N, C, H, W). |
|
|
|
Returns: |
|
tuple: A tuple of features from ``rpn_head`` and ``roi_head`` |
|
forward. |
|
""" |
|
if self.semi_test_cfg.get('forward_on', 'teacher') == 'teacher': |
|
return self.teacher( |
|
batch_inputs, batch_data_samples, mode='tensor') |
|
else: |
|
return self.student( |
|
batch_inputs, batch_data_samples, mode='tensor') |
|
|
|
def extract_feat(self, batch_inputs: Tensor) -> Tuple[Tensor]: |
|
"""Extract features. |
|
|
|
Args: |
|
batch_inputs (Tensor): Image tensor with shape (N, C, H ,W). |
|
|
|
Returns: |
|
tuple[Tensor]: Multi-level features that may have |
|
different resolutions. |
|
""" |
|
if self.semi_test_cfg.get('extract_feat_on', 'teacher') == 'teacher': |
|
return self.teacher.extract_feat(batch_inputs) |
|
else: |
|
return self.student.extract_feat(batch_inputs) |
|
|
|
def _load_from_state_dict(self, state_dict: dict, prefix: str, |
|
local_metadata: dict, strict: bool, |
|
missing_keys: Union[List[str], str], |
|
unexpected_keys: Union[List[str], str], |
|
error_msgs: Union[List[str], str]) -> None: |
|
"""Add teacher and student prefixes to model parameter names.""" |
|
if not any([ |
|
'student' in key or 'teacher' in key |
|
for key in state_dict.keys() |
|
]): |
|
keys = list(state_dict.keys()) |
|
state_dict.update({'teacher.' + k: state_dict[k] for k in keys}) |
|
state_dict.update({'student.' + k: state_dict[k] for k in keys}) |
|
for k in keys: |
|
state_dict.pop(k) |
|
return super()._load_from_state_dict( |
|
state_dict, |
|
prefix, |
|
local_metadata, |
|
strict, |
|
missing_keys, |
|
unexpected_keys, |
|
error_msgs, |
|
) |
|
|
