[源码分析]rotate_roi_align前向过程

最新推荐文章于 2024-07-12 16:16:27 发布

K . U . I

最新推荐文章于 2024-07-12 16:16:27 发布

阅读量1k

点赞数 1

分类专栏： mmdetection 文章标签： python

本文链接：https://blog.csdn.net/qq_44109682/article/details/120381827

版权

mmdetection 专栏收录该内容

3 篇文章 1 订阅

订阅专栏

前言

最近在忙于学业的同时，终于将之前一直心心念念的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$
$f(x,y)=w_1*Q_{11}+w_2*Q_{21}+w_3*Q_{12}+w_4*Q{22}$
最终将采样得到的点做均值处理得到最终的输出值。

结束

以上便是我对该模块前向的一点浅显的解读，如有错误，欢迎各位指正，接下来我会继续更新反向过程！！！

K . U . I

关注

1
点赞
踩
3

收藏

觉得还不错? 一键收藏
0
评论
[源码分析]rotate_roi_align前向过程

前言最近在忙于学业的同时，终于将之前一直心心念念的rotate roi align的源码通读了一遍，通读完毕之后感觉真的收获颇丰。由于想限制篇幅，不给读者造成太多的阅读压力，所以本篇博客只介绍前向过程的源码，下篇博客再介绍反向过程的源码。(ps:本次阅读的源码主要参考自detectron2框架，mmdetection框架的源码也是基于此进行了些许的优化，etc：加入了旋转方向的判断)。正文通常提出的算子在pytorch本身不支持自动求导的情况下，都会利用torch.autograd.Function
复制链接

扫一扫