def test_custom_op_gpu():
# possible places to find library file
if (os.name == 'posix'):
lib = 'libcustomop_gpu_lib.so'
if os.path.exists(lib):
fname = lib
elif os.path.exists('build/' + lib):
fname = 'build/' + lib
else:
raise MXNetError("library %s not found " % lib)
elif (os.name == 'nt'):
lib = 'libcustomop_gpu_lib.dll'
if os.path.exists('windows_package\\lib\\' + lib):
fname = 'windows_package\\lib\\' + lib
else:
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
# load the library containing gemm custom operators
mx.library.load(fname)
# test symbol custom relu operator in gpu
b = mx.nd.array([[-2, -1], [1, 2]], ctx=mx.gpu())
c = mx.sym.Variable('c')
d = mx.sym.Variable('d')
e = mx.sym.my_relu(c)
base = mx.sym.relu(d)
in_grad = [mx.nd.empty((2, 2), ctx=mx.gpu())]
in_grad_base = [mx.nd.empty((2, 2), ctx=mx.gpu())]
exe = e.bind(ctx=mx.gpu(), args={'c': b}, args_grad=in_grad)
exe_base = base.bind(ctx=mx.gpu(), args={'d': b}, args_grad=in_grad_base)
out = exe.forward()
out_base = exe_base.forward()
assert_almost_equal(out_base[0].asnumpy(),
out[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
# test backward
out_grad = mx.nd.ones((2, 2), ctx=mx.gpu())
exe.backward([out_grad])
exe_base.backward([out_grad])
assert_almost_equal(in_grad_base[0].asnumpy(),
in_grad[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
def test_external_op():
# check if operator already exists
if hasattr(mx.nd, 'min_ex'):
raise MXNetError('Operator already loaded')
lib = 'libexternal_lib.so'
fname = os.path.join(base_path,'example/extensions/lib_external_ops/build/'+lib)
if not os.path.exists(fname):
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
mx.library.load(fname, False)
# execute operator
try:
mx.nd.min_ex()
except:
raise MXNetError('Operator not loaded successfully')
def test_library_loading():
if (os.name == 'posix'):
lib = 'libsample_lib.so'
if os.path.exists(lib):
fname = lib
elif os.path.exists('build/' + lib):
fname = 'build/' + lib
else:
raise MXNetError("library %s not found " % lib)
elif (os.name == 'nt'):
lib = 'libsample_lib.dll'
if os.path.exists('windows_package\\lib\\' + lib):
fname = 'windows_package\\lib\\' + lib
else:
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
mx.library.load(fname)
def test_custom_op_gpu():
# possible places to find library file
if (os.name == 'posix'):
lib = 'libcustomop_gpu_lib.so'
if os.path.exists(lib):
fname = lib
elif os.path.exists(os.path.join(base_path, 'build/' + lib)):
fname = os.path.join(base_path, 'build/' + lib)
else:
raise MXNetError("library %s not found " % lib)
elif (os.name == 'nt'):
lib = 'libcustomop_gpu_lib.dll'
if os.path.exists('windows_package\\lib\\' + lib):
fname = 'windows_package\\lib\\' + lib
else:
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
# load the library containing gemm custom operators
mx.library.load(fname)
# test symbol custom relu operator in gpu
b = mx.nd.array([[-2, -1], [1, 2]], ctx=mx.gpu())
c = mx.sym.Variable('c')
d = mx.sym.Variable('d')
e = mx.sym.my_relu(c)
base = mx.sym.relu(d)
in_grad = [mx.nd.empty((2, 2), ctx=mx.gpu())]
in_grad_base = [mx.nd.empty((2, 2), ctx=mx.gpu())]
exe = e.bind(ctx=mx.gpu(), args={'c': b}, args_grad=in_grad)
exe_base = base.bind(ctx=mx.gpu(), args={'d': b}, args_grad=in_grad_base)
out = exe.forward()
out_base = exe_base.forward()
assert_almost_equal(out_base[0].asnumpy(),
out[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
# test custom relu backward
out_grad = mx.nd.ones((2, 2), ctx=mx.gpu())
exe.backward([out_grad])
exe_base.backward([out_grad])
assert_almost_equal(in_grad_base[0].asnumpy(),
in_grad[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
# test custom noisy relu producing deterministic result given same seed managed by mxnet
d1 = mx.nd.ones(shape=(10, 10, 10), ctx=mx.cpu())
d2 = mx.nd.ones(shape=(10, 10, 10), ctx=mx.gpu())
mx.random.seed(128, ctx=mx.cpu())
r1 = mx.nd.my_noisy_relu(d1)
mx.random.seed(128, ctx=mx.cpu())
r2 = mx.nd.my_noisy_relu(d1)
assert_almost_equal(r1.asnumpy(), r2.asnumpy(), rtol=1e-3, atol=1e-3)
mx.random.seed(128, ctx=mx.gpu())
r3 = mx.nd.my_noisy_relu(d2)
mx.random.seed(128, ctx=mx.gpu())
r4 = mx.nd.my_noisy_relu(d2)
assert_almost_equal(r3.asnumpy(), r4.asnumpy(), rtol=1e-3, atol=1e-3)
def test_custom_op():
# possible places to find library file
if (os.name=='posix'):
lib = 'libcustomop_lib.so'
if os.path.exists(lib):
fname = lib
elif os.path.exists(os.path.join(base_path,'build/'+lib)):
fname = os.path.join(base_path,'build/'+lib)
else:
raise MXNetError("library %s not found " % lib)
elif (os.name=='nt'):
lib = 'libcustomop_lib.dll'
if os.path.exists('windows_package\\lib\\'+lib):
fname = 'windows_package\\lib\\'+lib
else:
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
# load the library containing gemm custom operators
mx.library.load(fname)
# test symbol 2D gemm custom operators
s = mx.sym.Variable('s')
t = mx.sym.Variable('t')
c = mx.sym.my_gemm(s,t)
d = mx.sym.state_gemm(s,t)
# baseline gemm from MXNet
base = mx.sym.linalg.gemm2(s,t)
# get some random input matrices
dim_n, dim_k, dim_m = tuple(np.random.randint(1, 5, size=3))
mat1 = mx.nd.random.uniform(-10, 10, shape=(dim_n, dim_k), ctx=mx.cpu())
mat2 = mx.nd.random.uniform(-10, 10, shape=(dim_k, dim_m), ctx=mx.cpu())
# intermediate ndarrays to be populated by gradient compute
in_grad1 = [mx.nd.empty((dim_n,dim_k),ctx=mx.cpu()),mx.nd.empty((dim_k,dim_m),ctx=mx.cpu())]
in_grad2 = [mx.nd.empty((dim_n,dim_k),ctx=mx.cpu()),mx.nd.empty((dim_k,dim_m),ctx=mx.cpu())]
in_grad_base = [mx.nd.empty((dim_n,dim_k),ctx=mx.cpu()),mx.nd.empty((dim_k,dim_m),ctx=mx.cpu())]
exe1 = c._bind(ctx=mx.cpu(),args={'s':mat1,'t':mat2},args_grad=in_grad1)
exe2 = d._bind(ctx=mx.cpu(),args={'s':mat1,'t':mat2},args_grad=in_grad2)
exe_base = base._bind(ctx=mx.cpu(),args={'s':mat1,'t':mat2},args_grad=in_grad_base)
out1 = exe1.forward()
out2 = exe2.forward()
# test stateful operator by calling it multiple times
out2 = exe2.forward()
out_base = exe_base.forward()
# check that forward compute matches one executed by MXNet
assert_almost_equal(out_base[0].asnumpy(), out1[0].asnumpy(), rtol=1e-3, atol=1e-3)
assert_almost_equal(out_base[0].asnumpy(), out2[0].asnumpy(), rtol=1e-3, atol=1e-3)
# random output grad ndarray for gradient update
out_grad = mx.nd.ones((dim_n, dim_m), ctx=mx.cpu())
exe1.backward([out_grad])
exe2.backward([out_grad])
exe_base.backward([out_grad])
# check that gradient compute matches one executed by MXNet
assert_almost_equal(in_grad_base[0].asnumpy(), in_grad1[0].asnumpy(), rtol=1e-3, atol=1e-3)
assert_almost_equal(in_grad_base[0].asnumpy(), in_grad2[0].asnumpy(), rtol=1e-3, atol=1e-3)
def test_subgraph():
# possible places to find library file
if (os.name=='posix'):
lib = 'libsubgraph_lib.so'
if os.path.exists(lib):
# plain make build, when run in the CI
fname = lib
elif os.path.exists(os.path.join(base_path, 'build/'+lib)):
# plain cmake build when run in the CI
fname = os.path.join(base_path, 'build/'+lib)
else:
raise MXNetError("library %s not found " % lib)
elif (os.name=='nt'):
lib = 'libsubgraph_lib.dll'
if os.path.exists('windows_package\\lib\\'+lib):
# plain make build, when run in the CI
fname = 'windows_package\\lib\\'+lib
else:
# plain cmake build when run in the CI
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
mx.library.load(fname)
# test simple graph with add, exp and log operators, library supports exp/log
a = mx.sym.var('a')
b = mx.sym.var('b')
c = a + b
d = mx.sym.exp(c)
sym = mx.sym.log(d)
args = {'a':mx.nd.ones((3,2),ctx=mx.cpu()), 'b':mx.nd.ones((3,2),ctx=mx.cpu())}
# baseline - regular execution in MXNet
exe = sym._bind(ctx=mx.cpu(), args=args)
out = exe.forward()
# without propogating shapes/types, passing a custom option to subgraph prop "myOpt"
# should not create subgraph since subgraph prop requires type info
mysym1 = sym.optimize_for("myProp", myOpt='yello')
exe1 = mysym1._bind(ctx=mx.cpu(), args=args)
out1 = exe1.forward()
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(), out1[0].asnumpy(), rtol=1e-3, atol=1e-3)
# with propogating shapes/types, rejecting subgraph
# this tests creating the subgraph and having the subgraph prop reject it
mysym2 = sym.optimize_for("myProp", args, reject=True)
exe2 = mysym2._bind(ctx=mx.cpu(), args=args)
out2 = exe2.forward()
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(), out2[0].asnumpy(), rtol=1e-3, atol=1e-3)
# with propogating shapes/types
mysym3 = sym.optimize_for("myProp",args)
exe3 = mysym3._bind(ctx=mx.cpu(), args=args)
out3 = exe3.forward()
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(), out3[0].asnumpy(), rtol=1e-3, atol=1e-3)
# Gluon Hybridize partitioning with shapes/types
sym_block = nn.SymbolBlock(sym, [a,b])
sym_block.initialize()
sym_block.optimize_for(mx.nd.ones((3,2)),mx.nd.ones((3,2)),backend='myProp')
out4 = sym_block(mx.nd.ones((3,2)),mx.nd.ones((3,2)))
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(), out4[0].asnumpy(), rtol=1e-3, atol=1e-3)
# Gluon Hybridize partitioning with sym.var
sym_block2 = nn.SymbolBlock(sym, [a,b])
sym_block2.initialize()
a_var = mx.sym.var('a',shape=(3,2))
b_var = mx.sym.var('b',shape=(3,2))
sym_block2.optimize_for(a_var, b_var, backend='myProp')
# Gluon Hybridize partitioning with shapes/types
sym_block3 = nn.SymbolBlock(sym, [a,b])
sym_block3.initialize()
a_data = mx.nd.ones((3,2))
b_data = mx.nd.ones((3,2))
sym_block3.optimize_for(a_data, b_data, backend='myProp')
sym_filename, params_filename = sym_block3.export('optimized')
assert sym_filename == 'optimized-symbol.json'
assert params_filename is None
sym_block4 = nn.SymbolBlock.imports(sym_filename, ['a','b'], params_filename)
out5 = sym_block4(a_data, b_data)
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(), out5[0].asnumpy(), rtol=1e-3, atol=1e-3)
def test_custom_op():
if (os.name == 'posix'):
lib = 'libsample_lib.so'
if os.path.exists(lib):
fname = lib
elif os.path.exists('build/' + lib):
fname = 'build/' + lib
else:
raise MXNetError("library %s not found " % lib)
elif (os.name == 'nt'):
lib = 'libsample_lib.dll'
if os.path.exists('windows_package\\lib\\' + lib):
fname = 'windows_package\\lib\\' + lib
else:
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
mx.library.load(fname)
# test simple 2D gemm custom op loaded from sample library
s = mx.sym.Variable('s')
t = mx.sym.Variable('t')
c = mx.sym.my_gemm(s, t)
d = mx.sym.state_gemm(s, t)
base = mx.sym.linalg.gemm2(s, t) # baseline
dim_n, dim_k, dim_m = tuple(np.random.randint(1, 5, size=3))
mat1 = mx.nd.random.uniform(-10, 10, shape=(dim_n, dim_k), ctx=mx.cpu())
mat2 = mx.nd.random.uniform(-10, 10, shape=(dim_k, dim_m), ctx=mx.cpu())
in_grad1 = [
mx.nd.empty((dim_n, dim_k), ctx=mx.cpu()),
mx.nd.empty((dim_k, dim_m), ctx=mx.cpu())
]
in_grad2 = [
mx.nd.empty((dim_n, dim_k), ctx=mx.cpu()),
mx.nd.empty((dim_k, dim_m), ctx=mx.cpu())
]
in_grad_base = [
mx.nd.empty((dim_n, dim_k), ctx=mx.cpu()),
mx.nd.empty((dim_k, dim_m), ctx=mx.cpu())
]
exe1 = c.bind(ctx=mx.cpu(),
args={
's': mat1,
't': mat2
},
args_grad=in_grad1)
exe2 = d.bind(ctx=mx.cpu(),
args={
's': mat1,
't': mat2
},
args_grad=in_grad2)
exe_base = base.bind(ctx=mx.cpu(),
args={
's': mat1,
't': mat2
},
args_grad=in_grad_base)
out1 = exe1.forward()
out2 = exe2.forward()
out2 = exe2.forward() # stateful
out_base = exe_base.forward()
assert_almost_equal(out_base[0].asnumpy(),
out1[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
assert_almost_equal(out_base[0].asnumpy(),
out2[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
out_grad = mx.nd.ones((dim_n, dim_m), ctx=mx.cpu())
exe1.backward([out_grad])
exe2.backward([out_grad])
exe_base.backward([out_grad])
assert_almost_equal(in_grad_base[0].asnumpy(),
in_grad1[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
assert_almost_equal(in_grad_base[0].asnumpy(),
in_grad2[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
def test_subgraph():
# possible places to find library file
if (os.name == 'posix'):
lib = 'libsubgraph_lib.so'
if os.path.exists(lib):
# plain make build, when run in the CI
fname = lib
elif os.path.exists('build/' + lib):
# plain cmake build when run in the CI
fname = 'build/' + lib
else:
raise MXNetError("library %s not found " % lib)
elif (os.name == 'nt'):
lib = 'libsubgraph_lib.dll'
if os.path.exists('windows_package\\lib\\' + lib):
# plain make build, when run in the CI
fname = 'windows_package\\lib\\' + lib
else:
# plain cmake build when run in the CI
raise MXNetError("library %s not found " % lib)
fname = os.path.abspath(fname)
mx.library.load(fname)
# test simple graph with add, exp and log operators, library supports exp/log
a = mx.sym.var('a')
b = mx.sym.var('b')
c = a + b
d = mx.sym.exp(c)
sym = mx.sym.log(d)
args = {
'a': mx.nd.ones((3, 2), ctx=mx.cpu()),
'b': mx.nd.ones((3, 2), ctx=mx.cpu())
}
arg_array = [
mx.nd.ones((3, 2), dtype='float32', ctx=mx.cpu()),
mx.nd.ones((3, 2), dtype='float32', ctx=mx.cpu())
]
# baseline - regular execution in MXNet
exe = sym.bind(ctx=mx.cpu(), args=args)
out = exe.forward()
# without propogating shapes/types, passing a custom option to subgraph prop "myOpt"
# should not create subgraph since subgraph prop requires type info
mysym1 = sym.optimize_for("myProp", myOpt='yello')
exe1 = mysym1.bind(ctx=mx.cpu(), args=args)
out1 = exe1.forward()
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(),
out1[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
# with propogating shapes/types, rejecting subgraph
# this tests creating the subgraph and having the subgraph prop reject it
mysym2 = sym.optimize_for("myProp", arg_array, reject=True)
exe2 = mysym2.bind(ctx=mx.cpu(), args=args)
out2 = exe2.forward()
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(),
out2[0].asnumpy(),
rtol=1e-3,
atol=1e-3)
# with propogating shapes/types
mysym3 = sym.optimize_for("myProp", arg_array)
exe3 = mysym3.bind(ctx=mx.cpu(), args=args)
out3 = exe3.forward()
# check that result matches one executed by MXNet
assert_almost_equal(out[0].asnumpy(),
out3[0].asnumpy(),
rtol=1e-3,
atol=1e-3)