Spaces:
Build error
Build error
File size: 1,473 Bytes
8f87579 |
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 57 |
# -*- coding: utf-8 -*-
# File : test_prroi_pooling2d.py
# Author : Jiayuan Mao
# Email : [email protected]
# Date : 18/02/2018
#
# This file is part of Jacinle.
import unittest
import torch
import torch.nn as nn
import torch.nn.functional as F
from jactorch.utils.unittest import TorchTestCase
from prroi_pool import PrRoIPool2D
class TestPrRoIPool2D(TorchTestCase):
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())
if __name__ == '__main__':
unittest.main()
|