學習Faster R-CNN代碼roi_pooling（二）

roi_pooling理解起來比較簡單，所以我就先看了一下這部分的代碼。

roi_pooling目錄下

-src文件夾下是c和cuda版本的源碼。

-functions文件夾下的roi_pool.py是繼承了torch.autograd.Function類，實現RoI層的foward和backward函數。class RoIPoolFunction(Function)。

-modules文件夾下的roi_pool.py是繼承了torch.nn.Modules類，實現了對RoI層的封裝，此時RoI層就跟ReLU層一樣的使用了。class _RoIPooling(Module)。

-_ext文件夾下還有個roi_pooling文件夾，這個文件夾是存儲src中c，cuda編譯過后的文件的，編譯過后就可以被funcitons中的roi_pool.py調用了。

具體代碼：
如下圖所示：為roi_pooling/src/roi_pooling.c文件中的一段，可以根據該代碼了解到各變量包含哪些內容。如果只看py代碼，有些地方真不知道包含哪些信息，看起來很費勁，仔細看了下C代碼，解決了自己的疑惑。

rois[0] num_rois即roi的個數；
rois[1] size_rois即大小；
features[0] batch_size，批大小；
features[1] data_height，高；
features[2] data_width，寬；
features[3] num_channels，通道數。

細致一點的理解看我的另一篇博客：https://blog.csdn.net/weixin_43872578/article/details/86628515

functions/roi_pool.py

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 is a context object that can be used to stash information for backward computation
        #上下文對象，可用於存儲信息以進行反向計算
        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_()
        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)
            roi_pooling.roi_pooling_forward(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale,
                                            _features, rois, output)
        else:
            roi_pooling.roi_pooling_forward_cuda(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale,
                                                 features, rois, output, ctx.argmax)

        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)

        return grad_input, None

modules/roi_pool.py

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可以調用。

【未完，待更新…】

ref：https://www.jianshu.com/p/d674e16ce896

https://blog.csdn.net/weixin_43872578/article/details/86616801