| # 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 os |
| import sys |
| import ctypes |
| import mxnet as mx |
| from mxnet.base import SymbolHandle, check_call, _LIB, mx_uint, c_str_array, c_str, mx_real_t |
| from mxnet.symbol import Symbol |
| import numpy as onp |
| from mxnet.test_utils import assert_almost_equal, environment |
| from mxnet import gluon |
| from mxnet.gluon import nn |
| from mxnet import nd, npx |
| import pytest |
| import tempfile |
| |
| mx.npx.reset_np() |
| |
| def network_structure_1(): |
| data1 = mx.sym.var('data1', shape=(2, 3, 10, 10)) |
| data2 = mx.sym.var('data2') |
| conv1 = mx.sym.Convolution(data=data1, weight=data2, no_bias=True, kernel=(2, 2), num_filter=1) |
| conv2 = mx.sym.Convolution(data=data2, no_bias=True, kernel=(1, 1), num_filter=1) |
| out = mx.sym.Group([conv1, conv2]) |
| return (out, ['data1'], [(2, 3, 10, 10)]) |
| |
| def network_structure_2(): |
| # this tests whether the partitioning algorithm can deal with cycles |
| data = mx.sym.var('data', shape=(2, 3, 10, 10)) |
| ret = mx.sym.exp(data) |
| ret1 = mx.sym.cos(ret) |
| ret2 = mx.sym.sin(ret) |
| ret = ret1 + ret2 |
| return (ret, ['data'], [(2, 3, 10, 10)]) |
| |
| def network_structure_3(): |
| # this tests whether the partitioned sym can distinguish in_args and aux_states |
| data = mx.sym.var('data', shape=(2, 3, 10, 10)) |
| ret = mx.sym.exp(data) |
| ret1 = mx.sym.cos(ret) |
| ret2 = mx.sym.sin(ret) |
| ret = ret1 + ret2 |
| ret = mx.sym.BatchNorm(ret) |
| ret = mx.sym.BatchNorm(ret) |
| # Return the same and shape of 'data' and auxiliary states |
| return (ret, ['data'] + ret.list_auxiliary_states(), [(2, 3, 10, 10), (3,), (3,), (3,), (3,)]) |
| |
| def network_structure_4(): |
| # the last op has multiple duplicate outputs |
| data = mx.sym.var('data', shape=(2, 3, 10, 10)) |
| ret = mx.sym.exp(data) |
| ret = mx.sym.Group([ret, ret, ret]) |
| return (ret, ['data'], [(2, 3, 10, 10)]) |
| |
| def network_structure_5(): |
| # the subgraph has two duplicate input entries |
| data = mx.sym.var('data', shape=(2, 3, 10, 10)) |
| ret = data + data |
| return (ret, ['data'], [(2, 3, 10, 10)]) |
| |
| def network_structure_6(): |
| data1 = mx.sym.Variable('data1', shape=(3, 3, 10, 10), dtype=onp.float32) |
| data2 = mx.sym.Variable('data2', shape=(1, 0, 2, 2)) |
| data3 = mx.sym.sin(data2) |
| conv = mx.sym.Convolution(data=data1, weight=data3, kernel=(2, 2), num_filter=1) |
| return (conv, ['data1'], [(3, 3, 10, 10)]) |
| |
| def network_structure_7(): |
| # in this graph, the subgraph node and the other two external nodes form a cycle |
| data = mx.sym.Variable('data', shape=(1,)) |
| ret1 = mx.sym.sin(data) |
| ret2 = mx.sym.cos(ret1) |
| for _ in range(5): |
| ret2 = mx.sym.cos(ret2) |
| ret = ret1 + ret2 |
| return (ret, ['data'], [(1,)]) |
| |
| def network_structure_8(): |
| # in this graph, two nodes in the subgraph consume the same input, and |
| # and two nodes outside the subgraph consume a single output from the subgraph |
| data = mx.sym.Variable('data', shape=(1,)) |
| sin1 = mx.sym.sin(data) |
| sin2 = mx.sym.sin(data) |
| plus = sin1 + sin2 |
| ret1 = mx.sym.cos(plus) |
| ret2 = mx.sym.cos(plus) |
| ret = ret1 - ret2 |
| return (ret, ['data'], [(1,)]) |
| |
| def get_graphs(): |
| return [ |
| (network_structure_1(), ['Convolution']), |
| (network_structure_2(), ['exp', 'sin', '_Plus', 'elemwise_add', '_plus']), |
| (network_structure_2(), ['exp', 'cos', '_Plus', 'elemwise_add', '_plus']), |
| (network_structure_3(), ['exp', 'sin', '_Plus', 'elemwise_add', '_plus']), |
| (network_structure_3(), ['exp', 'cos', '_Plus', 'elemwise_add', '_plus']), |
| (network_structure_3(), ['exp', 'sin', '_Plus', 'elemwise_add', '_plus', 'BatchNorm']), |
| (network_structure_3(), ['exp', 'cos', '_Plus', 'elemwise_add', '_plus', 'BatchNorm']), |
| (network_structure_3(), ['exp', 'BatchNorm']), |
| (network_structure_3(), ['BatchNorm']), |
| (network_structure_4(), ['exp']), |
| (network_structure_5(), ['_plus', '_Plus', 'elemwise_add']), |
| (network_structure_6(), []), |
| (network_structure_6(), [mx.sym.sin.__name__]), |
| (network_structure_6(), [mx.sym.Convolution.__name__]), |
| (network_structure_6(), [mx.sym.sin.__name__, mx.sym.Convolution.__name__]), |
| (network_structure_7(), ['sin', 'elemwise_add', '_plus', '_Plus']), |
| (network_structure_8(), ['sin', 'elemwise_add']) |
| ] |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe1(sym, subgraph_backend, op_names): |
| """Use the partitioned sym to _simple_bind an executor and compare the outputs |
| with those of the original executor""" |
| sym, _, _ = sym |
| out = SymbolHandle() |
| check_call(_LIB.MXBuildSubgraphByOpNames(sym.handle, c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names), ctypes.byref(out))) |
| |
| partitioned_sym = Symbol(out) |
| assert partitioned_sym.list_inputs() == sym.list_inputs() |
| assert partitioned_sym.list_arguments() == sym.list_arguments() |
| assert partitioned_sym.list_auxiliary_states() == sym.list_auxiliary_states() |
| exe = sym._simple_bind(ctx=mx.current_context(), grad_req='null') |
| partitioned_exe = partitioned_sym._simple_bind(ctx=mx.current_context(), grad_req='null') |
| input_names = sym.list_inputs() |
| for name in input_names: |
| if name in exe.arg_dict: |
| exe.arg_dict[name][:] = mx.nd.random.uniform(shape=exe.arg_dict[name].shape) |
| partitioned_exe.arg_dict[name][:] = exe.arg_dict[name] |
| else: |
| assert name in exe.aux_dict |
| exe.aux_dict[name][:] = mx.nd.random.uniform(shape=exe.aux_dict[name].shape) |
| partitioned_exe.aux_dict[name][:] = exe.aux_dict[name] |
| exe.forward() |
| partitioned_exe.forward() |
| assert len(exe.outputs) == len(partitioned_exe.outputs) |
| for i in range(len(exe.outputs)): |
| assert_almost_equal((exe.outputs[i] - partitioned_exe.outputs[i]).abs().sum().asnumpy(), |
| onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe2(sym, subgraph_backend, op_names): |
| def get_executor(sym, subgraph_backend=None, op_names=None, original_exec=None): |
| exe = sym._simple_bind(ctx=mx.current_context(), grad_req='null') |
| input_names = sym.list_inputs() |
| for name in input_names: |
| if name in exe.arg_dict: |
| exe.arg_dict[name][:] = mx.nd.random.uniform(shape=exe.arg_dict[name].shape)\ |
| if original_exec is None else original_exec.arg_dict[name] |
| else: |
| assert name in exe.aux_dict |
| exe.aux_dict[name][:] = mx.nd.random.uniform(shape=exe.aux_dict[name].shape)\ |
| if original_exec is None else original_exec.aux_dict[name] |
| exe.forward() |
| return exe |
| sym, _, _ = sym |
| original_exec = get_executor(sym) |
| check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| partitioned_exec = get_executor(sym, subgraph_backend, op_names, original_exec) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend))) |
| outputs1 = original_exec.outputs |
| outputs2 = partitioned_exec.outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe3(sym, subgraph_backend, op_names): |
| """Use the partitioned sym to bind an executor and compare the outputs |
| with those of the original executor""" |
| sym, _, _ = sym |
| out = SymbolHandle() |
| check_call(_LIB.MXBuildSubgraphByOpNames(sym.handle, c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names), ctypes.byref(out))) |
| |
| partitioned_sym = Symbol(out) |
| input_names = sym.list_inputs() |
| arg_names = sym.list_arguments() |
| aux_names = sym.list_auxiliary_states() |
| assert partitioned_sym.list_inputs() == input_names |
| assert partitioned_sym.list_arguments() == arg_names |
| assert partitioned_sym.list_auxiliary_states() == aux_names |
| arg_shapes, _, aux_shapes = sym.infer_shape() |
| arg_array = [mx.nd.random.uniform(shape=shape) for shape in arg_shapes] |
| aux_array = [mx.nd.random.uniform(shape=shape) for shape in aux_shapes] |
| exe = sym._bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null') |
| partitioned_exe = partitioned_sym._bind(ctx=mx.current_context(), args=arg_array, |
| aux_states=aux_array, grad_req='null') |
| exe.forward() |
| partitioned_exe.forward() |
| assert len(exe.outputs) == len(partitioned_exe.outputs) |
| for i in range(len(exe.outputs)): |
| assert_almost_equal((exe.outputs[i] - partitioned_exe.outputs[i]).abs().sum().asnumpy(), |
| onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe4(sym, subgraph_backend, op_names): |
| def get_executor(sym, subgraph_backend=None, op_names=None, original_exec=None): |
| arg_shapes, _, aux_shapes = sym.infer_shape() |
| if subgraph_backend is None: |
| arg_array = [mx.nd.random.uniform(shape=shape) for shape in arg_shapes] |
| aux_array = [mx.nd.random.uniform(shape=shape) for shape in aux_shapes] |
| else: |
| arg_array = None |
| aux_array = None |
| exe = sym._bind(ctx=mx.current_context(), |
| args=arg_array if subgraph_backend is None else original_exec.arg_arrays, |
| aux_states=aux_array if subgraph_backend is None else original_exec.aux_arrays, |
| grad_req='null') |
| exe.forward() |
| return exe |
| |
| sym, _, _ = sym |
| original_exec = get_executor(sym) |
| check_call(_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| partitioned_exec = get_executor(sym, subgraph_backend, op_names, original_exec) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend))) |
| outputs1 = original_exec.outputs |
| outputs2 = partitioned_exec.outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| def set_random_inputs(exe1, input_names): |
| """Sets random values to exe1's args and auxs""" |
| for name in input_names: |
| if name in exe1.arg_dict: |
| exe1.arg_dict[name][:] = mx.nd.random.uniform(shape=exe1.arg_dict[name].shape) |
| else: |
| assert name in exe1.aux_dict |
| exe1.aux_dict[name][:] = mx.nd.random.uniform(shape=exe1.aux_dict[name].shape) |
| |
| def copy_inputs_between_executors(exe1, exe2, input_names): |
| """Copies values of args and auxs from exe1 to exe2""" |
| for name in input_names: |
| if name in exe2.arg_dict: |
| exe2.arg_dict[name][:] = exe1.arg_dict[name] |
| else: |
| assert name in exe2.aux_dict |
| exe2.aux_dict[name][:] = exe1.aux_dict[name] |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe5(sym, subgraph_backend, op_names): |
| """Call optimize_for to trigger graph partitioning without infer shapes/types before, |
| then _simple_bind and compare results of the partitioned sym and the original sym.""" |
| # _simple_bind |
| sym, _, _ = sym |
| exe1 = sym._simple_bind(ctx=mx.current_context(), grad_req='null') |
| input_names = sym.list_inputs() |
| set_random_inputs(exe1, input_names) |
| exe1.forward() |
| |
| # partition before _simple_bind |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| part_sym = sym.optimize_for(subgraph_backend) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| exe2 = part_sym._simple_bind(ctx=mx.current_context(), grad_req='null') |
| copy_inputs_between_executors(exe1, exe2, input_names) |
| exe2.forward() |
| |
| # compare outputs |
| outputs1 = exe1.outputs |
| outputs2 = exe2.outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe6(sym, subgraph_backend, op_names): |
| """Call optimize_for to trigger graph partitioning with shapes/types, then _simple_bind |
| and compare results of the partitioned sym and the original sym.""" |
| # _simple_bind |
| sym, _, _ = sym |
| exe1 = sym._simple_bind(ctx=mx.current_context(), grad_req='null') |
| input_names = sym.list_inputs() |
| set_random_inputs(exe1, input_names) |
| exe1.forward() |
| |
| # infer shape/type before partition before _simple_bind |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| part_sym = sym.optimize_for(subgraph_backend, exe1.arg_dict, exe1.aux_dict) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| exe2 = part_sym._simple_bind(ctx=mx.current_context(), grad_req='null') |
| copy_inputs_between_executors(exe1, exe2, input_names) |
| exe2.forward() |
| |
| # compare outputs |
| outputs1 = exe1.outputs |
| outputs2 = exe2.outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe7(sym, subgraph_backend, op_names): |
| """Call optimize_for to trigger graph partitioning without infer shapes/types before, |
| then bind and compare results of the partitioned sym and the original sym.""" |
| # bind |
| sym, _, _ = sym |
| arg_shapes, _, aux_shapes = sym.infer_shape() |
| arg_array = [mx.nd.random.uniform(shape=shape) for shape in arg_shapes] |
| aux_array = [mx.nd.random.uniform(shape=shape) for shape in aux_shapes] |
| exe1 = sym._bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null') |
| exe1.forward() |
| |
| # partition before bind |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| part_sym = sym.optimize_for(subgraph_backend) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| exe2 = part_sym._bind(ctx=mx.current_context(), args=arg_array, aux_states=aux_array, grad_req='null') |
| exe2.forward() |
| |
| # compare outputs |
| outputs1 = exe1.outputs |
| outputs2 = exe2.outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe8(sym, subgraph_backend, op_names): |
| """Call optimize_for to infer shapes, types and dtypes followed by graph partitioning, |
| then bind and compare results of the partitioned sym and the original sym.""" |
| # bind |
| sym, _, _ = sym |
| arg_shapes, _, aux_shapes = sym.infer_shape() |
| arg_names = sym.list_arguments() |
| aux_names = sym.list_auxiliary_states() |
| arg_dict = {name:mx.nd.random.uniform(shape=shape) for name,shape in zip(arg_names,arg_shapes)} |
| aux_dict = {name:mx.nd.random.uniform(shape=shape) for name,shape in zip(aux_names,aux_shapes)} |
| exe1 = sym._bind(ctx=mx.current_context(), args=arg_dict, aux_states=aux_dict, grad_req='null') |
| exe1.forward() |
| |
| # infer shape/type before partition before bind |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| part_sym = sym.optimize_for(subgraph_backend, arg_dict, aux_dict) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| exe2 = part_sym._bind(ctx=mx.current_context(), args=arg_dict, aux_states=aux_dict, grad_req='null') |
| exe2.forward() |
| |
| # compare outputs |
| outputs1 = exe1.outputs |
| outputs2 = exe2.outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_exe9(sym, subgraph_backend, op_names): |
| """Call optimize_for to infer shapes, types and dtypes followed by graph partitioning and |
| dedup subgraph, then bind and compare results of the partitioned sym and the original sym.""" |
| # bind |
| sym, _, _ = sym |
| arg_shapes, _, aux_shapes = sym.infer_shape() |
| arg_names = sym.list_arguments() |
| aux_names = sym.list_auxiliary_states() |
| arg_dict = {name:mx.nd.random.uniform(shape=shape) for name,shape in zip(arg_names,arg_shapes)} |
| aux_dict = {name:mx.nd.random.uniform(shape=shape) for name,shape in zip(aux_names,aux_shapes)} |
| exe1 = sym._bind(ctx=mx.current_context(), args=arg_dict, aux_states=aux_dict, grad_req='null') |
| exe1.forward() |
| |
| # infer shape/type before partition before bind |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| part_sym = sym.optimize_for(subgraph_backend, arg_dict, aux_dict, dedup_subgraph=True) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| exe2 = part_sym._bind(ctx=mx.current_context(), args=arg_dict, aux_states=aux_dict, grad_req='null') |
| exe2.forward() |
| |
| # compare outputs |
| outputs1 = exe1.outputs |
| outputs2 = exe2.outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| @pytest.mark.parametrize('subgraph_backend', ['default', 'default_v2']) |
| @pytest.mark.parametrize('sym,op_names', get_graphs()) |
| @pytest.mark.skipif(sys.platform == "win32", reason='https://github.com/apache/incubator-mxnet/issues/19915') |
| def test_subgraph_backend_gluon(sym, subgraph_backend, op_names, tmp_path): |
| """Call hybridize() to partition the graph, and then compare results of the partitioned |
| sym and the original sym. Here do an inference before hybridizing with the subgraph_backend |
| which means we'll pass shapes/types""" |
| # create Gluon block for given symbol |
| inputs = [mx.sym.var(i, dtype=mx_real_t) for i in sym[1]] |
| sym_block = nn.SymbolBlock(sym[0], inputs) |
| sym_block.initialize(ctx=mx.current_context()) |
| x = [mx.nd.random.uniform(shape=s,ctx=mx.current_context()) for s in sym[2]] |
| # hybridize and export to get baseline |
| sym_block.hybridize() |
| outputs1 = sym_block(*x) |
| |
| sym_filename, params_filename = sym_block.export(str(tmp_path / 'sym-block')) |
| |
| # load model and partition |
| sym_block = nn.SymbolBlock.imports(sym_filename, sym[1], params_filename, |
| ctx=mx.current_context()) |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| sym_block.optimize_for(*x, backend=subgraph_backend) |
| outputs2 = sym_block(*x) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| # compare outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal((outputs1[i] - outputs2[i]).abs().sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| # Test Gluon HybridBlocks for graph partitioning a network created by HybridSequential. |
| @mx.util.use_np |
| @pytest.mark.serial |
| def test_subgraph_backend_gluon_ext1(tmpdir): |
| def get_net(): |
| net = nn.HybridSequential() # Here we use the class HybridSequential. |
| net.add(nn.Dense(256, activation='relu'), |
| nn.Dense(128, activation='relu'), |
| nn.Dense(2)) |
| return net |
| |
| # regular inference |
| x = mx.np.random.normal(size=(1, 512),ctx=mx.current_context()) |
| net = get_net() |
| net.initialize(ctx=mx.current_context()) |
| outputs1 = net(x) |
| param_path = os.path.join(str(tmpdir), 'test_subgraph_backend_gluon_ext1.params') |
| net.save_parameters(param_path) |
| |
| # after partitioning |
| net = get_net() |
| net.load_parameters(param_path,ctx=mx.current_context()) |
| subgraph_backend = 'default' |
| op_names = ['FullyConnected'] |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| net.optimize_for(x, backend = subgraph_backend) |
| outputs2 = net(x) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| # compare outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal(mx.np.abs((outputs1[i] - outputs2[i])).sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| # Test Gluon HybridBlocks for graph partitioning a network created by HybridBlock. |
| @mx.util.use_np |
| @pytest.mark.serial |
| def test_subgraph_backend_gluon_ext2(tmpdir): |
| class Net(gluon.HybridBlock): |
| def __init__(self, **kwargs): |
| super(Net, self).__init__(**kwargs) |
| self.fc1 = nn.Dense(256) |
| self.fc2 = nn.Dense(128) |
| self.fc3 = nn.Dense(2) |
| |
| def forward(self, x): |
| x = npx.relu(self.fc1(x)) |
| x = npx.relu(self.fc2(x)) |
| return self.fc3(x) |
| # regular inference |
| x = mx.np.random.normal(size=(1, 512),ctx=mx.current_context()) |
| net = Net() |
| net.initialize(ctx=mx.current_context()) |
| outputs1 = net(x) |
| param_path = os.path.join(str(tmpdir), 'test_subgraph_backend_gluon_ext2.params') |
| net.save_parameters(param_path) |
| |
| # after partitioning |
| net = Net() |
| net.load_parameters(param_path, ctx=mx.current_context()) |
| subgraph_backend = 'default' |
| op_names = ['FullyConnected'] |
| check_call(_LIB.MXSetSubgraphPropertyOpNamesV2(c_str(subgraph_backend), mx_uint(len(op_names)), |
| c_str_array(op_names))) |
| net.optimize_for(x, backend = subgraph_backend) |
| outputs2 = net(x) |
| check_call(_LIB.MXRemoveSubgraphPropertyOpNamesV2(c_str(subgraph_backend))) |
| |
| # compare outputs |
| assert len(outputs1) == len(outputs2) |
| for i in range(len(outputs1)): |
| assert_almost_equal(mx.np.abs(outputs1[i] - outputs2[i]).sum().asnumpy(), onp.zeros(shape=(1,))) |
| |
| |
| if __name__ == "__main__": |
| import datetime |
| tmpdir = datetime.datetime.now().strftime('mylogfile_%H_%M_%S_%f_%d_%m_%Y.log') |
| os.mkdir(tmpdir) |
| subgraph_backends = ['default', 'default_v2'] |
| graphs = get_graphs() |
| for subgraph_backend in subgraph_backends: |
| for sym,op_names in graphs: |
| test_subgraph_exe1(sym, subgraph_backend, op_names) |
| test_subgraph_exe2(sym, subgraph_backend, op_names) |
| test_subgraph_exe3(sym, subgraph_backend, op_names) |
| test_subgraph_exe4(sym, subgraph_backend, op_names) |
| test_subgraph_exe5(sym, subgraph_backend, op_names) |
| test_subgraph_exe6(sym, subgraph_backend, op_names) |
| test_subgraph_exe7(sym, subgraph_backend, op_names) |
| test_subgraph_exe8(sym, subgraph_backend, op_names) |
| test_subgraph_exe9(sym, subgraph_backend, op_names) |
| test_subgraph_backend_gluon(sym, subgraph_backend, op_names, tmpdir) |
| test_subgraph_backend_gluon_ext1(tmpdir) |
| test_subgraph_backend_gluon_ext2(tmpdir) |
