(Oriented) 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:
p_module = tf.load_op_library(so_lib_path)
s_roi_align = op_module.ps_roi_align
ops.RegisterGradient("PsRoiAlign")
ef _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]
ool_method = 'max' # or 'mean'
ool_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:
Future Work:
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