add lib (zhanghang1989#61)

.gitignore

@@ -6,7 +6,5 @@ build/
 data/
 docs/src/
 docs/html/
-encoding/lib/
 encoding/_ext/
 encoding.egg-info/
-experiments/segmentation/

encoding/lib/__init__.py

@@ -0,0 +1,20 @@
+import os
+import torch
+from torch.utils.cpp_extension import load
+
+cwd = os.path.dirname(os.path.realpath(__file__))
+cpu_path = os.path.join(cwd, 'cpu')
+gpu_path = os.path.join(cwd, 'gpu')
+
+cpu = load( 'enclib_cpu', [
+        os.path.join(cpu_path, 'roi_align.cpp'),
+        os.path.join(cpu_path, 'roi_align_cpu.cpp'),
+    ], build_directory=cpu_path, verbose=False)
+
+if torch.cuda.is_available():
+    gpu = load( 'enclib_gpu', [
+            os.path.join(gpu_path, 'operator.cpp'),
+            os.path.join(gpu_path, 'encoding_kernel.cu'),
+            os.path.join(gpu_path, 'syncbn_kernel.cu'),
+            os.path.join(gpu_path, 'roi_align_kernel.cu'),
+        ], build_directory=gpu_path, verbose=False)

encoding/lib/cpu/roi_align.cpp

@@ -0,0 +1,28 @@
+#include <torch/torch.h>
+// CPU declarations
+
+at::Tensor ROIAlignForwardCPU(
+  const at::Tensor& input,
+  const at::Tensor& bottom_rois,
+  int64_t pooled_height,
+  int64_t pooled_width,
+  double spatial_scale,
+  int64_t sampling_ratio);
+
+at::Tensor ROIAlignBackwardCPU(
+  const at::Tensor& bottom_rois,
+  const at::Tensor& grad_output, // gradient of the output of the layer
+  int64_t b_size,
+  int64_t channels,
+  int64_t height,
+  int64_t width,
+  int64_t pooled_height,
+  int64_t pooled_width,
+  double spatial_scale,
+  int64_t sampling_ratio);
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("roi_align_forward", &ROIAlignForwardCPU, "ROI Align forward (CPU)");
+  m.def("roi_align_backward", &ROIAlignBackwardCPU, "ROI Align backward (CPU)");
+}