def test_forward(self):
pool = PrRoIPool2D(7, 7, spatial_scale=0.5)
features = torch.rand((4, 16, 24, 32)).cuda()
rois = torch.tensor([
[0, 0, 0, 14, 14],
[1, 14, 14, 28, 28],
]).float().cuda()
out = pool(features, rois)
out_gold = F.avg_pool2d(features, kernel_size=2, stride=1)
self.assertTensorClose(out, torch.stack((
out_gold[0, :, :7, :7],
out_gold[1, :, 7:14, 7:14],
), dim=0))
def test_backward_shapeonly(self):
pool = PrRoIPool2D(2, 2, spatial_scale=0.5)
features = torch.rand((4, 2, 24, 32)).cuda()
rois = torch.tensor([
[0, 0, 0, 4, 4],
[1, 14, 14, 18, 18],
]).float().cuda()
features.requires_grad = rois.requires_grad = True
out = pool(features, rois)
loss = out.sum()
loss.backward()
self.assertTupleEqual(features.size(), features.grad.size())
self.assertTupleEqual(rois.size(), rois.grad.size())
def __init__(self):
super(IoUNet, self).__init__()
pool = PrRoIPool2D(7, 7, spatial_scale=0.5)
from prroi_pool import PrRoIPool2D
class IoUNet(nn.Module):
def __init__(self):
super(IoUNet, self).__init__()
pool = PrRoIPool2D(7, 7, spatial_scale=0.5)
def forward():
out = pool(features, rois)
if __name__ == '__main__':
pool = PrRoIPool2D(2, 2, spatial_scale=0.5)
features = torch.rand((4, 2, 24, 32)).cuda()
rois = torch.tensor([
[0, 0, 0, 4, 4],
[1, 14, 14, 18, 18],
]).float().cuda()
features.requires_grad = rois.requires_grad = True
out = pool(features, rois)
loss = out.sum()
loss.backward()
print(out)
print(loss)
print(rois.grad)