Faster RCNN目标检测部署加速

Posted on | | Visitors: | 本文总阅读量 次

前言

先说一下，这个博客不涉及TensorRt加速，不涉及半精度或者int8精度加速，仅仅是libtorch部署时，NMS和RoI Pool的加速问题。当然，最后实操下来目前没有做到比python下调用torchvision的api做到的总体结果快，算是一次失败的双阶段网络搭建尝试，后续有时间再更新找找原因吧。

关注我的朋友可能都知道我最近在写libtorch教程，有一段时间没有更新最后的关键章节——目标检测章了。我自己是写过写单阶段的目标检测器的C++实现的，没有RoI Pooling或者RoI Align，没有高性能要求的NMS，写起来用起来都很爽。

但是考虑到Faster RCNN毕竟是目标检测最最经典的方法（可能之一）了，还是硬着头皮整一下libtorch版本的Faster RCNN。过程果然如我所料，比单阶段的网络整起来麻烦很多，最后效果也不一定有（大概率没有）单阶段的目标检测网络好用。毕竟实际项目，速度为王，看YOLO系列多火就知道了。后面有机会整个YOLO全家桶吧….

RoI Pool

RoI Pool的作用及原理

进入正题，今天主要有两个点需要重点说说，RoI Pool和NMS的加速问题。先说说RoI Pool，了解双阶段目标检测的，或者直白点，了解Faster RCNN的朋友应该都知道它或者用过它（不了解的可以去b站搜下，有bub开头的up主讲得挺不错的）。RoI Pool (or Region of Interest Pool)，它在Faster RCNN算法中扮演的角色相当于一个中间步骤吧，双阶段检测就靠它承上启下。

RoI Pool之前，输入图像经过一个骨干网络卷积，得到特征图，这个特征图再经过一个RPN（Region Proposal Network）网络得到粗定位的一堆预测结果。其实这里涉及到一个点，就是卷积做分类的能力要高于做回归。早期的研究者尝试直接用卷积的方式做单阶段目标检测，直接通过特征图判断位置和类别，但是最后效果都不好。直到YOLO的横空出世，用了分治的思想，把单阶段目标检测算法提了一个台阶。

这里的RPN网络就可以看作一个只有有无目标判断能力，定位效果相对不够的单阶段检测算法，通过卷积得到的特征图，配合预先设定好的先验框，得到位置的粗略信息。这些粗略信息可以看成是一堆建议框，框里就是目标，目标的类别未知。

RoI Pool的作用就是将这些建议框从特征图中裁剪出来，然后经过池化(Pool)操作统一到同一个尺寸，再将这些相同尺寸的特征图拼接起来输出。后续就是第二阶段的分类和精定位了。

几种python实现的RoI Pool

在了解了RoI Pool的原理之后，我在GitHub上找到了一些python的实现，帮助进一步了解。这个项目帮助我很多，里面有四种pytorch的RoI Pool实现，并且提供了速度对比。

简单说下该项目的四种实现方式，具体细节及代码就不展开了。四种方式：cffi实现，cupy实现，chainer实现以及pytorch实现。下面是四个方式的速度测试结果：

1
2
3
4
5
use_cuda: True, has_backward: True
method0: 0.0344547653198242, batch_size: 1, size: 50, num_rois: 300
method1: 0.1322056961059570, batch_size: 1, size: 50, num_rois: 300
method2: 0.1307379817962646, batch_size: 1, size: 50, num_rois: 300
method3: 0.2016681671142578, batch_size: 1, size: 50, num_rois: 300

可以发现，cffi实现的速度是最快的，甚至比项目里提供的cupy方式还要快，cupy是python调用cuda编程加速。下面我先在RTX 2070s下pytorch版本的RoI Pool速度如何。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import time
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import torch
from torch import nn, Tensor
def roi_pooling(input, rois, size=(7, 7), spatial_scale=1.0):
	assert rois.dim() == 2
	assert rois.size(1) == 5
	output = []
	rois = rois.data.float()
	num_rois = rois.size(0)
 
	rois[:, 1:].mul_(spatial_scale)
	rois = rois.long()
	for i in range(num_rois):
		roi = rois[i]
		im_idx = roi[0]
		im = input.narrow(0, im_idx, 1)[..., roi[2]:(roi[4] + 1), roi[1]:(roi[3] + 1)]
		output.append(F.adaptive_max_pool2d(im, size))
 
	output = torch.cat(output, 0)
	if has_backward:
		#        output.backward(output.data.clone())
		output.sum().backward()
	return output

def create_rois(config):
	rois = torch.rand((config[2], 5))
	rois[:, 0] = rois[:, 0] * config[0]
	rois[:, 1:] = rois[:, 1:] * config[1]
	for j in range(config[2]):
		max_, min_ = max(rois[j, 1], rois[j, 3]), min(rois[j, 1], rois[j, 3])
		rois[j, 1], rois[j, 3] = min_, max_
		max_, min_ = max(rois[j, 2], rois[j, 4]), min(rois[j, 2], rois[j, 4])
		rois[j, 2], rois[j, 4] = min_, max_
	rois = torch.floor(rois)
	rois = Variable(rois, requires_grad=False)
	return rois

if __name__ == '__main__':
	# batch_size, img_size, num_rois
	config = [1, 50, 300]
	T = 50
	has_backward = False
 
	start = time.time()
	x = Variable(torch.rand((config[0], 512, config[1], config[1])),requires_grad=True).cuda()
	rois = create_rois(config).cuda()
	for t in range(T):
		# output = roi_pool_torchvision(x, rois, 7, 1.0)
		output = roi_pooling(x,rois)
	print('time: {}, batch_size: {}, size: {}, num_rois: {}'.format((time.time() - start) / T,
		                                                                config[0],
		                                                                config[1],
		                                                                config[2]))

以上代码基本是从项目中复制来的，没有太多改变，输出打印如下：

1
time: 0.12641457080841065, batch_size: 1, size: 50, num_rois: 300

耗时126ms，这比项目输出的速度要好一些。但是这显然不能用于生产环境中，而且还没算训练时的RoI Pool耗时，但是应该比这高得多。

libtorch版本的RoI Pool

然后我不信邪，在C++里实现了上述的pytorch代码。代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class RoIPoolImpl : public torch::nn::Module {
public:
	RoIPoolImpl(int output_size, float spatial_scale);
	torch::Tensor forward(torch::Tensor input, torch::Tensor rois);
private:
	float spatial_scale = 1.0;
	int output_size;
}; TORCH_MODULE(RoIPool);

RoIPoolImpl::RoIPoolImpl(int _output_size, float _spatial_scale) {
	spatial_scale = _spatial_scale;
	output_size = _output_size;
}
torch::Tensor RoIPoolImpl::forward(torch::Tensor input, torch::Tensor rois) {
	assert(rois.dim() == 2);
	assert(rois.size(1) == 5);
	int num_rois = rois.size(0);
	rois.narrow(1, 1, rois.size(1)-1) = rois.narrow(1, 1, rois.size(1) - 1).mul_(spatial_scale);
	rois = rois.to(torch::kLong);

	rois = torch::clamp(rois, 0);
	rois.select(1, 1) = torch::clamp_max(rois.select(1, 1), input.size(3) -1);
	rois.select(1, 2) = torch::clamp_max(rois.select(1, 2), input.size(2) - 1);
	rois.select(1, 3) = torch::clamp_max(rois.select(1, 3), input.size(3) - 1);
	rois.select(1, 4) = torch::clamp_max(rois.select(1, 4), input.size(2) - 1);

	rois.select(1, 3) = rois.select(1, 3) - rois.select(1, 1);
	rois.select(1, 4) = rois.select(1, 4) - rois.select(1, 2);
	rois = torch::clamp(rois, 0);
	std::vector<torch::Tensor> output;
	auto poolOption = torch::nn::functional::AdaptiveMaxPool2dFuncOptions(output_size);
	for (int i = 0; i < num_rois; i++) {//此循环外可以将rois拷贝到数组，再循环内取数，应该提速
		auto roi = rois[i];
		auto im_idx = roi[0];
		int x1 = roi[1].item().toInt();
		int y1 = roi[2].item().toInt();
		int w = roi[3].item().toInt();
		int h = roi[4].item().toInt();

		auto im = input.narrow(0, im_idx, 1).narrow(2, y1, h + 1).narrow(3, x1, w + 1);
		output.push_back(torch::nn::functional::adaptive_max_pool2d(im, poolOption));
	}
	return torch::cat(output, 0);
}

测试了一下速度，测试代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
torch::Tensor create_rois(int config[]) {
	auto rois = torch::rand({ config[2], 5 });
	rois.select(1,0) = rois.select(1, 0) * config[0];
	rois.narrow(1, 1, rois.size(1) - 1) = rois.narrow(1, 1, rois.size(1) - 1) * config[1];
	for (int j = 0; j < config[2]; j++) {
		auto max_ = max(rois[j][1], rois[j][3]);
		auto min_ = min(rois[j][1], rois[j][3]);
		rois[j][1] = min_;
		rois[j][3] = max_;
		max_ = max(rois[j][2], rois[j][4]);
		min_ = min(rois[j][2], rois[j][4]);
		rois[j][2] = min_;
		rois[j][4] = max_;
	}
	rois = torch::floor(rois);
	return rois;
}
int main(int argc, char *argv[])
{

	int config[3] = { 1, 50, 300 };
	auto roipool = RoIPool(7, 1.0);
	roipool->eval();
	roipool->to(torch::Device(torch::kCUDA));
	auto x = torch::rand({ config[0], 512, config[1], config[1] }).to(torch::Device(torch::kCUDA));
	auto rois = create_rois(config).to(torch::Device(torch::kCUDA));
	roipool->forward(x, rois);
	//model->forward(inputs);
	int T = 10;
	int64 t0 = cv::getCPUTickCount();
	for (int i = 0; i < T; i++) {
		//model->forward(inputs);
		auto output = roipool->forward(x, rois);
	}
	int64 t1 = cv::getCPUTickCount();
	std::cout << "execution time is " << (t1 - t0) / (double)cv::getTickFrequency() / T << " seconds" << std::endl;
}

代码里用了libtorch和opencv，后者主要用于计时，如果不能运行可以参考网络的配置博客或者我自己的pytorch部署博客和libtorch的QT部署。

可以发现我的代码主要添加了一些安全性设置，原来的python代码，输入受create_rois影响，已经将rois内部的大小关系定义成方框的左上角点坐标小于右下角点。但是实际的Faster RCNN中，RPN网络出来的结果并非如此，所以加了安全性设置。

在同样的RTX 2070s显卡上，该libtorch代码的执行速度高达0.08s，比python下同样代码快个50%吧。除去安全性代码可能还会更快。但是，我们发现，这个速度最多只能和开源项目中的cupy或者chainer相对提速优越一点点。（0.2s->0.13s和0.126s->0.08s）。它的速度依然是比cffi实现慢许多的，如果在同样显卡上甚至慢相当多。

libtorch版本的RoI Pool提速

那我们对这个代码进行提速，看看最后能到什么效果。重新构建RoIPool类，代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
class RoIPoolImpl : public torch::nn::Module {//our RoIPool should be a little slower than cuda version, not tested
public:
	RoIPoolImpl(int output_size, float spatial_scale);
	torch::Tensor forward(torch::Tensor input, torch::Tensor rois);
private:
	float spatial_scale = 1.0;
	int output_size;
	torch::nn::functional::AdaptiveMaxPool2dFuncOptions poolOption = torch::nn::functional::AdaptiveMaxPool2dFuncOptions(output_size);
}; TORCH_MODULE(RoIPool);

RoIPoolImpl::RoIPoolImpl(int _output_size, float _spatial_scale) {
	spatial_scale = _spatial_scale;
	output_size = _output_size;
	poolOption = torch::nn::functional::AdaptiveMaxPool2dFuncOptions(output_size);
}
torch::Tensor RoIPoolImpl::forward(torch::Tensor input, torch::Tensor rois) {
	assert(rois.dim() == 2);
	assert(rois.size(1) == 5);
	int num_rois = rois.size(0);
	rois.narrow(1, 1, rois.size(1)-1) = rois.narrow(1, 1, rois.size(1) - 1).mul_(spatial_scale);
	rois = rois.to(torch::kInt);

	rois.select(1, 1) = torch::clamp_max(rois.select(1, 1), input.size(3) -1);
	rois.select(1, 2) = torch::clamp_max(rois.select(1, 2), input.size(2) - 1);
	rois.select(1, 3) = torch::clamp_max(rois.select(1, 3), input.size(3) - 1);
	rois.select(1, 4) = torch::clamp_max(rois.select(1, 4), input.size(2) - 1);

	rois.select(1, 3) = rois.select(1, 3) - rois.select(1, 1);
	rois.select(1, 4) = rois.select(1, 4) - rois.select(1, 2);
	rois = torch::clamp(rois, 0);

	int* roi_p = new int[5 * num_rois];
	std::memcpy(roi_p, rois.to(at::kCPU).detach().data_ptr(), 5 * num_rois * sizeof(int));
	std::vector<torch::Tensor> output;
	output.reserve(num_rois);

	for (int i = 0; i < num_rois; i++) {
		auto roi = rois[i];
		auto im_idx = roi[0];
		
		int x1 = roi_p[i * 5 + 1];
		int y1 = roi_p[i * 5 + 2];
		int w = roi_p[i * 5 + 3];
		int h = roi_p[i * 5 + 4];

		auto im = input.narrow(0, im_idx, 1).narrow(2, y1, h + 1).narrow(3, x1, w + 1);
		output.push_back(torch::nn::functional::adaptive_max_pool2d(im, poolOption));
	}
	delete[] roi_p;
	return torch::cat(output, 0);
}

这次，我们将从libtorch中tensor中取数的操作变化一下，先在循环外围将张量拷贝到数组指针中，然后再在循环内取数。重新计时，这次我们测试结果为：0.03s。显然，我们在2070s显卡上，将pytorch版本的代码转到libtorch中并添加安全性检查，速度变化为：0.126s->0.08s->0.03s。总体提速四倍以上。这个速度从值上要比项目中cffi实现更快，但是从比例上，cffi比pytorch实现提速高于四倍(0.2s->0.34s)。所以仍旧有差距，但是已经是本人目前的最优实现了。

cuda编程提速RoI Pool

实际经验告诉我，pytorch实现的Faster RCNN执行速度应该可以比0.126s要更优。这是为何？我们定位到pytorch的RoI Pool，发现它其实也是通过调用cuda编程加速了该过程。那么显然，libtorch部署同样可以采用这个代码。

直接用torchvision的源代码可能有些问题，要做些修改，主要是需要将FP16精度的部分代码注释掉，

NMS

libtorch的NMS

前面的博客NMS的几种写法中我有分享libtorch版本的NMS。代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
//输入boxes：Nx4; socres: N; thresh:介于(0,1)的float变量
//输出vector<int>向量，存储保留的box的index值
vector<int> nms_libtorch(torch::Tensor boxes, torch::Tensor scores, float thresh);\\函数声明

vector<int> nms_libtorch(torch::Tensor bboxes, torch::Tensor scores, float thresh) {\\函数定义
	auto x1 = bboxes.select(-1, 0);
	auto y1 = bboxes.select(-1, 1);
	auto x2 = bboxes.select(-1, 2);
	auto y2 = bboxes.select(-1, 3);
	auto areas = (x2 - x1)*(y2 - y1);   //[N, ] 每个bbox的面积
	auto tuple_sorted = scores.sort(0, true);    //降序排列
	auto order = std::get<1>(tuple_sorted);

	vector<int>	keep;
	while (order.numel() > 0) {// torch.numel()返回张量元素个数
		if (order.numel() == 1) {//    保留框只剩一个
			auto i = order.item();
			keep.push_back(i.toInt());
			break;
		}
		else {
			auto i = order[0].item();// 保留scores最大的那个框box[i]
			keep.push_back(i.toInt());
		}
		//计算box[i]与其余各框的IOU(思路很好)
		auto order_mask = order.narrow(0, 1, order.size(-1) - 1);
		x1.index({ order_mask });
		x1.index({ order_mask }).clamp(x1[keep.back()].item().toFloat(), 1e10);
		auto xx1 = x1.index({ order_mask }).clamp(x1[keep.back()].item().toFloat(),1e10);// [N - 1, ]
		auto yy1 = y1.index({ order_mask }).clamp(y1[keep.back()].item().toFloat(), 1e10);
		auto xx2 = x2.index({ order_mask }).clamp(0, x2[keep.back()].item().toFloat());
		auto yy2 = y2.index({ order_mask }).clamp(0, y2[keep.back()].item().toFloat());
		auto inter = (xx2 - xx1).clamp(0, 1e10) * (yy2 - yy1).clamp(0, 1e10);// [N - 1, ]

		auto iou = inter / (areas[keep.back()] + areas.index({ order.narrow(0,1,order.size(-1) - 1) }) - inter);//[N - 1, ]
		auto idx = (iou <= thresh).nonzero().squeeze();//注意此时idx为[N - 1, ] 而order为[N, ]
		if (idx.numel() == 0) {
			break;
		}
		order = order.index({ idx + 1 }); //修补索引之间的差值
	}
	return keep;
}

这个代码，看似只有一个循环，时间复杂度近似为O(n)。但是实际上，每次循环内部，都要计算当前框与其他框的IOU，时间复杂度近似为O(n²)。如果框的数量一旦增加，速度上将大打折扣。对于单阶段的目标检测如YOLO，最后经过阈值筛选后再NMS时，框的数量已经很少，几乎不影响速度。但是对于双阶段目标检测如Faster RCNN而言，框的数量一定是很高的，一般可以指定到1000以上。两者不在一个量级上，NMS的速度将极慢，我的2070s上可能将近有几百毫秒。这无疑不能忍受。

cuda编程提速NMS

同样，我们想到了torchvision中的NMS实现，直接拷贝过来，嗯，真香。实操可能会有些编译错误，调试改改应该就能通过。

最后，本文代码只经过速度测试，未上线到生产环境，可能有许多需要改进。