add external OPs from MaskRCNN

Author: jiayisunx

intel_pytorch_extension_py/ops/__init__.py

@@ -6,3 +6,6 @@
 from .jit import *
 from .save import *
 from .to import *
+from .roi_align import ROIAlign
+from .roi_align import roi_align
+from .nms import nms

intel_pytorch_extension_py/ops/nms.py

@@ -0,0 +1,3 @@
+import _torch_ipex as core
+
+nms = core.nms

intel_pytorch_extension_py/ops/roi_align.py

@@ -0,0 +1,68 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+import torch
+from torch import nn
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+from torch.nn.modules.utils import _pair
+
+import _torch_ipex as core
+
+
+class _ROIAlign(Function):
+    @staticmethod
+    def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio):
+        ctx.save_for_backward(roi)
+        ctx.output_size = _pair(output_size)
+        ctx.spatial_scale = spatial_scale
+        ctx.sampling_ratio = sampling_ratio
+        ctx.input_shape = input.size()
+        output = core.roi_align_forward(
+            input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio
+        )
+        return output
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        rois, = ctx.saved_tensors
+        output_size = ctx.output_size
+        spatial_scale = ctx.spatial_scale
+        sampling_ratio = ctx.sampling_ratio
+        bs, ch, h, w = ctx.input_shape
+        grad_input = core.roi_align_backward(
+            grad_output,
+            rois,
+            spatial_scale,
+            output_size[0],
+            output_size[1],
+            bs,
+            ch,
+            h,
+            w,
+            sampling_ratio,
+        )
+        return grad_input, None, None, None, None
+
+
+roi_align = _ROIAlign.apply
+
+
+class ROIAlign(nn.Module):
+    def __init__(self, output_size, spatial_scale, sampling_ratio):
+        super(ROIAlign, self).__init__()
+        self.output_size = output_size
+        self.spatial_scale = spatial_scale
+        self.sampling_ratio = sampling_ratio
+
+    def forward(self, input, rois):
+        return roi_align(
+            input, rois, self.output_size, self.spatial_scale, self.sampling_ratio
+        )
+
+    def __repr__(self):
+        tmpstr = self.__class__.__name__ + "("
+        tmpstr += "output_size=" + str(self.output_size)
+        tmpstr += ", spatial_scale=" + str(self.spatial_scale)
+        tmpstr += ", sampling_ratio=" + str(self.sampling_ratio)
+        tmpstr += ")"
+        return tmpstr

torch_ipex/csrc/cpu/ExternalOPs.h

@@ -0,0 +1,33 @@
+// Autogenerated file by gen-common-ops.py. Do not edit directly!
+#pragma once
+
+#include <ATen/Tensor.h>
+
+namespace torch_ipex {
+
+class IpexExternal {
+ public:
+  static at::Tensor ROIAlign_forward(const at::Tensor& input,
+                                     const at::Tensor& rois,
+                                     const float spatial_scale,
+                                     const int pooled_height,
+                                     const int pooled_width,
+                                     const int sampling_ratio);
+
+  static at::Tensor ROIAlign_backward(const at::Tensor& grad,
+                                      const at::Tensor& rois,
+                                      const float spatial_scale,
+                                      const int pooled_height,
+                                      const int pooled_width,
+                                      const int batch_size,
+                                      const int channels,
+                                      const int height,
+                                      const int width,
+                                      const int sampling_ratio);
+
+  static at::Tensor nms(const at::Tensor& dets,
+                        const at::Tensor& scores,
+                        const float threshold);
+};
+
+}  // namespace torch_ipex