| | |
| | 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 |
| |
|
