dnt 0.2.1__py3-none-any.whl → 0.3.1.8__py3-none-any.whl

dnt/third_party/fast-reid/fastreid/modeling/meta_arch/baseline.py

@@ -0,0 +1,188 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+import torch
+from torch import nn
+
+from fastreid.config import configurable
+from fastreid.modeling.backbones import build_backbone
+from fastreid.modeling.heads import build_heads
+from fastreid.modeling.losses import *
+from .build import META_ARCH_REGISTRY
+
+
+@META_ARCH_REGISTRY.register()
+class Baseline(nn.Module):
+    """
+    Baseline architecture. Any models that contains the following two components:
+    1. Per-image feature extraction (aka backbone)
+    2. Per-image feature aggregation and loss computation
+    """
+
+    @configurable
+    def __init__(
+            self,
+            *,
+            backbone,
+            heads,
+            pixel_mean,
+            pixel_std,
+            loss_kwargs=None
+    ):
+        """
+        NOTE: this interface is experimental.
+
+        Args:
+            backbone:
+            heads:
+            pixel_mean:
+            pixel_std:
+        """
+        super().__init__()
+        # backbone
+        self.backbone = backbone
+
+        # head
+        self.heads = heads
+
+        self.loss_kwargs = loss_kwargs
+
+        self.register_buffer('pixel_mean', torch.Tensor(pixel_mean).view(1, -1, 1, 1), False)
+        self.register_buffer('pixel_std', torch.Tensor(pixel_std).view(1, -1, 1, 1), False)
+
+    @classmethod
+    def from_config(cls, cfg):
+        backbone = build_backbone(cfg)
+        heads = build_heads(cfg)
+        return {
+            'backbone': backbone,
+            'heads': heads,
+            'pixel_mean': cfg.MODEL.PIXEL_MEAN,
+            'pixel_std': cfg.MODEL.PIXEL_STD,
+            'loss_kwargs':
+                {
+                    # loss name
+                    'loss_names': cfg.MODEL.LOSSES.NAME,
+
+                    # loss hyperparameters
+                    'ce': {
+                        'eps': cfg.MODEL.LOSSES.CE.EPSILON,
+                        'alpha': cfg.MODEL.LOSSES.CE.ALPHA,
+                        'scale': cfg.MODEL.LOSSES.CE.SCALE
+                    },
+                    'tri': {
+                        'margin': cfg.MODEL.LOSSES.TRI.MARGIN,
+                        'norm_feat': cfg.MODEL.LOSSES.TRI.NORM_FEAT,
+                        'hard_mining': cfg.MODEL.LOSSES.TRI.HARD_MINING,
+                        'scale': cfg.MODEL.LOSSES.TRI.SCALE
+                    },
+                    'circle': {
+                        'margin': cfg.MODEL.LOSSES.CIRCLE.MARGIN,
+                        'gamma': cfg.MODEL.LOSSES.CIRCLE.GAMMA,
+                        'scale': cfg.MODEL.LOSSES.CIRCLE.SCALE
+                    },
+                    'cosface': {
+                        'margin': cfg.MODEL.LOSSES.COSFACE.MARGIN,
+                        'gamma': cfg.MODEL.LOSSES.COSFACE.GAMMA,
+                        'scale': cfg.MODEL.LOSSES.COSFACE.SCALE
+                    }
+                }
+        }
+
+    @property
+    def device(self):
+        return self.pixel_mean.device
+
+    def forward(self, batched_inputs):
+        images = self.preprocess_image(batched_inputs)
+        features = self.backbone(images)
+
+        if self.training:
+            assert "targets" in batched_inputs, "Person ID annotation are missing in training!"
+            targets = batched_inputs["targets"]
+
+            # PreciseBN flag, When do preciseBN on different dataset, the number of classes in new dataset
+            # may be larger than that in the original dataset, so the circle/arcface will
+            # throw an error. We just set all the targets to 0 to avoid this problem.
+            if targets.sum() < 0: targets.zero_()
+
+            outputs = self.heads(features, targets)
+            losses = self.losses(outputs, targets)
+            return losses
+        else:
+            outputs = self.heads(features)
+            return outputs
+
+    def preprocess_image(self, batched_inputs):
+        """
+        Normalize and batch the input images.
+        """
+        if isinstance(batched_inputs, dict):
+            images = batched_inputs['images']
+        elif isinstance(batched_inputs, torch.Tensor):
+            images = batched_inputs
+        else:
+            raise TypeError("batched_inputs must be dict or torch.Tensor, but get {}".format(type(batched_inputs)))
+
+        images.sub_(self.pixel_mean).div_(self.pixel_std)
+        return images
+
+    def losses(self, outputs, gt_labels):
+        """
+        Compute loss from modeling's outputs, the loss function input arguments
+        must be the same as the outputs of the model forwarding.
+        """
+        # model predictions
+        # fmt: off
+        pred_class_logits = outputs['pred_class_logits'].detach()
+        cls_outputs       = outputs['cls_outputs']
+        pred_features     = outputs['features']
+        # fmt: on
+
+        # Log prediction accuracy
+        log_accuracy(pred_class_logits, gt_labels)
+
+        loss_dict = {}
+        loss_names = self.loss_kwargs['loss_names']
+
+        if 'CrossEntropyLoss' in loss_names:
+            ce_kwargs = self.loss_kwargs.get('ce')
+            loss_dict['loss_cls'] = cross_entropy_loss(
+                cls_outputs,
+                gt_labels,
+                ce_kwargs.get('eps'),
+                ce_kwargs.get('alpha')
+            ) * ce_kwargs.get('scale')
+
+        if 'TripletLoss' in loss_names:
+            tri_kwargs = self.loss_kwargs.get('tri')
+            loss_dict['loss_triplet'] = triplet_loss(
+                pred_features,
+                gt_labels,
+                tri_kwargs.get('margin'),
+                tri_kwargs.get('norm_feat'),
+                tri_kwargs.get('hard_mining')
+            ) * tri_kwargs.get('scale')
+
+        if 'CircleLoss' in loss_names:
+            circle_kwargs = self.loss_kwargs.get('circle')
+            loss_dict['loss_circle'] = pairwise_circleloss(
+                pred_features,
+                gt_labels,
+                circle_kwargs.get('margin'),
+                circle_kwargs.get('gamma')
+            ) * circle_kwargs.get('scale')
+
+        if 'Cosface' in loss_names:
+            cosface_kwargs = self.loss_kwargs.get('cosface')
+            loss_dict['loss_cosface'] = pairwise_cosface(
+                pred_features,
+                gt_labels,
+                cosface_kwargs.get('margin'),
+                cosface_kwargs.get('gamma'),
+            ) * cosface_kwargs.get('scale')
+
+        return loss_dict

dnt/third_party/fast-reid/fastreid/modeling/meta_arch/build.py

@@ -0,0 +1,26 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+import torch
+
+from fastreid.utils.registry import Registry
+
+META_ARCH_REGISTRY = Registry("META_ARCH")  # noqa F401 isort:skip
+META_ARCH_REGISTRY.__doc__ = """
+Registry for meta-architectures, i.e. the whole model.
+The registered object will be called with `obj(cfg)`
+and expected to return a `nn.Module` object.
+"""
+
+
+def build_model(cfg):
+    """
+    Build the whole model architecture, defined by ``cfg.MODEL.META_ARCHITECTURE``.
+    Note that it does not load any weights from ``cfg``.
+    """
+    meta_arch = cfg.MODEL.META_ARCHITECTURE
+    model = META_ARCH_REGISTRY.get(meta_arch)(cfg)
+    model.to(torch.device(cfg.MODEL.DEVICE))
+    return model

dnt/third_party/fast-reid/fastreid/modeling/meta_arch/distiller.py

@@ -0,0 +1,140 @@
+# encoding: utf-8
+"""
+@author:  l1aoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+import logging
+
+import torch
+import torch.nn.functional as F
+
+from fastreid.config import get_cfg
+from fastreid.modeling.meta_arch import META_ARCH_REGISTRY, build_model, Baseline
+from fastreid.utils.checkpoint import Checkpointer
+
+logger = logging.getLogger(__name__)
+
+
+@META_ARCH_REGISTRY.register()
+class Distiller(Baseline):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+        # Get teacher model config
+        model_ts = []
+        for i in range(len(cfg.KD.MODEL_CONFIG)):
+            cfg_t = get_cfg()
+            cfg_t.merge_from_file(cfg.KD.MODEL_CONFIG[i])
+            cfg_t.defrost()
+            cfg_t.MODEL.META_ARCHITECTURE = "Baseline"
+            # Change syncBN to BN due to no DDP wrapper
+            if cfg_t.MODEL.BACKBONE.NORM == "syncBN":
+                cfg_t.MODEL.BACKBONE.NORM = "BN"
+            if cfg_t.MODEL.HEADS.NORM == "syncBN":
+                cfg_t.MODEL.HEADS.NORM = "BN"
+
+            model_t = build_model(cfg_t)
+
+            # No gradients for teacher model
+            for param in model_t.parameters():
+                param.requires_grad_(False)
+
+            logger.info("Loading teacher model weights ...")
+            Checkpointer(model_t).load(cfg.KD.MODEL_WEIGHTS[i])
+
+            model_ts.append(model_t)
+
+        self.ema_enabled = cfg.KD.EMA.ENABLED
+        self.ema_momentum = cfg.KD.EMA.MOMENTUM
+        if self.ema_enabled:
+            cfg_self = cfg.clone()
+            cfg_self.defrost()
+            cfg_self.MODEL.META_ARCHITECTURE = "Baseline"
+            if cfg_self.MODEL.BACKBONE.NORM == "syncBN":
+                cfg_self.MODEL.BACKBONE.NORM = "BN"
+            if cfg_self.MODEL.HEADS.NORM == "syncBN":
+                cfg_self.MODEL.HEADS.NORM = "BN"
+            model_self = build_model(cfg_self)
+            # No gradients for self model
+            for param in model_self.parameters():
+                param.requires_grad_(False)
+
+            if cfg_self.MODEL.WEIGHTS != '':
+                logger.info("Loading self distillation model weights ...")
+                Checkpointer(model_self).load(cfg_self.MODEL.WEIGHTS)
+            else:
+                # Make sure the initial state is same
+                for param_q, param_k in zip(self.parameters(), model_self.parameters()):
+                    param_k.data.copy_(param_q.data)
+
+            model_ts.insert(0, model_self)
+
+        # Not register teacher model as `nn.Module`, this is
+        # make sure teacher model weights not saved
+        self.model_ts = model_ts
+
+    @torch.no_grad()
+    def _momentum_update_key_encoder(self, m=0.999):
+        """
+        Momentum update of the key encoder
+        """
+        for param_q, param_k in zip(self.parameters(), self.model_ts[0].parameters()):
+            param_k.data = param_k.data * m + param_q.data * (1. - m)
+
+    def forward(self, batched_inputs):
+        if self.training:
+            images = self.preprocess_image(batched_inputs)
+            # student model forward
+            s_feat = self.backbone(images)
+            assert "targets" in batched_inputs, "Labels are missing in training!"
+            targets = batched_inputs["targets"].to(self.device)
+
+            if targets.sum() < 0: targets.zero_()
+
+            s_outputs = self.heads(s_feat, targets)
+
+            t_outputs = []
+            # teacher model forward
+            with torch.no_grad():
+                if self.ema_enabled:
+                    self._momentum_update_key_encoder(self.ema_momentum)  # update self distill model
+                for model_t in self.model_ts:
+                    t_feat = model_t.backbone(images)
+                    t_output = model_t.heads(t_feat, targets)
+                    t_outputs.append(t_output)
+
+            losses = self.losses(s_outputs, t_outputs, targets)
+            return losses
+
+        # Eval mode, just conventional reid feature extraction
+        else:
+            return super().forward(batched_inputs)
+
+    def losses(self, s_outputs, t_outputs, gt_labels):
+        """
+        Compute loss from modeling's outputs, the loss function input arguments
+        must be the same as the outputs of the model forwarding.
+        """
+        loss_dict = super().losses(s_outputs, gt_labels)
+
+        s_logits = s_outputs['pred_class_logits']
+        loss_jsdiv = 0.
+        for t_output in t_outputs:
+            t_logits = t_output['pred_class_logits'].detach()
+            loss_jsdiv += self.jsdiv_loss(s_logits, t_logits)
+
+        loss_dict["loss_jsdiv"] = loss_jsdiv / len(t_outputs)
+
+        return loss_dict
+
+    @staticmethod
+    def _kldiv(y_s, y_t, t):
+        p_s = F.log_softmax(y_s / t, dim=1)
+        p_t = F.softmax(y_t / t, dim=1)
+        loss = F.kl_div(p_s, p_t, reduction="sum") * (t ** 2) / y_s.shape[0]
+        return loss
+
+    def jsdiv_loss(self, y_s, y_t, t=16):
+        loss = (self._kldiv(y_s, y_t, t) + self._kldiv(y_t, y_s, t)) / 2
+        return loss