def get_executor(sym,
subgraph_backend=None,
op_names=None,
original_exec=None):
if subgraph_backend is not None:
os.environ['MXNET_SUBGRAPH_BACKEND'] = subgraph_backend
check_call(
_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend),
mx_uint(len(op_names)),
c_str_array(op_names)))
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()
if subgraph_backend is not None:
check_call(
_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))
del os.environ['MXNET_SUBGRAPH_BACKEND']
return exe
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, )))
def get_executor(sym,
subgraph_backend=None,
op_names=None,
original_exec=None):
if subgraph_backend is not None:
os.environ['MXNET_SUBGRAPH_BACKEND'] = subgraph_backend
check_call(
_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend),
mx_uint(len(op_names)),
c_str_array(op_names)))
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()
if subgraph_backend is not None:
check_call(
_LIB.MXRemoveSubgraphPropertyOpNames(c_str(subgraph_backend)))
del os.environ['MXNET_SUBGRAPH_BACKEND']
return exe
def _check_subgraph_exe6(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
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.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend),
mx_uint(len(op_names)),
c_str_array(op_names)))
part_sym = sym.optimize_for(subgraph_backend, exe1.arg_dict)
check_call(
_LIB.MXRemoveSubgraphPropertyOpNames(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(),
np.zeros(shape=(1, )))
def test_subgraph_exe4(sym, subgraph_backend, op_names):
"""Use env var MXNET_SUBGRAPH_BACKEND=default to trigger graph partitioning in bind
and compare results of the partitioned sym and the original sym."""
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)
with environment('MXNET_SUBGRAPH_BACKEND', subgraph_backend):
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 test_subgraph_exe2(sym, subgraph_backend, op_names):
"""Use env var MXNET_SUBGRAPH_BACKEND=default to trigger graph partitioning in _simple_bind
and compare results of the partitioned sym and the original sym."""
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)
with environment('MXNET_SUBGRAPH_BACKEND', subgraph_backend):
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(),
np.zeros(shape=(1, )))
def _check_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
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()
# infer shape/type before partition before bind
check_call(
_LIB.MXSetSubgraphPropertyOpNames(c_str(subgraph_backend),
mx_uint(len(op_names)),
c_str_array(op_names)))
part_sym = sym.optimize_for(subgraph_backend, arg_array)
check_call(
_LIB.MXRemoveSubgraphPropertyOpNames(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(),
np.zeros(shape=(1, )))
