HiKapok / PSROIAlign

(Oriented) PsRoIAlign Operation In Tensorflow C++ API

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

PsRoIAlign Operation In Tensorflow C++ API

PsRoIAlign involves interpolation techniques for PsRoiPooling (position-sensitive RoI pooling operation), the interpolation idea is proposed in RoIAlign to avoid any quantization of the RoI boundaries. The first adoption of PsRoIAlign might be in this paper Light-Head R-CNN: In Defense of Two-Stage Object Detector.

You can find more details about the RoiPooling technique in Fast R-CNN and Spatial Pyramid Pooling in Deep Convolutional Networks for Visual Recognition.

This repository contains code of the implement of PsRoIAlign operation in Tensorflow C++ API. You can use this operation in many popular two-stage object detector. Both research work using PsRoIAlign and contribution to this repository are welcomed.

For using this op in your own machine, just following these steps:

  • copy the header file "cuda_config.h" from "your_python_path/site-packages/external/local_config_cuda/cuda/cuda/cuda_config.h" to "your_python_path/site-packages/tensorflow/include/tensorflow/stream_executor/cuda/cuda_config.h".

  • run the following script:

mkdir build
cd build && cmake ..
make
  • run "test_op.py" and check the numeric errors to test your install

  • follow the below codes snippet to integrate this Op into your own code:

     op_module = tf.load_op_library(so_lib_path)
     ps_roi_align = op_module.ps_roi_align
    
     @ops.RegisterGradient("PsRoiAlign")
     def _ps_roi_align_grad(op, grad, _):
       '''The gradients for `PsRoiAlign`.
       '''
       inputs_features = op.inputs[0]
       rois = op.inputs[1]
       pooled_features_grad = op.outputs[0]
       pooled_index = op.outputs[1]
       grid_dim_width = op.get_attr('grid_dim_width')
       grid_dim_height = op.get_attr('grid_dim_height')
     
       return [op_module.ps_roi_align_grad(inputs_features, rois, grad, pooled_index, grid_dim_width, grid_dim_height, pool_method), None]
    
     pool_method = 'max' # or 'mean'
     pool_result = ps_roi_align(features, rois, 2, 2, pool_method)

The code is tested under TensorFlow 1.6 with CUDA 8.0 using Ubuntu 16.04. This PsRoIAlign Op had been used to train Xception based Light-Head RCNN successfully with performance at ~75%mAP on PASCAL VOC 2007 Test dataset, you can see codes here.

Update:

  • Added support for mean pooling (default is max pooling)
  • PsRoIAlign now support oriented RoI inputs (for both max and mean pooling).

Future Work:

  • Check if there is need to ensure the convex of polygon
  • Improve performance

If you encountered some linkage problem when generating or loading *.so, you are highly recommended to read this section in the official tourial to make sure you were using the same C++ ABI version.

The MIT License

About

(Oriented) PsRoIAlign Operation In Tensorflow C++ API

License:MIT License


Languages

Language:C++ 63.4%Language:Cuda 26.3%Language:Python 7.3%Language:CMake 2.3%Language:C 0.6%Language:Shell 0.1%