前言
最近在忙于学业的同时，终于将之前一直心心念念的rotate roi align的源码通读了一遍，通读完毕之后感觉真的收获颇丰。由于想限制篇幅，不给读者造成太多的阅读压力，所以本篇博客只介绍前向过程的源码，下篇博客再介绍反向过程的源码。
(ps:本次阅读的源码主要参考自detectron2框架，mmdetection框架的源码也是基于此进行了些许的优化，etc：加入了旋转方向的判断)。
正文
通常提出的算子在pytorch本身不支持自动求导的情况下，都会利用torch.autograd.Function进行算子的重新定义，并且考虑到运行效率通常还会利用到torch.utils.cpp_extension结合c++和cuda编程提升在实际训练过程中的速度。
本次分析的源码也不例外，并且由于c++源码和cuda源码本质上都是一样的，只不过cuda源码多了一些额外的核函数定义等。所以本次分析的是c++源码。
First
at::Tensor ROIAlignRotated_forward_cpu(
const at::Tensor& input,
const at::Tensor& rois,
const float spatial_scale,
const int pooled_height,
const int pooled_width,
const int sampling_ratio) {
// 确保变量符合要求
AT_ASSERTM(input.device().is_cpu(), "input must be a CPU tensor");
AT_ASSERTM(rois.device().is_cpu(), "rois must be a CPU tensor");
at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2};
at::CheckedFrom c = "ROIAlign_forward_cpu";
at::checkAllSameType(c, {input_t, rois_t});
// 得到相关尺寸信息
auto num_rois = rois.size(0);
auto channels = input.size(1);
auto height = input.size(2);
auto width = input.size(3);
// 定义输出，初始值为0
at::Tensor output = at::zeros(
{num_rois, channels, pooled_height, pooled_width}, input.options());
// 得到输出的尺寸大小
auto output_size = num_rois * pooled_height * pooled_width * channels;
if (output.numel() == 0) {
return output;
}
// 确保变量的内存是连续的
auto input_ = input.contiguous(), rois_ = rois.contiguous();
// 这个宏定义本身其实是一个匿名函数，使得函数能够支持半精度计算
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
input.scalar_type(), "ROIAlignRotated_forward", [&] {
ROIAlignRotatedForward<scalar_t>(
output_size,
input_.data_ptr<scalar_t>(),
spatial_scale,
channels,
height,
width,
pooled_height,
pooled_width,
sampling_ratio,
rois_.data_ptr<scalar_t>(),
output.data_ptr<scalar_t>());
});
return output;
}
上述函数是从python程序传出的参数需要经历的第一个c++程序，除了程序中所标注的注释，有两个部分值得注意：
1、传入参数的含义：
const at::Tensor& input // 传入的张量
const at::Tensor& rois // 传入的rpn生成的候选区域(batch_ind,x,y,w,h,theta)
const float spatial_scale // 每层特征图相对于原图缩放的尺度
const int pooled_height // 池化后的特征图尺寸高度，通常目标检测为7
const int pooled_width // 池化后的特征图尺寸宽度
const int sampling_ratio // 采样率，目前经常为0，后续会结合程序说明
2、contiguous()
这个函数对于张量的作用简单概括为，使得张量在内存中的存储变为连续的，方便后续使用指针取到相应的值。(ps：张量在内存中的表现形式一般是以行优先的一维展开形式)，详情可以参考：这篇知乎
second
接下来就来到了具体前向流程：
template <typename T>
void ROIAlignRotatedForward(
const int nthreads, const T* input, const T& spatial_scale, const int channels,
const int height, const int width, const int pooled_height, const int pooled_width,
const int sampling_ratio, const T* rois, T* output) {
// 得到候选区域的个数
int n_rois = nthreads / channels / pooled_width / pooled_height;
// 循环，对单个roi进行处理
for (int n = 0; n < n_rois; n++) {
// 由于在底层存储中，多维张量只是一排按特定顺序存储的数，所以为了得到相应的位置
// 只要将指针位置相加
int index_n = n * channels * pooled_width * pooled_height;
// 得到候选区域的参数，并进行简单的处理
const T* current_roi = rois + n * 6;
int roi_batch_ind = current_roi[0];
T offset = (T)0.5;
T roi_center_w = current_roi[1] * spatial_scale - offset;
T roi_center_h = current_roi[2] * spatial_scale - offset;
T roi_width = current_roi[3] * spatial_scale;
T roi_height = current_roi[4] * spatial_scale;
T theta = current_roi[5];
T cos_theta = cos(theta);
T sin_theta = sin(theta);
// 确保候选区域的宽和高都大于0
AT_ASSERTM(
roi_width >= 0 && roi_height >= 0,
"ROIs in ROIAlignRotated do not have non-negative size!");
// 得到一个bin内会有多少个参数
T bin_size_h = static_cast<T>(roi_height) / static_cast<T>(pooled_height);
T bin_size_w = static_cast<T>(roi_width) / static_cast<T>(pooled_width);
// 这里就需要用到采样率参数，来判断每个bin内需要有多少个采样点了，如果为0的话，就意味着需要采集和bin内像素点个数相同的采样点个数
int roi_bin_grid_h =
(sampling_ratio > 0) ? sampling_ratio : ceil(roi_height / pooled_height); // e.g., = 2
int roi_bin_grid_w =
(sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / pooled_width);
// 得到采样点的个数
const T count = std::max(roi_bin_grid_h * roi_bin_grid_w, 1); // e.g. = 4
// 定义了一个顺序容器，该容器内的每一个元素都是一个结构体，存储了四个点的位置坐标（指针的偏移量）以及四个面积
std::vector<PreCalc<T>> pre_calc(
roi_bin_grid_h * roi_bin_grid_w * pooled_width * pooled_height);
// 旋转之前的偏移量，之后的双线性插值法可以用到
T roi_start_h = -roi_height / 2.0;
T roi_start_w = -roi_width / 2.0;
// 双线性插值法进行处理
pre_calc_for_bilinear_interpolate(
height, width, pooled_height, pooled_width, roi_bin_grid_h, roi_bin_grid_w,roi_start_h,
roi_start_w, bin_size_h, bin_size_w, roi_bin_grid_h, roi_bin_grid_w, roi_center_h,
roi_center_w, cos_theta, sin_theta, pre_calc);
// 对特征图的每一个通道进行处理
for (int c = 0; c < channels; c++) {
// 定位到输入的特定特征图，一个是输出的索引一个是输入的指针
int index_n_c = index_n + c * pooled_width * pooled_height;
const T* offset_input =
input + (roi_batch_ind * channels + c) * height * width;
int pre_calc_index = 0;
for (int ph = 0; ph < pooled_height; ph++) {
for (int pw = 0; pw < pooled_width; pw++) {
// 定位到输出的单个像素位置
int index = index_n_c + ph * pooled_width + pw;
T output_val = 0.;
// 利用双线性插值法的计算公式，得到每一个采样点的具体的值求和
for (int iy = 0; iy < roi_bin_grid_h; iy++) {
for (int ix = 0; ix < roi_bin_grid_w; ix++) {
PreCalc<T> pc = pre_calc[pre_calc_index];
output_val += pc.w1 * offset_input[pc.pos1] +
pc.w2 * offset_input[pc.pos2] +
pc.w3 * offset_input[pc.pos3] + pc.w4 * offset_input[pc.pos4];
pre_calc_index += 1;
}
}
// 均值处理
output_val /= count;
// 最终得到输出像素的值
output[index] = output_val;
} // for pw
} // for ph
} // for c
} // for n
}
上述程序值得注意的点在于
1、双线性插值的子函数
// 结构体模板，用于存储每个采样点的信息
template <typename T>
struct PreCalc {
int pos1;
int pos2;
int pos3;
int pos4;
T w1;
T w2;
T w3;
T w4;
};
// 双线性插值采样
template <typename T>
void pre_calc_for_bilinear_interpolate(
const int height, const int width, const int pooled_height, const int pooled_width, const int iy_upper,
const int ix_upper, T roi_start_h, T roi_start_w, T bin_size_h, T bin_size_w, int roi_bin_grid_h,
int roi_bin_grid_w, T roi_center_h, T roi_center_w, T cos_theta, T sin_theta,
std::vector<PreCalc<T>>& pre_calc) {
int pre_calc_index = 0;
for (int ph = 0; ph < pooled_height; ph++) {
// 先取输出宽度bin坐标
for (int pw = 0; pw < pooled_width; pw++) {
// 再取输出高度bin坐标
for (int iy = 0; iy < iy_upper; iy++) {
// 这层循环指的是高度上需要采样的点的个数
// 算出需要采样点的纵坐标
const T yy = roi_start_h + ph * bin_size_h +
static_cast<T>(iy + .5f) * bin_size_h /
static_cast<T>(roi_bin_grid_h); // e.g., 0.5, 1.5
for (int ix = 0; ix < ix_upper; ix++) {
// 这层循环指的是宽度上需要采样的点的个数
// 算出需要采样点的横坐标
const T xx = roi_start_w + pw * bin_size_w +
static_cast<T>(ix + .5f) * bin_size_w /
static_cast<T>(roi_bin_grid_w);
// 上述坐标之所以加上start是为了方便之后旋转是围绕原点旋转的
// 这里的旋转矩阵按照常规的算法会发现，是顺时针矩阵。因为图片的坐标系是左手系，所以在右手系中坐标顺序通常是(y,x)
T y = yy * cos_theta - xx * sin_theta + roi_center_h;
T x = yy * sin_theta + xx * cos_theta + roi_center_w;
// 判断是否超出边界
if (y < -1.0 || y > height || x < -1.0 || x > width) {
// empty
PreCalc<T> pc;
pc.pos1 = 0;
pc.pos2 = 0;
pc.pos3 = 0;
pc.pos4 = 0;
pc.w1 = 0;
pc.w2 = 0;
pc.w3 = 0;
pc.w4 = 0;
pre_calc[pre_calc_index] = pc;
pre_calc_index += 1;
continue;
}
// 考虑特殊情况
if (y < 0) {
y = 0;
}
if (x < 0) {
x = 0;
}
// 定义采样点附近四个点变量
int y_low = (int)y;
int x_low = (int)x;
int y_high;
int x_high;
if (y_low >= height - 1) {
y_high = y_low = height - 1;
y = (T)y_low;
} else {
y_high = y_low + 1;
}
if (x_low >= width - 1) {
x_high = x_low = width - 1;
x = (T)x_low;
} else {
x_high = x_low + 1;
}
// 计算采样点分割的四个小矩形的面积
T ly = y - y_low;
T lx = x - x_low;
T hy = 1. - ly, hx = 1. - lx;
T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
// 将四个点的位置偏移量和四个小矩形存储
PreCalc<T> pc;
pc.pos1 = y_low * width + x_low;
pc.pos2 = y_low * width + x_high;
pc.pos3 = y_high * width + x_low;
pc.pos4 = y_high * width + x_high;
pc.w1 = w1;
pc.w2 = w2;
pc.w3 = w3;
pc.w4 = w4;
pre_calc[pre_calc_index] = pc;
pre_calc_index += 1;
}
}
}
}
}
上述程序考察的唯一点就是双线性插值法。
上图中，右上，左上，右下，左下的面积分别为：
w
1
,
w
2
,
w
3
,
w
4
w_1,w_2,w_3,w_4
w1?,w2?,w3?,w4?
f
(
x
,
y
)
=
w
1
?
Q
11
+
w
2
?
Q
21
+
w
3
?
Q
12
+
w
4
?
Q
22
f(x,y)=w_1*Q_{11}+w_2*Q_{21}+w_3*Q_{12}+w_4*Q{22}
f(x,y)=w1??Q11?+w2??Q21?+w3??Q12?+w4??Q22
最终将采样得到的点做均值处理得到最终的输出值。
结束
以上便是我对该模块前向的一点浅显的解读，如有错误，欢迎各位指正，接下来我会继续更新反向过程！！！