Enable to use mmcv.jit in parrots environment (open-mmlab#4192)

* add mmcv.jit

* only use derivate and coderize

* fix for isort

* small modify for decorator order

Author: lml131

mmdet/core/bbox/coder/bucketing_bbox_coder.py

@@ -1,3 +1,4 @@
+import mmcv
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -90,6 +91,7 @@ def decode(self, bboxes, pred_bboxes, max_shape=None):
         return decoded_bboxes
 
 
+@mmcv.jit(coderize=True)
 def generat_buckets(proposals, num_buckets, scale_factor=1.0):
     """Generate buckets w.r.t bucket number and scale factor of proposals.
 
@@ -138,6 +140,7 @@ def generat_buckets(proposals, num_buckets, scale_factor=1.0):
     return bucket_w, bucket_h, l_buckets, r_buckets, t_buckets, d_buckets
 
 
+@mmcv.jit(coderize=True)
 def bbox2bucket(proposals,
                 gt,
                 num_buckets,
@@ -261,6 +264,7 @@ def bbox2bucket(proposals,
     return offsets, offsets_weights, bucket_labels, bucket_cls_weights
 
 
+@mmcv.jit(coderize=True)
 def bucket2bbox(proposals,
                 cls_preds,
                 offset_preds,

mmdet/core/bbox/coder/delta_xywh_bbox_coder.py

@@ -1,3 +1,4 @@
+import mmcv
 import numpy as np
 import torch
 
@@ -75,6 +76,7 @@ def decode(self,
         return decoded_bboxes
 
 
+@mmcv.jit(coderize=True)
 def bbox2delta(proposals, gt, means=(0., 0., 0., 0.), stds=(1., 1., 1., 1.)):
     """Compute deltas of proposals w.r.t. gt.
 
@@ -120,6 +122,7 @@ def bbox2delta(proposals, gt, means=(0., 0., 0., 0.), stds=(1., 1., 1., 1.)):
     return deltas
 
 
+@mmcv.jit(coderize=True)
 def delta2bbox(rois,
                deltas,
                means=(0., 0., 0., 0.),

mmdet/core/bbox/coder/legacy_delta_xywh_bbox_coder.py

@@ -1,3 +1,4 @@
+import mmcv
 import numpy as np
 import torch
 
@@ -79,6 +80,7 @@ def decode(self,
         return decoded_bboxes
 
 
+@mmcv.jit(coderize=True)
 def legacy_bbox2delta(proposals,
                       gt,
                       means=(0., 0., 0., 0.),
@@ -127,6 +129,7 @@ def legacy_bbox2delta(proposals,
     return deltas
 
 
+@mmcv.jit(coderize=True)
 def legacy_delta2bbox(rois,
                       deltas,
                       means=(0., 0., 0., 0.),

mmdet/core/bbox/coder/tblr_bbox_coder.py

@@ -1,3 +1,4 @@
+import mmcv
 import torch
 
 from ..builder import BBOX_CODERS
@@ -68,6 +69,7 @@ def decode(self, bboxes, pred_bboxes, max_shape=None):
         return decoded_bboxes
 
 
+@mmcv.jit(coderize=True)
 def bboxes2tblr(priors, gts, normalizer=4.0, normalize_by_wh=True):
     """Encode ground truth boxes to tblr coordinate.
 
@@ -114,6 +116,7 @@ def bboxes2tblr(priors, gts, normalizer=4.0, normalize_by_wh=True):
     return loc / normalizer
 
 
+@mmcv.jit(coderize=True)
 def tblr2bboxes(priors,
                 tblr,
                 normalizer=4.0,

mmdet/core/bbox/coder/yolo_bbox_coder.py

@@ -1,3 +1,4 @@
+import mmcv
 import torch
 
 from ..builder import BBOX_CODERS
@@ -21,6 +22,7 @@ def __init__(self, eps=1e-6):
         super(BaseBBoxCoder, self).__init__()
         self.eps = eps
 
+    @mmcv.jit(coderize=True)
     def encode(self, bboxes, gt_bboxes, stride):
         """Get box regression transformation deltas that can be used to
         transform the ``bboxes`` into the ``gt_bboxes``.
@@ -55,6 +57,7 @@ def encode(self, bboxes, gt_bboxes, stride):
             [x_center_target, y_center_target, w_target, h_target], dim=-1)
         return encoded_bboxes
 
+    @mmcv.jit(coderize=True)
     def decode(self, bboxes, pred_bboxes, stride):
         """Apply transformation `pred_bboxes` to `boxes`.
 

mmdet/models/losses/accuracy.py

@@ -1,6 +1,8 @@
+import mmcv
 import torch.nn as nn
 
 
+@mmcv.jit(coderize=True)
 def accuracy(pred, target, topk=1, thresh=None):
     """Calculate accuracy according to the prediction and target.
 

mmdet/models/losses/ae_loss.py

@@ -1,10 +1,12 @@
+import mmcv
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 from ..builder import LOSSES
 
 
+@mmcv.jit(derivate=True, coderize=True)
 def ae_loss_per_image(tl_preds, br_preds, match):
     """Associative Embedding Loss in one image.
 

mmdet/models/losses/balanced_l1_loss.py

@@ -1,3 +1,4 @@
+import mmcv
 import numpy as np
 import torch
 import torch.nn as nn
@@ -6,6 +7,7 @@
 from .utils import weighted_loss
 
 
+@mmcv.jit(derivate=True, coderize=True)
 @weighted_loss
 def balanced_l1_loss(pred,
                      target,

mmdet/models/losses/gaussian_focal_loss.py

@@ -1,9 +1,11 @@
+import mmcv
 import torch.nn as nn
 
 from ..builder import LOSSES
 from .utils import weighted_loss
 
 
+@mmcv.jit(derivate=True, coderize=True)
 @weighted_loss
 def gaussian_focal_loss(pred, gaussian_target, alpha=2.0, gamma=4.0):
     """`Focal Loss <https://arxiv.org/abs/1708.02002>`_ for targets in gaussian

mmdet/models/losses/gfocal_loss.py

@@ -1,10 +1,12 @@
+import mmcv
 import torch.nn as nn
 import torch.nn.functional as F
 
 from ..builder import LOSSES
 from .utils import weighted_loss
 
 
+@mmcv.jit(derivate=True, coderize=True)
 @weighted_loss
 def quality_focal_loss(pred, target, beta=2.0):
     r"""Quality Focal Loss (QFL) is from `Generalized Focal Loss: Learning
@@ -49,6 +51,7 @@ def quality_focal_loss(pred, target, beta=2.0):
     return loss
 
 
+@mmcv.jit(derivate=True, coderize=True)
 @weighted_loss
 def distribution_focal_loss(pred, label):
     r"""Distribution Focal Loss (DFL) is from `Generalized Focal Loss: Learning