#! /usr/bin/env python3 | |
# -*- coding: utf-8 -*- | |
# File : prroi_pool.py | |
# Author : Jiayuan Mao, Tete Xiao | |
# Email : maojiayuan@gmail.com, jasonhsiao97@gmail.com | |
# Date : 07/13/2018 | |
# | |
# This file is part of PreciseRoIPooling. | |
# Distributed under terms of the MIT license. | |
# Copyright (c) 2017 Megvii Technology Limited. | |
import torch.nn as nn | |
from .functional import prroi_pool2d | |
__all__ = ['PrRoIPool2D'] | |
class PrRoIPool2D(nn.Module): | |
def __init__(self, pooled_height, pooled_width, spatial_scale): | |
super().__init__() | |
self.pooled_height = int(pooled_height) | |
self.pooled_width = int(pooled_width) | |
self.spatial_scale = float(spatial_scale) | |
def forward(self, features, rois): | |
return prroi_pool2d(features, rois, self.pooled_height, self.pooled_width, self.spatial_scale) | |