pytorch

1

2

3

4

5
import numpy as np
6
from numpy.testing import assert_array_equal
7

8
from caffe2.python import core, workspace
9
from caffe2.python.test_util import TestCase
10
from caffe2.proto import caffe2_pb2
11

12

13
class TestLengthsToShapeOps(TestCase):
14
    def test_lengths_to_shape_ops(self):
15
        workspace.FeedBlob('l', np.array([200, 200, 200], dtype=np.int32))
16
        workspace.RunOperatorOnce(core.CreateOperator(
17
            'LengthsToShape', ['l'], ['s']))
18
        workspace.FeedBlob('res', np.array([3, 200], dtype=np.int32))
19
        assert_array_equal(workspace.FetchBlob('s'), workspace.FetchBlob('res'))
20

21
    def test_reshape_ops(self):
22
        workspace.FeedBlob('res', np.array([[0, 0, 0, 0]], dtype=np.float32))
23
        workspace.FeedBlob('shape', np.array([1, 4], dtype=np.int32))
24
        workspace.FeedBlob('input', np.zeros((2, 2), dtype=np.float32))
25
        workspace.RunOperatorOnce(core.CreateOperator(
26
            'Reshape', ['input', 'shape'], ['output', 'old_shape']))
27
        assert_array_equal(workspace.FetchBlob('output'),
28
                           workspace.FetchBlob('res'))
29

30
    def test_basic_reshape(self):
31
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(2, 4))
32
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(2, 4), arg_shape=False)
33

34
    def test_missing_dim(self):
35
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(-1, 8))
36
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(-1, 8), arg_shape=False)
37

38
    def test_in_place(self):
39
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(-1, 8), in_place=True)
40
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(-1, 8),
41
                     in_place=True, arg_shape=False)
42

43
    def test_zero_dim(self):
44
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(0, 0, 0),
45
                     expected_shape=(4, 2, 1))
46
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(0, 0, 0),
47
                     expected_shape=(4, 2, 1), arg_shape=False)
48
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(0, 2, 1),
49
                     expected_shape=(4, 2, 1))
50
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(0, 2, 1),
51
                     expected_shape=(4, 2, 1), arg_shape=False)
52
        _test_reshape_output_and_gradient(old_shape=(0, 0), new_shape=(0, 0, 0),
53
                     expected_shape=(0, 0, 0), arg_shape=False)
54

55
    def test_zero_dim_and_missing_dim(self):
56
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(0, -1, 0),
57
                     expected_shape=(4, 2, 1))
58
        _test_reshape_output_and_gradient(old_shape=(4, 2, 1), new_shape=(0, -1, 0),
59
                     expected_shape=(4, 2, 1), arg_shape=False)
60
        _test_reshape_output_and_gradient(old_shape=(4, 3, 2), new_shape=(-1, 0),
61
                     expected_shape=(8, 3))
62
        _test_reshape_output_and_gradient(old_shape=(4, 3, 2), new_shape=(-1, 0),
63
                     expected_shape=(8, 3), arg_shape=False)
64

65
        # empty tensor will just have -1 dim = 0
66
        _test_reshape_output_and_gradient(
67
            old_shape=(2, 0),
68
            new_shape=(-1, 0),
69
            expected_shape=(0, 0),
70
            arg_shape=False
71
        )
72

73
    def test_backprop(self):
74
        old_shape = (4, 2, 1)
75
        new_shape = (1, 8)
76
        X = np.random.rand(*old_shape).astype(np.float32)
77
        Y = np.random.rand(*new_shape).astype(np.float32)
78

79
        net = core.Net('net')
80

81
        net.GivenTensorFill([], 'X', shape=old_shape, values=X.flatten())
82
        net.GivenTensorFill([], 'Y', shape=new_shape, values=Y.flatten())
83

84
        net.Reshape(['X'], ['X_out', 'old_shape'], shape=new_shape)
85
        net.DotProduct(['X_out', 'Y'], 'Z')
86
        net.AddGradientOperators(['Z'])
87

88
        workspace.RunNetOnce(net)
89

90
        Z = workspace.FetchBlob('Z')
91
        X_grad = workspace.FetchBlob('X_grad')
92

93
        # Check forward computation
94
        np.testing.assert_allclose(
95
            Z.squeeze(), X.reshape(new_shape).dot(Y.T).squeeze(), rtol=1e-5)
96

97
        # Check the shape of the gradient
98
        np.testing.assert_array_equal(X_grad.shape, X.shape)
99

100
        # Check the gradient
101
        np.testing.assert_allclose(X_grad, Y.reshape(old_shape), rtol=1e-5)
102

103
    def test_input_shape_changes(self):
104
        workspace.FeedBlob(
105
            'input_blob',
106
            np.array(np.random.rand(10, 20, 10), dtype=np.float32))
107
        net = core.Net('mynet')
108
        z, _ = net.Reshape('input_blob',
109
                           ['z_reshape', 'dummy_size'],
110
                           shape=(-1, 10))
111
        workspace.CreateNet(net)
112
        workspace.RunNet(net)
113
        workspace.FeedBlob(
114
            'input_blob',
115
            np.array(np.random.rand(10, 40, 10), dtype=np.float32))
116
        workspace.RunNet(net)
117

118
    def test_nonempty_tensor_gradient(self):
119
        old_shape = [4, 2]
120
        new_shape = [1, 2, -1]
121
        expected_new_shape = [1, 2, 4]
122
        _test_reshape_output_and_gradient(
123
            old_shape=old_shape,
124
            new_shape=new_shape,
125
            expected_shape=expected_new_shape,
126
            expected_gradient=np.ones(shape=old_shape)
127
        )
128

129
    def test_empty_tensor(self):
130
        old_shape = [4, 0]
131
        new_shape = [1, -1]
132
        expected_new_shape = [1, 0]
133
        _test_reshape_output_and_gradient(
134
            old_shape=old_shape,
135
            new_shape=new_shape,
136
            expected_shape=expected_new_shape,
137
            expected_gradient=np.empty(shape=old_shape)
138
        )
139

140
    def test_one_dim_empty_tensor_gradient(self):
141
        old_shape = (0,)
142
        new_shape = [1, -1]
143
        expected_new_shape = [1, 0]
144
        _test_reshape_output_and_gradient(
145
            old_shape=old_shape,
146
            new_shape=new_shape,
147
            expected_shape=expected_new_shape,
148
            expected_gradient=np.empty(shape=old_shape)
149
        )
150

151
    def test_one_dim_and_empty_tensor(self):
152
        old_shape = (0,)
153
        new_shape = [0, -1]
154
        expected_new_shape = [0, 0]
155
        _test_reshape_output_and_gradient(old_shape=old_shape, new_shape=new_shape, expected_shape=expected_new_shape)
156

157
    def test_scalar_to_tensor(self):
158
        old_shape = ()
159
        new_shape = [1, -1]
160
        expected_new_shape = [1, 1]
161
        _test_reshape_output_and_gradient(old_shape=old_shape, new_shape=new_shape, expected_shape=expected_new_shape)
162

163

164
def _test_reshape_output_and_gradient(
165
    old_shape,
166
    new_shape,
167
    expected_shape=None,
168
    arg_shape=True,
169
    in_place=False,
170
    expected_gradient=None
171
):
172
    devices = [core.DeviceOption(caffe2_pb2.CPU, 0)]
173
    if workspace.NumGpuDevices() > 0:
174
        devices.append(core.DeviceOption(workspace.GpuDeviceType, 0))
175

176
    for device_opt in devices:
177
        with core.DeviceScope(device_opt):
178
            if expected_shape is None:
179
                expected_shape = new_shape
180
            net = core.Net('net')
181

182
            if len(old_shape) == 0:
183
                # scalar, convert to tensor before feeding to blob
184
                X = np.atleast_1d(np.random.rand(*old_shape))
185
            else:
186
                X = np.random.rand(*old_shape).astype(np.float32)
187
            blob_in = 'X'
188
            blob_out = blob_in if in_place else blob_in + '_out'
189

190
            if arg_shape:
191
                out, _ = net.Reshape([blob_in], [blob_out, 'old_shape'], shape=new_shape)
192
            else:
193
                out, _ = net.Reshape([blob_in, 'new_shape'], [blob_out, 'old_shape'])
194
                workspace.FeedBlob('new_shape', np.asarray(new_shape))
195

196
            workspace.FeedBlob(blob_in, X)
197
            if expected_gradient is not None:
198
                net.AddGradientOperators([out])
199
            workspace.CreateNet(net)
200
            workspace.RunNetOnce(net)
201

202
            Y = workspace.FetchBlob(blob_out)
203
            np.testing.assert_allclose(Y, X.reshape(expected_shape))
204
            if expected_gradient is not None:
205
                data_grad = workspace.FetchBlob(blob_in + '_grad')
206
                np.testing.assert_array_equal(data_grad, expected_gradient)
207

208

209
if __name__ == "__main__":
210
    import unittest
211
    unittest.main()
212