faster_rcnn_pytorch中的roi_pooling源码解析

源码来自于https://github.com/longcw/faster_rcnn_pytorch
这次只看了roi_pooling文件夹下的代码，其他的需要编译的文件跟这个类似。

roi_pooling目录

roi_pooling目录

-src文件夹下是c和cuda版本的源码，其中roi_pooling的操作的foward是c和cuda版本都有的，而backward仅写了cuda版本的代码。
-functions文件夹下的roi_pool.py是继承了torch.autograd.Function类，实现RoI层的foward和backward函数。
-modules文件夹下的roi_pool.py是继承了torch.nn.Modules类，实现了对RoI层的封装，此时RoI层就跟ReLU层一样的使用了。
-_ext文件夹下还有个roi_pooling文件夹，这个文件夹是存储src中c，cuda编译过后的文件的，编译过后就可以被funcitons中的roi_pool.py调用了。

functions/roi_pool.py

# -*- coding:utf8 -*-
import torch
from torch.autograd import Function
from .._ext import roi_pooling
import pdb

# 重写函数实现RoI层的正向传播和反向传播 modules中的roi_pool实现层的封装

class RoIPoolFunction(Function):
    def __init__(ctx, pooled_height, pooled_width, spatial_scale):
        ctx.pooled_width = pooled_width
        ctx.pooled_height = pooled_height
        ctx.spatial_scale = spatial_scale
        ctx.feature_size = None

    def forward(ctx, features, rois): 
        ctx.feature_size = features.size()           
        batch_size, num_channels, data_height, data_width = ctx.feature_size
        num_rois = rois.size(0)
        output = features.new(num_rois, num_channels, ctx.pooled_height, ctx.pooled_width).zero_() #new是torch.tensor的方法
        ctx.argmax = features.new(num_rois, num_channels, ctx.pooled_height, ctx.pooled_width).zero_().int()
        ctx.rois = rois
        if not features.is_cuda:
            _features = features.permute(0, 2, 3, 1)  # permute = transform 也是torch.tensor的方法
            roi_pooling.roi_pooling_forward(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale,
                                            _features, rois, output) #调用_ext下的编译好的cpu版本函数
        else:
            roi_pooling.roi_pooling_forward_cuda(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale,
                                                 features, rois, output, ctx.argmax) #调用_ext下的编译好的gpu版本函数

        return output

    def backward(ctx, grad_output):
        assert(ctx.feature_size is not None and grad_output.is_cuda)
        batch_size, num_channels, data_height, data_width = ctx.feature_size
        grad_input = grad_output.new(batch_size, num_channels, data_height, data_width).zero_()

        roi_pooling.roi_pooling_backward_cuda(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale,
                                              grad_output, ctx.rois, grad_input, ctx.argmax) #这个地方就只有gpu版本的了

        return grad_input, None

modules\roi_pool.py

# -*- coding:utf8 -*-

from torch.nn.modules.module import Module
from ..functions.roi_pool import RoIPoolFunction

# 对roi_pooling层的封装，就是ROI Pooling Layer了
class _RoIPooling(Module):
    def __init__(self, pooled_height, pooled_width, spatial_scale):
        super(_RoIPooling, self).__init__()

        self.pooled_width = int(pooled_width)
        self.pooled_height = int(pooled_height)
        self.spatial_scale = float(spatial_scale)

    def forward(self, features, rois):
        return RoIPoolFunction(self.pooled_height, self.pooled_width, self.spatial_scale)(features, rois)
               # 直接调用了functions中的函数，此时已经实现了foward，backward操作

剩下的src，_ext文件的代码就可以自己读读了，就是用c，cuda对roi_pooling实现了foward和backward，目的就是为了让python可以调用。