pytorch

1

2

3

4

5
from caffe2.python import core
6
import caffe2.python.hypothesis_test_util as hu
7
import caffe2.python.serialized_test.serialized_test_util as serial
8
import hypothesis.strategies as st
9
import numpy as np
10
import unittest
11

12

13
class TestPad(serial.SerializedTestCase):
14
    @serial.given(pad_t=st.integers(-5, 0),
15
           pad_l=st.integers(-5, 0),
16
           pad_b=st.integers(-5, 0),
17
           pad_r=st.integers(-5, 0),
18
           mode=st.sampled_from(["constant", "reflect", "edge"]),
19
           size_w=st.integers(16, 128),
20
           size_h=st.integers(16, 128),
21
           size_c=st.integers(1, 4),
22
           size_n=st.integers(1, 4),
23
           **hu.gcs)
24
    def test_crop(self,
25
                  pad_t, pad_l, pad_b, pad_r,
26
                  mode,
27
                  size_w, size_h, size_c, size_n,
28
                  gc, dc):
29
        op = core.CreateOperator(
30
            "PadImage",
31
            ["X"],
32
            ["Y"],
33
            pad_t=pad_t,
34
            pad_l=pad_l,
35
            pad_b=pad_b,
36
            pad_r=pad_r,
37
        )
38
        X = np.random.rand(
39
            size_n, size_c, size_h, size_w).astype(np.float32)
40

41
        def ref(X):
42
            return (X[:, :, -pad_t:pad_b or None, -pad_l:pad_r or None],)
43

44
        self.assertReferenceChecks(gc, op, [X], ref)
45
        self.assertDeviceChecks(dc, op, [X], [0])
46

47

48
if __name__ == "__main__":
49
    unittest.main()
50