pytorch

1
## @package onnx
2
# Module caffe2.python.onnx.tests.conversion_test
3

4

5

6

7

8
import json
9
import tempfile
10
import textwrap
11
import traceback
12
import unittest
13
import zipfile
14

15
from caffe2.proto import caffe2_pb2
16
from caffe2.python import brew, core
17
from caffe2.python.model_helper import ModelHelper
18
from click.testing import CliRunner
19
import numpy as np
20
from onnx import helper, ModelProto, TensorProto
21
from caffe2.python.onnx.helper import c2_native_run_net
22

23
from caffe2.python.onnx.bin.conversion import caffe2_to_onnx, onnx_to_caffe2
24
import caffe2.python.onnx.backend as c2
25
from caffe2.python.onnx.tests.test_utils import TestCase
26

27

28
class TestConversion(TestCase):
29
    def _run_command(self, cmd, *args, **kwargs):
30
        runner = CliRunner()
31
        result = runner.invoke(cmd, *args, **kwargs)
32
        self.assertEqual(result.exit_code, 0, textwrap.dedent('''
33
        Command exited with non-zero exit code:
34
        output: {}
35
        exception: {}
36
        exc_info: {}
37
        '''.format(result.output,
38
                   result.exception,
39
                   traceback.format_exception(*result.exc_info))))
40
        return result
41

42
    def test_caffe2_to_onnx(self):
43
        caffe2_net = tempfile.NamedTemporaryFile()
44
        caffe2_init_net = tempfile.NamedTemporaryFile()
45
        output = tempfile.NamedTemporaryFile()
46

47
        model = ModelHelper(name='caffe2-to-onnx-test')
48
        brew.relu(model, ["X"], "Y")
49
        caffe2_net.write(model.net.Proto().SerializeToString())
50
        caffe2_net.flush()
51

52
        init_model = ModelHelper(name='caffe2-to-onnx-init-test')
53
        init_model.net.GivenTensorFill([], 'X', shape=[2, 2],
54
                                       values=np.zeros((2, 2)).flatten().astype(float))
55
        caffe2_init_net.write(init_model.net.Proto().SerializeToString())
56
        caffe2_init_net.flush()
57

58
        self._run_command(
59
            caffe2_to_onnx, [
60
                caffe2_net.name,
61
                '--caffe2-init-net', caffe2_init_net.name,
62
                '--output', output.name,
63
            ],
64
            catch_exceptions=False,
65
        )
66

67
        onnx_model = ModelProto()
68
        onnx_model.ParseFromString(output.read())
69
        self.assertEqual(len(onnx_model.graph.node), 1)
70
        self.assertEqual(onnx_model.graph.node[0].op_type, 'Relu')
71
        self.assertEqual(len(onnx_model.graph.initializer), 1)
72
        self.assertEqual(onnx_model.graph.initializer[0].name, onnx_model.graph.input[0].name)
73

74
    def test_caffe2_to_onnx_value_info(self):
75
        caffe2_net = tempfile.NamedTemporaryFile()
76
        output = tempfile.NamedTemporaryFile()
77

78
        model = ModelHelper(name='caffe2-to-onnx-test')
79
        brew.relu(model, ["X"], "Y")
80
        caffe2_net.write(model.net.Proto().SerializeToString())
81
        caffe2_net.flush()
82

83
        args = [caffe2_net.name, '--output', output.name]
84
        self.assertRaisesRegex(Exception,
85
                               'value info',
86
                               self._run_command, caffe2_to_onnx, args)
87

88
        args.extend([
89
            '--value-info',
90
            json.dumps({
91
                'X': (TensorProto.FLOAT, (2, 2)),
92
            })])
93
        self._run_command(caffe2_to_onnx, args)
94

95
        onnx_model = ModelProto()
96
        onnx_model.ParseFromString(output.read())
97
        self.assertEqual(len(onnx_model.graph.node), 1)
98
        self.assertEqual(onnx_model.graph.node[0].op_type, 'Relu')
99
        self.assertEqual(len(onnx_model.graph.initializer), 0)
100

101
    @unittest.skip("Disabled due to onnx optimizer deprecation")
102
    def test_onnx_to_caffe2(self):
103
        onnx_model = tempfile.NamedTemporaryFile()
104
        output = tempfile.NamedTemporaryFile()
105
        init_net_output = tempfile.NamedTemporaryFile()
106

107
        node_def = helper.make_node(
108
            "Mul", ["X", "W"], ["Y"])
109
        graph_def = helper.make_graph(
110
            [node_def],
111
            "test",
112
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3)),
113
             helper.make_tensor_value_info("W", TensorProto.FLOAT, (1, 3))],
114
            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 3))],
115
            initializer=[helper.make_tensor("W",
116
                                            TensorProto.FLOAT,
117
                                            [1, 3],
118
                                            np.zeros((1, 3)).flatten().astype(float))])
119
        model_def = helper.make_model(graph_def, producer_name='onnx-to-caffe2-test')
120
        onnx_model.write(model_def.SerializeToString())
121
        onnx_model.flush()
122

123
        self._run_command(
124
            onnx_to_caffe2, [
125
                onnx_model.name,
126
                '--output', output.name,
127
                '--init-net-output', init_net_output.name,
128
            ])
129

130
        caffe2_net = caffe2_pb2.NetDef()
131
        caffe2_net.ParseFromString(output.read())
132
        self.assertEqual(len(caffe2_net.op), 1)
133
        self.assertEqual(caffe2_net.op[0].type, 'Mul')
134

135
        caffe2_init_net = caffe2_pb2.NetDef()
136
        caffe2_init_net.ParseFromString(init_net_output.read())
137
        self.assertEqual(len(caffe2_init_net.op), 1)
138
        self.assertEqual(set(sum([list(init_op.output)
139
                                  for init_op in caffe2_init_net.op], [])),
140
                         {'W'})
141

142
    def test_onnx_to_caffe2_zipfile(self):
143
        buf = tempfile.NamedTemporaryFile()
144
        onnx_model = zipfile.ZipFile(buf, 'w')
145

146
        node_def = helper.make_node(
147
            "MatMul", ["X", "W"], ["Y"])
148
        X = np.random.rand(2, 3).astype(np.float32)
149
        W = np.random.rand(3, 2).flatten().astype(np.float32)
150
        graph_def = helper.make_graph(
151
            [node_def],
152
            "test",
153
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3)),
154
             helper.make_tensor_value_info("W", TensorProto.FLOAT, (3, 2))],
155
            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2))],
156
            initializer=[helper.make_tensor("W",
157
                                            TensorProto.FLOAT,
158
                                            [3, 2],
159
                                            W.tobytes(),
160
                                            raw=True)])
161
        model_def = helper.make_model(graph_def, producer_name='onnx-to-caffe2-test')
162
        onnx_model.writestr('__MODEL_PROTO', model_def.SerializeToString())
163
        onnx_model.writestr('W', W.tobytes())
164
        onnx_model.close()
165

166
        W = W.reshape((3, 2))
167
        Y_expect = np.matmul(X, W)
168

169
        c2_model = c2.prepare_zip_archive(buf)
170
        Y = c2_model.run(X).Y
171
        np.testing.assert_allclose(Y, Y_expect)
172

173
    def _make_fake_if_op(self, true_nodes, false_nodes, output_types):
174
        true = helper.make_tensor("condition", TensorProto.BOOL, (), [True])
175
        true_graph = helper.make_graph(true_nodes, "true_graph", [], [
176
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2)),
177
        ])
178
        false_graph = helper.make_graph(false_nodes, "false_graph", [], [
179
            helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2)),
180
        ])
181
        if_inputs = ["condition"]
182
        if_outputs = [name for _, _, name in output_types]
183
        retval_nodes = [
184
            helper.make_node("Constant", [], ["condition"], value=true),
185
            helper.make_node("If", if_inputs, if_outputs, then_branch=true_graph,
186
                             else_branch=false_graph)
187
        ]
188
        return retval_nodes
189

190
    def test_onnx_to_caffe2_if(self):
191
        true_nodes = [helper.make_node(
192
            "MatMul", ["X", "W"], ["Y"])]
193
        false_nodes = [helper.make_node("Slice", ["X"], ["Y"], axes=[0, 1],
194
                                        starts=[0, 0], ends=[2, 2])]
195
        nodes = self._make_fake_if_op(true_nodes, false_nodes, [(TensorProto.FLOAT, (2, 2), "Y")])
196
        X = np.random.rand(2, 3).astype(np.float32)
197
        W = np.random.rand(3, 2).flatten().astype(np.float32)
198
        graph_def = helper.make_graph(
199
            nodes,
200
            "test",
201
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3)),
202
             helper.make_tensor_value_info("W", TensorProto.FLOAT, (3, 2))],
203
            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2))],
204
            initializer=[helper.make_tensor("W",
205
                                            TensorProto.FLOAT,
206
                                            [3, 2],
207
                                            W.tolist())]
208
        )
209
        onnx_id = helper.make_opsetid("", 9)
210
        model_def = helper.make_model(graph_def, producer_name='onnx-to-caffe2-test',
211
                                      opset_imports=[onnx_id])
212

213
        p = c2.prepare(model_def)
214
        Y = np.matmul(X, W.reshape(3, 2))
215
        out = p.run(X)
216
        np.testing.assert_allclose(out.Y, Y)
217

218
    # input_types and output_types are lists of triples of (name, type, shape)
219
    def _make_fake_loop_op(self, body_nodes, input_types, output_types):
220
        ten = helper.make_tensor("trip_count_value", TensorProto.INT64, (1,), [10])
221
        true = helper.make_tensor("condition", TensorProto.BOOL, (1,), [True])
222
        # lcd is a dummy loop-carried dependency that only exists because
223
        # right now the schema checker is broken and assumes a variadic
224
        # input needs at least one value.
225
        graph_inputs = [helper.make_tensor_value_info("i", TensorProto.INT64, (1,)),
226
                        helper.make_tensor_value_info("cond", TensorProto.BOOL, (1,))]
227
        for type, shape, name in input_types:
228
            graph_inputs.append(helper.make_tensor_value_info("_" + name, type, shape))
229
        graph_outputs = [helper.make_tensor_value_info("cond", TensorProto.BOOL, (1,))]
230
        for type, shape, name in output_types:
231
            graph_outputs.append(helper.make_tensor_value_info("_" + name, type, shape))
232
        body_graph = helper.make_graph(body_nodes, "body_graph", graph_inputs,
233
                                       graph_outputs)
234
        loop_inputs = ["trip_count", "condition"]
235
        loop_inputs.extend([name for _, _, name in input_types])
236
        loop_outputs = [name for _, _, name in output_types]
237
        retval_nodes = [
238
            helper.make_node("Constant", [], ["trip_count"], value=ten),
239
            helper.make_node("Constant", [], ["condition"], value=true),
240
            helper.make_node("Loop", loop_inputs, loop_outputs, body=body_graph)
241
        ]
242
        return retval_nodes
243

244
    @unittest.skip("Disabled due to onnx optimizer deprecation")
245
    def test_onnx_to_caffe2_loop(self):
246
        body_nodes = [helper.make_node(
247
            "MatMul", ["_X", "W"], ["_Y"])]
248
        nodes = self._make_fake_loop_op(body_nodes,
249
                                        [(TensorProto.FLOAT, (2, 2), "X")],
250
                                        [(TensorProto.FLOAT, (2, 2), "Y")])
251
        X = np.random.rand(2, 2).astype(np.float32)
252
        W = np.random.rand(2, 2).flatten().astype(np.float32)
253
        graph_def = helper.make_graph(
254
            nodes,
255
            "test",
256
            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 2)),
257
             helper.make_tensor_value_info("W", TensorProto.FLOAT, (2, 2))],
258
            [helper.make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2))],
259
            initializer=[helper.make_tensor("W",
260
                                            TensorProto.FLOAT,
261
                                            [2, 2],
262
                                            W.tolist())]
263
        )
264
        model_def = helper.make_model(graph_def, producer_name='onnx-to-caffe2-test')
265
        Y = X
266
        for _ in range(10):
267
            Y = np.matmul(Y, W.reshape(2, 2))
268
        p = c2.prepare(model_def)
269
        out = p.run(X)
270
        np.testing.assert_allclose(out.Y, Y)
271

272
    # TODO investigate why this is failing after changing Reshape
273
    # operator from taking the new shape as attribute to as input
274
    @unittest.skip('Start failing after Reshape op change')
275
    def test_convert_end2end(self):
276
        predict_net_f = tempfile.NamedTemporaryFile()
277
        init_net_f = tempfile.NamedTemporaryFile()
278
        onnx_model_f = tempfile.NamedTemporaryFile()
279

280
        x = 'X'
281
        w = 'W'
282
        b = 'b'
283
        y = 'Y'
284

285
        predict_net = caffe2_pb2.NetDef()
286
        predict_net.name = 'test-convert-end2end'
287
        predict_net.external_input[:] = [x, w, b]
288
        predict_net.external_output[:] = [y]
289
        predict_net.op.extend([
290
            core.CreateOperator(
291
                'FC',
292
                inputs=[x, w, b],
293
                outputs=[y],
294
                axis=2,
295
            ),
296
        ])
297
        predict_net_f.write(predict_net.SerializeToString())
298
        predict_net_f.flush()
299

300
        init_net = caffe2_pb2.NetDef()
301
        init_net.name = 'test-convert-end2end-init'
302
        init_net.external_output[:] = [w, b]
303
        x_val = np.random.randn(1, 3, 2).astype(np.float32)
304
        w_val = np.random.randn(4, 2).astype(np.float32)
305
        b_val = np.random.randn(4).astype(np.float32)
306
        init_net.op.extend([
307
            core.CreateOperator(
308
                'GivenTensorFill',
309
                [],
310
                [w],
311
                values=w_val,
312
                shape=w_val.shape,
313
            ),
314
            core.CreateOperator(
315
                'GivenTensorFill',
316
                [],
317
                [b],
318
                values=b_val,
319
                shape=b_val.shape,
320
            ),
321
        ])
322
        init_net_f.write(init_net.SerializeToString())
323
        init_net_f.flush()
324

325
        y_val = np.matmul(x_val, w_val.transpose()) + b_val
326
        for _ in range(5):
327
            self._run_command(
328
                caffe2_to_onnx, [
329
                    predict_net_f.name,
330
                    '--caffe2-init-net', init_net_f.name,
331
                    '--output', onnx_model_f.name,
332
                    '--value-info',
333
                    json.dumps({
334
                        x: (TensorProto.FLOAT, (1, 3, 2)),
335
                    }),
336
                ],
337
                catch_exceptions=False,
338
            )
339

340
            onnx_model_f.seek(0)
341
            onnx_model = ModelProto()
342
            onnx_model.ParseFromString(onnx_model_f.read())
343
            np.testing.assert_almost_equal(
344
                c2.run_model(
345
                    onnx_model, {onnx_model.graph.input[0].name: x_val}),
346
                [y_val])
347

348
            self._run_command(
349
                onnx_to_caffe2, [
350
                    onnx_model_f.name,
351
                    '--output', predict_net_f.name,
352
                    '--init-net-output', init_net_f.name,
353
                ])
354
            predict_net_f.seek(0)
355
            predict_net = caffe2_pb2.NetDef()
356
            predict_net.ParseFromString(predict_net_f.read())
357
            init_net_f.seek(0)
358
            init_net = caffe2_pb2.NetDef()
359
            init_net.ParseFromString(init_net_f.read())
360
            x = predict_net.external_input[0]
361
            np.testing.assert_almost_equal(c2_native_run_net(init_net=init_net,
362
                                                             predict_net=predict_net,
363
                                                             inputs={x: x_val})[1],
364
                                           [y_val])
365