tests/python/unittest/test_deferred_compute.py - master - githubmirror/incubator-mxnet

tests/python/unittest/test_deferred_compute.py
570 строк16 KB
Dick Carter
[WIP] [BUGFIX] Fix flakey TemporaryDirectory() cleanup on Windows (#21107)
01 авг 2022, 20:21
01 авг 2022, 20:21dedb8c9
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import functools
import operator

import numpy as np

import mxnet as mx
import mxnet._deferred_compute as dc
from mxnet.base import MXNetError
from mxnet.util import TemporaryDirectory
import pytest


def _all_same(arrays1, arrays2, message=''):
    same = all(np.array_equal(a1, a2) for a1, a2 in zip(arrays1, arrays2))
    if not same:
        raise AssertionError('Arrays not equal ({}):\n{}\n\n{}'.format(message, arrays1, arrays2))


def _assert_dc(setup, compute, mode='all', setup_is_deterministic=True, numpy=True):
    """Compare results of deferred compute and normal imperative mode.

    Parameters
    ----------
    setup : callable
        Setup function computing inputs for compute function. Always called
        outside of deferred compute.
    compute : callable
        Compute function. We compare the output between normal computation and
        deferred compute.
    mode : {'all', 'symbolic', 'imperative', 'imperativewithnondccompute'}
        Compare deferred compute outputs triggered via imperative computation
        (eg. asnumpy() conversion) or obtained from the exported symbol or
        both.
    setup_is_deterministic : bool
        If True, setup function may be called multiple times. If False, will
        only be called once.
    numpy : bool
        If True, use mx.np. Otherwise mx.nd.

    """
    try:
        nd = mx.np if numpy else mx.nd
        if numpy:
            mx.npx.set_np()

        xs = setup(nd=nd)
        ys = compute(*xs, nd=nd)

        ys_np = [y.asnumpy() for y in ys]

        if setup_is_deterministic:
            xs = setup(nd=nd)

        xs_names = list(map(str, range(len(xs))))
        symbol_inputs = [
            mx.symbol.var(name).as_np_ndarray()
            if numpy else mx.symbol.var(name)
            for arg, name in zip(xs, xs_names)
        ]
        dc.set_variable(xs, symbol_inputs)
        with dc.context():
            ys_dc = compute(*xs, nd=nd)

        assert mode in ('all', 'symbolic', 'imperative', 'imperativewithnondccompute')
        if mode in ('all', 'imperativewithnondccompute'):
            ys_dc_np = [(y + 0).asnumpy() for y in ys_dc]
            _all_same(ys_np, ys_dc_np)

        if mode in ('all', 'imperative'):
            ys_dc_np = [y.asnumpy() for y in ys_dc]
            _all_same(ys_np, ys_dc_np)

        if mode in ('all', 'symbolic'):
            sym = dc.get_symbol(ys_dc, sym_cls=mx.sym.np._Symbol if numpy else mx.sym.Symbol)

            if setup_is_deterministic:
                xs = setup(nd=nd)

            args = {name: x for name, x in zip(xs_names, xs)}
            ys_sym = sym._bind(mx.device.current_device(), args=args).forward()

            ys_sym_np = [y.asnumpy() for y in ys_sym]
            _all_same(ys_np, ys_sym_np)
    finally:
        if numpy:
            mx.npx.reset_np()


def _all_assert_dc(setup, compute, setup_is_deterministic=True, numpy=(False, True)):
    for mode in ('all', 'symbolic', 'imperative', 'imperativewithnondccompute'):
        for numpy_ in numpy:
            _assert_dc(setup, compute, mode=mode, setup_is_deterministic=True, numpy=numpy_)


###############################################################################
# Test cases without inputs
###############################################################################
def _dc_empty_setup(*, nd):
    return []


def test_dc_no_inputs_single_output():
    def f(*, nd):
        a = nd.arange(10)
        b = a + nd.arange(a.shape[0])
        c = b - 1
        return [c]

    _all_assert_dc(_dc_empty_setup, f)


def test_dc_no_inputs_reshape():
    def f(*, nd):
        a = nd.arange(10)
        b = a + nd.arange(a.shape[0])
        c = b.reshape((5, 2))
        d = b.reshape((2, 5))
        e = (c.reshape((-1, )) + d.reshape((-1, ))) / 2
        return [c + 1, d + 1, e]

    _all_assert_dc(_dc_empty_setup, f)


def test_dc_no_inputs_slice():
    def f(*, nd):
        a = nd.arange(10)
        b = a[:5]
        if nd is mx.nd:
            c = nd.concat(b, b, dim=0)
        else:
            c = nd.concatenate([b, b], axis=0)
        return [c + a]

    _all_assert_dc(_dc_empty_setup, f)


def test_dc_no_inputs_subset_of_output():
    def f(*, nd):
        a = nd.arange(10)
        if nd is mx.nd:
            b, c = mx.nd.split(a, 2, axis=0)
        else:
            b, c = mx.np.array_split(a, 2, axis=0)
        return [b]

    _all_assert_dc(_dc_empty_setup, f)


def test_dc_numpy_tril():
    def f(a, *, nd):
        assert nd is mx.np
        a = nd.ones((2, 2))
        b = nd.tril(a, 1)
        c = nd.tril(a, -1)
        return [b, c]

    for mode in ('all', 'symbolic', 'imperative', 'imperativewithnondccompute'):
        _assert_dc(_dc_simple_setup, f, mode=mode)


###############################################################################
# Test cases with inputs
###############################################################################
def _dc_simple_setup(shape=(10, ), *, nd):
    n = functools.reduce(operator.mul, shape, 1)
    return [nd.arange(n).reshape(shape)]


def test_dc_single_output():
    def f(a, *, nd):
        b = a + nd.arange(a.shape[0])
        c = b - 1
        return [c]

    _all_assert_dc(_dc_simple_setup, f)


def test_dc_reshape():
    def f(a, *, nd):
        b = a + nd.arange(a.shape[0])
        c = b.reshape((5, 2))
        d = b.reshape((2, 5))
        e = (c.reshape((-1, )) + d.reshape((-1, ))) / 2
        return [c + 1, d + 1, e]

    _all_assert_dc(_dc_simple_setup, f)


def test_dc_slice():
    def f(a, *, nd):
        b = a[:5]
        if nd is mx.nd:
            c = nd.concat(b, b, dim=0)
        else:
            c = nd.concatenate([b, b], axis=0)
        return [c + a]

    _all_assert_dc(_dc_simple_setup, f)


def test_dc_subset_of_output():
    def f(a, *, nd):
        if nd is mx.nd:
            b, c = mx.nd.split(a, 2, axis=0)
        else:
            b, c = mx.np.array_split(a, 2, axis=0)
        return [b]

    _all_assert_dc(_dc_simple_setup, f)


def test_dc_inplace_error():
    def f(a, *, nd):
        a[:5] = 0
        b = a + 1
        return [a, b]

    with pytest.raises(MXNetError):
    # TODO(leezu): Should raise NotImplementedError https://github.com/apache/incubator-mxnet/issues/17522
        _all_assert_dc(_dc_simple_setup, f)


###############################################################################
# Special cases
###############################################################################
def test_dc_input_part_of_output():
    a = mx.np.arange(10)
    dc.set_variable(a, mx.sym.var('a'))
    with dc.context():
        b = a + 1
    dc.get_symbol([a, b])


def test_dc_get_symbol_called_twice():
    a = mx.np.arange(10)
    dc.set_variable(a, mx.sym.var('a'))
    with dc.context():
        b = a + 1
    sym1 = dc.get_symbol(b)
    sym2 = dc.get_symbol(b)
    assert sym1.list_inputs() == ['a']
    assert sym2.list_inputs() == ['a']


def test_dc_set_variable_called_twice_error():
    a = mx.np.arange(10)
    dc.set_variable(a, mx.sym.var('a'))
    with pytest.raises(MXNetError):
    # TODO(leezu): Should raise ValueError https://github.com/apache/incubator-mxnet/issues/17522
        dc.set_variable(a, mx.sym.var('b'))


def test_dc_no_inputs_context_switch():
    def f(*, nd):
        a = nd.arange(10)
        if nd is mx.nd:
            b = a.as_in_context(mx.cpu(1))
            c = (b - 1).as_in_context(mx.device.current_device())
        else:
            b = a.to_device(mx.cpu(1))
            c = (b - 1).to_device(mx.device.current_device())
        return [c]

    _assert_dc(_dc_empty_setup, f)


def test_dc_context_switch():
    def f(a, *, nd):
        if nd is mx.nd:
            b = a.as_in_context(mx.cpu(1))
            c = (b - 1).as_in_context(mx.device.current_device())
        else:
            b = a.to_device(mx.cpu(1))
            c = (b - 1).to_device(mx.device.current_device())
        return [c]

    _assert_dc(_dc_simple_setup, f)


def test_dc_astype():
    def f(a, *, nd):
        a = a.astype(np.int32)
        b = nd.zeros_like(a)
        return [a + b]

    _assert_dc(_dc_simple_setup, f)


def test_dc_dynamic_shape():
    def f(a, *, nd):
        return [mx.nd.np.flatnonzero(a)]

    for mode in ('imperative', 'imperativewithnondccompute', 'symbolic', 'all'):
        _assert_dc(_dc_simple_setup, f, mode=mode, numpy=True)


###############################################################################
# Indexing specific tests
###############################################################################
def test_dc_integer_indexing():
    def f(a, *, nd):
        return [a[1] + 1]

    _all_assert_dc(_dc_simple_setup, f)


def test_dc_slice_indexing():
    def f(a, *, nd):
        b = a.reshape((5, 2))
        return [b[:2, 1] + 1]

    _all_assert_dc(_dc_simple_setup, f)


def test_dc_tuple_indexing():
    def f(a, *, nd):
        b = a.reshape((5, 2))
        return [b[(1, 1)] + 1]

    _all_assert_dc(_dc_simple_setup, f)


def test_dc_simple_boolean_indexing():
    def setup(*, nd):
        assert nd is mx.np
        x = mx.np.array([[0, 1], [1, 1], [2, 2]])
        return [x, x < 2]

    def f(a, idx, *, nd):
        assert nd is mx.np
        return [a[idx].reshape((2, 2))]


def test_dc_list_indexing_error():
    def f(a, *, nd):
        assert nd is mx.np
        return [a[[1, 2, 3]]]

    for mode in ('all', 'symbolic', 'imperative', 'imperativewithnondccompute'):
        with pytest.raises(TypeError):
            _assert_dc(_dc_simple_setup, f, mode=mode)


def test_dc_numpy_indexing_error():
    def f(a, *, nd):
        assert nd is mx.np
        return [a[np.array([1, 2, 3])]]

    for mode in ('all', 'symbolic', 'imperative', 'imperativewithnondccompute'):
        with pytest.raises(TypeError):
            _assert_dc(_dc_simple_setup, f, mode=mode)


###############################################################################
# Gluon
###############################################################################
def _assert_dc_gluon(setup, net, setup_is_deterministic=True, numpy=True, autograd=True, device=None):
    """Compare results of deferred compute and normal imperative mode.

    Parameters
    ----------
    setup : callable
        Setup function computing inputs for compute function. Always called
        outside of deferred compute.
    net : Block
    setup_is_deterministic : bool
        If True, setup function may be called multiple times. If False, will
        only be called once.
    numpy : bool
        If True, use mx.np. Otherwise mx.nd.
    autograd : bool
        Wrap in autograd

    """

    nd = mx.np if numpy else mx.nd

    if device is None:
        device = mx.device.current_device()
    with device:
        xs = setup(nd=nd)

    ys = net(*xs)
    ys_np = [y.asnumpy() for y in ys]

    net.hybridize()
    if setup_is_deterministic:
        xs = setup(nd=nd)

    if autograd:
        with mx.autograd.record():
            ys_hybrid = net(*xs)
        mx.autograd.backward(ys_hybrid)
        [p.grad() for p in net.collect_params().values()]
    else:
        ys_hybrid = net(*xs)

    assert all(
        isinstance(y, mx.numpy.ndarray) if numpy else isinstance(y, mx.ndarray.ndarray.NDArray)
        for y in ys_hybrid)

    ys_hybrid_np = [y.asnumpy() for y in ys_hybrid]

    _all_same(ys_np, ys_hybrid_np)

    with TemporaryDirectory() as root:
        with mx.util.np_shape(True), mx.util.np_array(True):
            net.export(root)

def _dc_gluon_simple_setup(shape=(8, 10), *, nd):
    return [nd.ones(shape=shape, device=mx.device.current_device())]


def test_dc_hybridblock():
    class MyBlock(mx.gluon.HybridBlock):
        def __init__(self):
            super().__init__()
            self.dense = mx.gluon.nn.Dense(units=10, in_units=10)
            self.weight = mx.gluon.Parameter('weight', shape=(10, ))

        def forward(self, x):
            assert x.shape[1] == 10  # due to in_units=10 above
            return self.dense(x) + self.weight.data(x.context)

    if mx.device.current_device() == mx.cpu(0):  # CPU tests
        devices = [mx.cpu(0), mx.cpu(1)]
    else:  # Use default device, GPU tests
        devices = [mx.device.current_device()]
    for device in devices:
        net = MyBlock()
        net.initialize(device=devices)
        _assert_dc_gluon(_dc_gluon_simple_setup, net, numpy=True, device=device)


def test_dc_hybridblock_wrapped():
    @mx.util.use_np
    class MyBlock(mx.gluon.HybridBlock):
        def __init__(self):
            super().__init__()
            self.dense = mx.gluon.nn.Dense(units=10, in_units=10)
            self.weight = mx.gluon.Parameter('weight', shape=(10, ))

        def forward(self, x):
            assert x.shape[1] == 10  # due to in_units=10 above
            return self.dense(x) + self.weight.data(x.context)

    net = MyBlock()
    net.initialize()
    _assert_dc_gluon(_dc_gluon_simple_setup, net, numpy=True)


def test_dc_hybridblock_deferred_init_no_infer_shape_error():
    class MyBlock(mx.gluon.HybridBlock):
        def __init__(self):
            super().__init__()
            self.dense = mx.gluon.nn.Dense(units=10)
            self.weight = mx.gluon.Parameter('weight', allow_deferred_init=True)

        def forward(self, x):
            return self.dense(x) + self.weight.data(x.context)

    net = MyBlock()
    net.initialize()
    data = mx.np.ones(shape=(8, 10), device=mx.device.current_device())
    with pytest.raises(RuntimeError):
        net(data)


def test_dc_hybridblock_deferred_init():
    class MyBlock(mx.gluon.HybridBlock):
        def __init__(self):
            super().__init__()
            self.dense = mx.gluon.nn.Dense(units=10)
            self.weight = mx.gluon.Parameter('weight', allow_deferred_init=True)

        def infer_shape(self, x):
            self.weight.shape = (x.shape[1], )

        def forward(self, x):
            return self.dense(x) + self.weight.data(x.device)

    net = MyBlock()
    net.initialize()
    _assert_dc_gluon(_dc_gluon_simple_setup, net, numpy=True)


def test_dc_hybridblock_dynamic_shape():
    class MyBlock(mx.gluon.HybridBlock):
        def __init__(self):
            super().__init__()
            self.dense = mx.gluon.nn.Dense(units=10)

        def forward(self, x, idx):
            return x[idx].reshape((2, 2)), mx.np.flatnonzero(self.dense(x))

    def setup(*, nd):
        assert nd is mx.np
        x = mx.np.array([[0, 1], [1, 1], [2, 2]])
        return [x, x < 2]

    with mx.util.np_shape(True), mx.util.np_array(True):
        net = MyBlock()
        net.initialize()
        _assert_dc_gluon(setup, net, numpy=True)

def test_dc_hybridblock_graph_partition():
    class MyBlock(mx.gluon.HybridBlock):
        def __init__(self):
            super().__init__()
            self.dense = mx.gluon.nn.Dense(units=4)

        def forward(self, x, idx):
            mask = mx.nd.np._internal.boolean_mask(self.dense(x), idx)
            return mx.np.sum(mask)

    def setup(*, nd):
        x = mx.np.array([[0, 1], [2, 3], [4, 5], [6, 7]])
        idx = mx.np.array([1, 1, 1, 1])
        return [x, idx]

    net = MyBlock()
    net.initialize()
    _assert_dc_gluon(setup, net, numpy=True, autograd=False)


def test_indexing_shape_change():
    class ConcatBlock(mx.gluon.nn.HybridBlock):
        def forward(self, inputs):
            return mx.np.concatenate([
                inputs,
                mx.np.pad(inputs[:,1:], ((0,0), (0,1))),
            ])

    net = ConcatBlock()
    net.hybridize()
    net(mx.np.random.uniform(size=(8, 16)))
    net(mx.np.random.uniform(size=(8, 8)))


def test_indexing_empty_shape():
    @mx.util.use_np
    class TestModel(mx.gluon.HybridBlock):
        def forward(self, x):
            return x[0]

    net = TestModel()
    net.hybridize()
    try:
        mx.npx.set_np()
        net(mx.np.zeros((2, 2, 4, 0, 128)))
        net(mx.np.zeros((2, 2, 4, 2, 128)))  # test indexing after input shape change
    finally:
        mx.npx.reset_np()