forked from pytorch/pytorch
/
upsample_nearest_op_test.py
43 lines (33 loc) · 1.27 KB
/
upsample_nearest_op_test.py
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
from __future__ import absolute_import, division, print_function, unicode_literals
import unittest
import caffe2.python.hypothesis_test_util as hu
import hypothesis.strategies as st
import numpy as np
from caffe2.python import core, dyndep
from hypothesis import given
dyndep.InitOpsLibrary("@/caffe2/modules/detectron:detectron_ops")
class TestUpsampleNearestOp(hu.HypothesisTestCase):
@given(
N=st.integers(1, 3),
H=st.integers(10, 300),
W=st.integers(10, 300),
scale=st.integers(1, 3),
**hu.gcs
)
def test_upsample_nearest_op(self, N, H, W, scale, gc, dc):
C = 32
X = np.random.randn(N, C, H, W).astype(np.float32)
op = core.CreateOperator("UpsampleNearest", ["X"], ["Y"], scale=scale)
def ref(X):
outH = H * scale
outW = W * scale
outH_idxs, outW_idxs = np.meshgrid(
np.arange(outH), np.arange(outW), indexing="ij"
)
inH_idxs = (outH_idxs / scale).astype(np.int32)
inW_idxs = (outW_idxs / scale).astype(np.int32)
Y = X[:, :, inH_idxs, inW_idxs]
return [Y]
self.assertReferenceChecks(device_option=gc, op=op, inputs=[X], reference=ref)
if __name__ == "__main__":
unittest.main()