test_tracing.py 10.5 KB
Newer Older
1 2 3 4 5 6 7 8
# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
M
Megvii Engine Team 已提交
9
import io
10
from tempfile import mkstemp
M
Megvii Engine Team 已提交
11

M
Megvii Engine Team 已提交
12
import numpy as np
13
import pytest
M
Megvii Engine Team 已提交
14

15
import megengine.core.tensor.megbrain_graph as G
16
import megengine.functional as F
17 18
import megengine.utils.comp_graph_tools as cgtools
from megengine import tensor
19
from megengine.core._trace_option import set_symbolic_shape
M
Megvii Engine Team 已提交
20
from megengine.core.ops import builtin as ops
21
from megengine.core.ops.builtin import Elemwise
22
from megengine.core.tensor.utils import isscalar
23
from megengine.functional import exp, log
M
Megvii Engine Team 已提交
24
from megengine.jit import exclude_from_trace, trace
25
from megengine.random import normal, uniform
M
Megvii Engine Team 已提交
26 27 28 29 30 31 32


def test_trace():
    for symbolic in [False, True]:

        @trace(symbolic=symbolic)
        def f(x):
33
            return -x
M
Megvii Engine Team 已提交
34

35 36
        x = tensor([1])
        y = f(x).numpy()
M
Megvii Engine Team 已提交
37 38

        for i in range(3):
39
            np.testing.assert_equal(f(x).numpy(), y)
M
Megvii Engine Team 已提交
40 41 42


def test_exclude_from_trace():
43
    for symbolic in [False]:
M
Megvii Engine Team 已提交
44 45 46

        @trace(symbolic=symbolic)
        def f(x):
47
            x = -x
M
Megvii Engine Team 已提交
48 49
            with exclude_from_trace():
                if i % 2:
50 51
                    x = -x
            x = -x
M
Megvii Engine Team 已提交
52 53
            return x

54
        x = tensor([1])
M
Megvii Engine Team 已提交
55 56

        for i in range(3):
57 58
            y = f(x).numpy()
            np.testing.assert_equal(f(x).numpy(), y)
M
Megvii Engine Team 已提交
59 60 61 62 63 64 65 66


def test_print_in_trace():
    for symbolic in [False]:  # cannot read value in symbolic mode

        @trace(symbolic=symbolic)
        def f(x):
            nonlocal buf
67
            x = -x
M
Megvii Engine Team 已提交
68
            buf = x.numpy()
69
            x = -x
M
Megvii Engine Team 已提交
70 71 72
            return x

        buf = None
73
        x = tensor([1])
M
Megvii Engine Team 已提交
74 75

        for i in range(3):
76
            y = f(x).numpy()
M
Megvii Engine Team 已提交
77 78
            z = buf
            buf = None
79
            np.testing.assert_equal(f(x).numpy(), y)
M
Megvii Engine Team 已提交
80
            np.testing.assert_equal(z, buf)
M
Megvii Engine Team 已提交
81 82 83


def test_dump():
84 85
    @trace(symbolic=True, capture_as_const=True)
    def f(a, b):
86
        return a + b
87

88 89 90
    a = tensor([2])
    b = tensor([4])
    y = f(a, b).numpy()
91 92

    for i in range(3):
93
        np.testing.assert_equal(f(a, b).numpy(), y)
94 95

    file = io.BytesIO()
96 97
    dump_info = f.dump(file)
    assert dump_info.nr_opr == 3
98 99
    np.testing.assert_equal(dump_info.inputs, ["arg_0", "arg_1"])
    np.testing.assert_equal(dump_info.outputs, ["ADD(arg_0,arg_1)[4]"])
100
    file.seek(0)
101
    result = cgtools.load_and_inference(file, [a, b])
102 103 104 105
    np.testing.assert_equal(result[0], y)


def test_capture_dump():
106
    a = tensor([2])
107 108 109

    @trace(symbolic=True, capture_as_const=True)
    def f(x):
110
        return x * a
111

112 113
    x = tensor([3])
    y = f(x).numpy()
114 115

    for i in range(3):
116
        np.testing.assert_equal(f(x).numpy(), y)
117 118 119 120

    file = io.BytesIO()
    f.dump(file)
    file.seek(0)
121
    result = cgtools.load_and_inference(file, [x])
122 123 124 125
    np.testing.assert_equal(result[0], y)


def test_dump_volatile():
126
    p = tensor([2])
127

M
Megvii Engine Team 已提交
128 129
    @trace(symbolic=True, capture_as_const=True)
    def f(x):
130
        return x * p
M
Megvii Engine Team 已提交
131

132 133
    x = tensor([3])
    y = f(x).numpy()
M
Megvii Engine Team 已提交
134 135

    for i in range(3):
136
        np.testing.assert_equal(f(x).numpy(), y)
M
Megvii Engine Team 已提交
137 138

    file = io.BytesIO()
139
    f.dump(file, optimize_for_inference=False)
140
    file.seek(0)
141
    cg, _, outputs = G.load_graph(file)
142 143 144
    (out,) = outputs
    assert (
        cgtools.get_owner_opr_type(cgtools.get_owner_opr_inputs(out)[1])
145
        == "ImmutableTensor"
146
    )
147 148 149 150 151 152 153


def test_trace_profiler():
    for symbolic in [False, True]:

        @trace(symbolic=symbolic, profiling=True)
        def f(x):
154
            return -x
155

156 157
        x = tensor([1])
        y = f(x).numpy()
158

159 160
        f(x)
        f(x)  # XXX: has to run twice
161 162 163

        out = f.get_profile()
        assert out.get("profiler")
164 165


166
@pytest.mark.skip(reason="force opt_level=0 when building graph")
167
def test_goptions():
168 169
    @trace(symbolic=True, opt_level=0, capture_as_const=True)
    def f(x):
170 171 172 173
        # directly return x / x will not trigger gopt
        # since there's no way to tell the two x are the same
        y = 2.0 * x
        return y / y
174 175 176

    @trace(symbolic=True, opt_level=1, capture_as_const=True)
    def g(x):
177 178
        y = 2.0 * x
        return y / y
179

180 181 182
    d = tensor(0.0)
    assert not np.isfinite(f(d).numpy())
    np.testing.assert_equal(g(d).numpy().item(), 1.0)
183 184


185
@pytest.mark.skip(reason="force opt_level=0 when building graph")
186 187 188 189 190 191 192 193 194
def test_goptions_log_sum_exp():
    @trace(symbolic=True, opt_level=0, capture_as_const=True)
    def f(x, y):
        return log(exp(x) + exp(y))

    @trace(symbolic=True, opt_level=1, capture_as_const=True)
    def g(x, y):
        return log(exp(x) + exp(y))

195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212
    val = 1.0e4
    d = tensor(val)
    o = tensor(0.0)
    assert not np.isfinite(f(d, o).numpy())
    np.testing.assert_almost_equal(g(d, o), val)


@pytest.mark.skip(reason="could not use opt_level=0 with dump")
def test_goptions_log_exp():
    @trace(symbolic=True, opt_level=0, capture_as_const=True)
    def f(x):
        return log(exp(x))

    @trace(symbolic=True, opt_level=1, capture_as_const=True)
    def g(x):
        return log(exp(x))

    f(tensor(1.0))
213
    _, out = mkstemp()
214 215
    f.dump(out, optimize_for_inference=False)
    *_, outputs = G.load_graph(out)
216 217
    oprs_1 = cgtools.get_oprs_seq(outputs)

218
    g(tensor(1.0))
219 220
    g.dump(out, optimize_for_inference=False)
    *_, outputs = G.load_graph(out)
221 222 223 224 225 226 227 228 229 230 231 232
    oprs_2 = cgtools.get_oprs_seq(outputs)

    assert len(oprs_1) - len(oprs_2) == 2


def test_optimize_for_inference():
    @trace(symbolic=True, capture_as_const=True)
    def f(x):
        return exp(x)

    _, out = mkstemp()
    f(tensor(5.0))
233
    f.dump(out, enable_io16xc32=True)
234

235
    res = G.load_graph(out)
236 237
    computing_input = res.output_vars_list[0].owner.inputs[0]
    assert computing_input.dtype == np.float16
238 239


240 241 242
def test_optimize_for_inference_broadcast():
    a = tensor(np.ones(1, dtype=np.float32))

243
    @trace(capture_as_const=True, symbolic_shape=True)
244
    def f():
245
        return a._broadcast(tensor([1, 10], dtype=np.int32))
246 247 248 249 250

    f()
    f.dump(io.BytesIO())


251 252 253 254 255
def test_trace_cvt_bool():
    x = tensor([0], dtype=np.int32)

    @trace(symbolic=True)
    def f(x):
256 257 258 259
        a = x.shape
        b = a[0]
        assert isscalar(b)
        return b == 0
260 261

    for i in range(3):
262
        np.testing.assert_equal(f(x).numpy(), False)
263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278


def test_trace_reshape():
    for symbolic in [False, True]:
        x1 = tensor(np.random.randn(2, 10, 10))
        x2 = tensor(np.random.randn(4, 10, 10))
        x3 = tensor(np.random.randn(8, 10, 10))

        @trace(symbolic=symbolic, capture_as_const=True)
        def f(x):
            y = x.reshape(x.shape[0], 100)
            return y

        f(x1)
        f(x2)
        f(x3)
279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311


def test_trace_topk():
    x = tensor([5, 2, 7, 1, 0, 3, 2])

    @trace(symbolic=True)
    def f(x):
        y = F.topk(x, 3)
        np.testing.assert_equal(y[0].shape.numpy(), np.array([3,]))
        return y

    for i in range(3):
        f(x)


def test_trace_warp_perspective():
    inp_shape = (1, 1, 4, 4)
    x = tensor(np.arange(16, dtype=np.float32).reshape(inp_shape))
    M_shape = (1, 3, 3)
    M = tensor(
        np.array(
            [[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32
        ).reshape(M_shape)
    )

    @trace(symbolic=True)
    def f(x, M):
        out = F.warp_perspective(x, M, (2, 2))
        np.testing.assert_equal(out.shape.numpy(), np.array([1, 1, 2, 2]))
        return out

    for i in range(1):
        f(x, M)
312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343


def test_raise_on_trace():
    step_count = 0
    catch_count = 0
    bad_step = 10

    class CatchMe(Exception):
        pass

    a = tensor([1, 2, 3, 4])
    b = tensor([5, 6, 7, 8])
    c = tensor([9, 0, 1, 2])

    @trace
    def add_abc(a, b, c):
        ps = a + b
        result = ps + c
        if step_count == bad_step:
            raise CatchMe("catch me")
        return result

    for i in range(100):
        try:
            d = add_abc(a, b, c)
        except CatchMe as e:
            catch_count += 1
        else:
            np.testing.assert_equal(d.numpy(), (a + b + c).numpy())
        step_count += 1

    assert catch_count == 1
344 345 346 347 348 349 350 351 352 353


def test_trace_broadcast():
    for symbolic in [False, True]:
        x1 = tensor(np.random.randn(3, 1, 1))
        x2 = tensor(np.random.randn(1, 4, 1))
        x3 = tensor(np.random.randn(1, 1, 5))

        @trace(symbolic=symbolic, capture_as_const=True)
        def f(x):
354
            y = F.broadcast_to(x, (3, 4, 5))
355 356 357 358 359
            return y

        f(x1)
        f(x2)
        f(x3)
360 361 362 363 364 365 366 367 368 369 370 371 372 373


def test_trace_nms():
    def make_inputs(n):
        boxes = np.zeros((n, 4))
        boxes[:, :2] = np.random.rand(n, 2) * 100
        boxes[:, 2:] = np.random.rand(n, 2) * 100 + 100

        scores = np.random.rand(n)

        return tensor(boxes), tensor(scores)

    @trace(symbolic=False)
    def f(boxes, scores):
374
        # with tracing, max_output must be specified
375
        results = F.nn.nms(boxes, scores=scores, iou_thresh=0.5, max_output=20)
376
        # without tracing, max output can be inferred inside nms
377 378 379 380 381 382 383
        with exclude_from_trace():
            _ = F.nn.nms(boxes, scores=scores, iou_thresh=0.5)
        return results

    f(*make_inputs(10))
    f(*make_inputs(20))
    f(*make_inputs(30))
384 385 386 387 388 389 390 391 392 393 394 395 396 397


def test_trace_valid_broadcast():
    x1 = tensor(np.random.randn(1, 1))
    x2 = tensor(np.random.randn(1, 2))
    shape = (tensor([2]), tensor([2]))

    @trace(symbolic=False)
    def f(x, shape):
        y = F.broadcast_to(x, shape)
        return y

    f(x1, shape)
    f(x2, shape)
398 399 400 401 402 403 404 405 406 407 408 409


def test_clip():
    x = tensor(np.random.randn(10, 10))

    @trace(symbolic=True)
    def f(x, lower, upper):
        y = F.clip(x, lower, upper)
        return y

    for i in range(3):
        f(x, tensor([0]), tensor([1]))
410 411 412 413 414 415 416 417 418 419 420 421 422


# test returning noncontiguous tensor from trace
def test_slice():
    @trace
    def f(x):
        return x[:, 1::2]

    x = F.arange(8).reshape(2, 4)
    f(x)
    y = f(x)
    np.testing.assert_array_equal(y.numpy(), x.numpy()[:, 1::2])
    y + y
423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442


def test_random():
    def run_test(op):
        for symbolic_shape in [True, False]:

            @trace(symbolic=True, symbolic_shape=symbolic_shape)
            def f():
                out = op(size=[10, 10])
                out_shape = out.shape
                assert out_shape is not None
                if not isinstance(out_shape, tuple):
                    assert out.shape.numpy() is not None
                return out

            for _ in range(3):
                f()

    run_test(uniform)
    run_test(normal)