add onnx namespace for custom ops (#1254)

open-mmlab · Aug 16, 2021 · 44edcdd · 44edcdd
1 parent 54907a3
commit 44edcdd
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 35 deletions.
diff --git a/mmcv/ops/carafe.py b/mmcv/ops/carafe.py
@@ -19,12 +19,12 @@ class CARAFENaiveFunction(Function):
     @staticmethod
     def symbolic(g, features, masks, kernel_size, group_size, scale_factor):
         return g.op(
-            'MMCVCARAFENaive',
+            'mmcv::MMCVCARAFENaive',
             features,
             masks,
-            kernel_size=kernel_size,
-            group_size=group_size,
-            scale_factor=scale_factor)
+            kernel_size_i=kernel_size,
+            group_size_i=group_size,
+            scale_factor_f=scale_factor)
 
     @staticmethod
     def forward(ctx, features, masks, kernel_size, group_size, scale_factor):
@@ -102,12 +102,12 @@ class CARAFEFunction(Function):
     @staticmethod
     def symbolic(g, features, masks, kernel_size, group_size, scale_factor):
         return g.op(
-            'MMCVCARAFE',
+            'mmcv::MMCVCARAFE',
             features,
             masks,
-            kernel_size=kernel_size,
-            group_size=group_size,
-            scale_factor=scale_factor)
+            kernel_size_i=kernel_size,
+            group_size_i=group_size,
+            scale_factor_f=scale_factor)
 
     @staticmethod
     def forward(ctx, features, masks, kernel_size, group_size, scale_factor):

diff --git a/mmcv/ops/cc_attention.py b/mmcv/ops/cc_attention.py
@@ -15,7 +15,7 @@ class CAWeightFunction(torch.autograd.Function):
 
     @staticmethod
     def symbolic(g, t, f):
-        return g.op('MMCVCAWeight', t, f)
+        return g.op('mmcv::MMCVCAWeight', t, f)
 
     @staticmethod
     def forward(ctx, t, f):
@@ -41,7 +41,7 @@ class CAMapFunction(torch.autograd.Function):
 
     @staticmethod
     def symbolic(g, weight, v):
-        return g.op('MMCVCAMap', weight, v)
+        return g.op('mmcv::MMCVCAMap', weight, v)
 
     @staticmethod
     def forward(ctx, weight, v):

diff --git a/mmcv/ops/deform_roi_pool.py b/mmcv/ops/deform_roi_pool.py
@@ -16,15 +16,15 @@ class DeformRoIPoolFunction(Function):
     def symbolic(g, input, rois, offset, output_size, spatial_scale,
                  sampling_ratio, gamma):
         return g.op(
-            'MMCVDeformRoIPool',
+            'mmcv::MMCVDeformRoIPool',
             input,
             rois,
             offset,
-            pooled_height=output_size[0],
-            pooled_width=output_size[1],
-            spatial_scale=spatial_scale,
-            sampling_ratio=sampling_ratio,
-            gamma=gamma)
+            pooled_height_i=output_size[0],
+            pooled_width_i=output_size[1],
+            spatial_scale_f=spatial_scale,
+            sampling_ratio_f=sampling_ratio,
+            gamma_f=gamma)
 
     @staticmethod
     def forward(ctx,

diff --git a/mmcv/ops/focal_loss.py b/mmcv/ops/focal_loss.py
@@ -17,13 +17,13 @@ class SigmoidFocalLossFunction(Function):
     @staticmethod
     def symbolic(g, input, target, gamma, alpha, weight, reduction):
         return g.op(
-            'MMCVSigmoidFocalLoss',
+            'mmcv::MMCVSigmoidFocalLoss',
             input,
             target,
-            gamma=gamma,
-            alpha=alpha,
-            weight=weight,
-            reduction=reduction)
+            gamma_f=gamma,
+            alpha_f=alpha,
+            weight_f=weight,
+            reduction_s=reduction)
 
     @staticmethod
     def forward(ctx,
@@ -111,13 +111,13 @@ class SoftmaxFocalLossFunction(Function):
     @staticmethod
     def symbolic(g, input, target, gamma, alpha, weight, reduction):
         return g.op(
-            'MMCVSoftmaxFocalLoss',
+            'mmcv::MMCVSoftmaxFocalLoss',
             input,
             target,
-            gamma=gamma,
-            alpha=alpha,
-            weight=weight,
-            reduction=reduction)
+            gamma_f=gamma,
+            alpha_f=alpha,
+            weight_f=weight,
+            reduction_s=reduction)
 
     @staticmethod
     def forward(ctx,

diff --git a/mmcv/ops/masked_conv.py b/mmcv/ops/masked_conv.py
@@ -18,13 +18,13 @@ class MaskedConv2dFunction(Function):
     @staticmethod
     def symbolic(g, features, mask, weight, bias, padding, stride):
         return g.op(
-            'MMCVMaskedConv2d',
+            'mmcv::MMCVMaskedConv2d',
             features,
             mask,
             weight,
             bias,
-            padding=padding,
-            stride=stride)
+            padding_i=padding,
+            stride_i=stride)
 
     @staticmethod
     def forward(ctx, features, mask, weight, bias, padding=0, stride=1):

diff --git a/mmcv/ops/psa_mask.py b/mmcv/ops/psa_mask.py
@@ -14,7 +14,10 @@ class PSAMaskFunction(Function):
     @staticmethod
     def symbolic(g, input, psa_type, mask_size):
         return g.op(
-            'MMCVPSAMask', input, psa_type=psa_type, mask_size=mask_size)
+            'mmcv::MMCVPSAMask',
+            input,
+            psa_type_i=psa_type,
+            mask_size_i=mask_size)
 
     @staticmethod
     def forward(ctx, input, psa_type, mask_size):

diff --git a/mmcv/ops/sync_bn.py b/mmcv/ops/sync_bn.py
@@ -22,16 +22,16 @@ class SyncBatchNormFunction(Function):
     def symbolic(g, input, running_mean, running_var, weight, bias, momentum,
                  eps, group, group_size):
         return g.op(
-            'MMCVSyncBatchNorm',
+            'mmcv::MMCVSyncBatchNorm',
             input,
             running_mean,
             running_var,
             weight,
             bias,
-            momentum=momentum,
-            eps=eps,
-            group=group,
-            group_size=group_size)
+            momentum_f=momentum,
+            eps_f=eps,
+            group_i=group,
+            group_size_i=group_size)
 
     @staticmethod
     def forward(self, input, running_mean, running_var, weight, bias, momentum,