| |
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| from mmengine.runner import load_checkpoint |
| from torch import Tensor |
|
|
| from mmdet.registry import MODELS |
| from mmdet.structures import SampleList |
| from mmdet.utils import ConfigType, OptConfigType |
| from ..utils.misc import unpack_gt_instances |
| from .kd_one_stage import KnowledgeDistillationSingleStageDetector |
|
|
|
|
| @MODELS.register_module() |
| class LAD(KnowledgeDistillationSingleStageDetector): |
| """Implementation of `LAD <https://arxiv.org/pdf/2108.10520.pdf>`_.""" |
|
|
| def __init__(self, |
| backbone: ConfigType, |
| neck: ConfigType, |
| bbox_head: ConfigType, |
| teacher_backbone: ConfigType, |
| teacher_neck: ConfigType, |
| teacher_bbox_head: ConfigType, |
| teacher_ckpt: Optional[str] = None, |
| eval_teacher: bool = True, |
| train_cfg: OptConfigType = None, |
| test_cfg: OptConfigType = None, |
| data_preprocessor: OptConfigType = None) -> None: |
| super(KnowledgeDistillationSingleStageDetector, self).__init__( |
| backbone=backbone, |
| neck=neck, |
| bbox_head=bbox_head, |
| train_cfg=train_cfg, |
| test_cfg=test_cfg, |
| data_preprocessor=data_preprocessor) |
| self.eval_teacher = eval_teacher |
| self.teacher_model = nn.Module() |
| self.teacher_model.backbone = MODELS.build(teacher_backbone) |
| if teacher_neck is not None: |
| self.teacher_model.neck = MODELS.build(teacher_neck) |
| teacher_bbox_head.update(train_cfg=train_cfg) |
| teacher_bbox_head.update(test_cfg=test_cfg) |
| self.teacher_model.bbox_head = MODELS.build(teacher_bbox_head) |
| if teacher_ckpt is not None: |
| load_checkpoint( |
| self.teacher_model, teacher_ckpt, map_location='cpu') |
|
|
| @property |
| def with_teacher_neck(self) -> bool: |
| """bool: whether the detector has a teacher_neck""" |
| return hasattr(self.teacher_model, 'neck') and \ |
| self.teacher_model.neck is not None |
|
|
| def extract_teacher_feat(self, batch_inputs: Tensor) -> Tensor: |
| """Directly extract teacher features from the backbone+neck.""" |
| x = self.teacher_model.backbone(batch_inputs) |
| if self.with_teacher_neck: |
| x = self.teacher_model.neck(x) |
| return x |
|
|
| def loss(self, batch_inputs: Tensor, |
| batch_data_samples: SampleList) -> dict: |
| """ |
| 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[str, Tensor]: A dictionary of loss components. |
| """ |
| outputs = unpack_gt_instances(batch_data_samples) |
| batch_gt_instances, batch_gt_instances_ignore, batch_img_metas \ |
| = outputs |
| |
| with torch.no_grad(): |
| x_teacher = self.extract_teacher_feat(batch_inputs) |
| outs_teacher = self.teacher_model.bbox_head(x_teacher) |
| label_assignment_results = \ |
| self.teacher_model.bbox_head.get_label_assignment( |
| *outs_teacher, batch_gt_instances, batch_img_metas, |
| batch_gt_instances_ignore) |
|
|
| |
| x = self.extract_feat(batch_inputs) |
| losses = self.bbox_head.loss(x, label_assignment_results, |
| batch_data_samples) |
| return losses |
|
|
