pytorch ROI_pooling最简实现

本文介绍了一个PyTorch版本的ROI Pooling实现方法,适用于需要修改或使用ROI Pooling功能的同学。代码中详细展示了如何通过PyTorch进行ROI Pooling操作,包括输入张量的处理、ROI区域的选择及池化操作。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

代码来源:https://github.com/SirLPS/roi_pooling

注释掉的是我工程中需要的代码,没注释的是pytorch版roi_pooling,供各位需要修改roi_pooling的同学参考

# -*- coding:UTF-8 -*-
import time
import torch
import torch.nn.functional as F
from torch.autograd import Variable


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 relation_pooling(input, roisA,roisB, size=(7, 7), spatial_scale=1.0):  # pytorch version use for loop !!!
#
# 	output = []
# 	roisA = roisA.data.float()
# 	roisB = roisB.data.float()
# 	num_rois = roisA.size(0)
#
# 	h = input.size(2)
# 	w = input.size(3)
#
# 	roisA[:, 1:].mul_(spatial_scale)
# 	roisB[:, 1:].mul_(spatial_scale)
# 	roisA = roisA.long()
# 	roisB = roisB.long()
#
#
# 	for i in range(num_rois):
# 		roiA = roisA[i]
# 		roiB = roisB[i]
# 		xmin = min(roiA[1], roiB[1])
# 		ymin = min(roiA[2], roiB[2])
# 		xmax = max(roiA[3], roiB[3])
# 		ymax = max(roiA[4], roiB[4])
# 		mask = Variable(torch.zeros(h,w)).cuda()
# 		mask[roiA[2]:(roiA[4] + 1), roiA[1]:(roiA[3] + 1)] = 1
# 		mask[roiB[2]:(roiB[4] + 1), roiB[1]:(roiB[3] + 1)] = 1
# 		input = input * mask.unsqueeze(0).unsqueeze(1)
# 		unionbox = torch.LongTensor([0,xmin,ymin,xmax,ymax])
# 		im = input[..., unionbox[2]:(unionbox[4] + 1), unionbox[1]:(unionbox[3] + 1)]
# 		output.append(F.adaptive_max_pool2d(im, size))
#
# 	output = torch.stack(output, dim=0).squeeze()
#
# 	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 = True

	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_pooling(x,rois)
	print('time: {}, batch_size: {}, size: {}, num_rois: {}'.format((time.time() - start) / T,
		                                                                config[0],
		                                                                config[1],
		                                                                config[2]))


	# start = time.time()
	# x = Variable(torch.rand((config[0], 512, config[1], config[1])),requires_grad=True).cuda()
	# roisA = create_rois(config).cuda()
	# roisB = create_rois(config).cuda()
	# for t in range(T):
	# 	output = relation_pooling(x, roisA, roisB)
	# print('time: {}, batch_size: {}, size: {}, num_rois: {}'.format((time.time() - start) / T,
	# 	                                                                config[0],
	# 	                                                                config[1],
	# 	                                                                config[2]))

 

### RoI Pooling in Deep Learning for Object Detection Region of Interest (RoI) Pooling 是一种用于处理不同大小区域的技术,在目标检测中扮演重要角色。该操作允许网络接受任意形状的输入并将其转换成固定尺寸的特征图,从而可以应用于后续的标准全连接层或其他分类器。 #### 工作原理 给定一系列候选框(即感兴趣区域),这些候选框可能来自Selective Search 或者其他提议算法,对于每一个候选框,RoI pooling会执行如下操作: - 将候选框映射到卷积层输出的空间位置上; - 对应于每个候选框的位置被划分为相同数量的小块或子窗口; - 应用大池化或者平均池化来提取每个子窗口内的特征向量;终得到一组长度固定的特征向量作为下一层的输入[^1]。 这种机制确保了无论原始图像中的物体如何缩放和平移变换,都能获得一致维度的特征描述符,提高了模型对尺度变化和空间布局差异的容忍度。 ```python import torch.nn.functional as F def roi_pooling(input, rois, pooled_height, pooled_width): """Apply ROI Pooling to input tensor.""" # Assume `input` is a feature map with shape [batch_size, channels, height, width] # and `rois` contains the coordinates of regions. batch_size = input.size(0) num_rois = rois.size(0) outputs = [] for i in range(num_rois): roi = rois[i] start_h, end_h = int(round(roi[0])), int(round(roi[2])) start_w, end_w = int(round(roi[1])), int(round(roi[3])) region = input[:, :, start_h:end_h, start_w:end_w] out = F.adaptive_max_pool2d(region, output_size=(pooled_height, pooled_width)) outputs.append(out.view(batch_size, -1)) return torch.cat(outputs, dim=1) ``` 上述代码展示了如何实现自适应的大池化版本的ROI Pooling函数。这里使用PyTorch库来进行张量运算,并利用内置的功能实现了灵活调整输出大小的能力。
评论 6
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值