def _reshape(inputs, attrs, params):
X, shape = inputs
graph = {}
for op in sutils.topo_sort(shape):
name, op_name = op.attr('name'), op.attr('op_name')
childs, attr = sutils.sym_iter(op.get_children()), op.list_attr()
if childs is not None:
childs = [graph[c.attr('name')] for c in childs]
if sutils.is_var(op, params):
pass
elif childs is None:
params[name] = sutils.get_nd_op(op_name)(**attr)
op = mx.sym.var(name, shape=params[name].shape)
else:
childs = [graph[c.attr('name')] for c in childs]
assert all([sutils.is_params(c, params) for c in childs])
in_params = [params[c.attr('name')] for c in childs]
if op_name == "expand_dims" and in_params[0].shape == ():
params[name] = nd.array([in_params[0].asnumpy()],
dtype=in_params[0].dtype)
elif op_name == "Reshape" and sutils.get_attr(attr, 'shape') == []:
assert in_params[0].shape == (1, )
params[name] = nd.array(in_params[0].asnumpy()[0],
dtype=in_params[0].dtype)
else:
params[name] = sutils.get_nd_op(op_name)(*in_params, **attr)
op = mx.sym.var(name, shape=params[name].shape)
graph[name] = op
assert sutils.is_params(graph[shape.attr('name')], params)
shape = params[shape.attr('name')].asnumpy().tolist()
shape[0] = -1 # since dim zero is batch, set -1 for flexiblity.
return mx.sym.reshape(X, shape)
def _impl(op, params, graph):
name, op_name = op.attr('name'), op.attr('op_name')
childs, attr = sym_iter(op.get_children()), op.list_attr()
if is_var(op, params):
pass
elif childs is None:
params[name] = get_nd_op(op_name)(**attr)
attr = {'precision': str(get_bit(params[name]))}
op = mx.sym.var(name, shape=params[name].shape, attr=attr)
elif all([is_params(c, params) for c in childs]):
in_params = [params[c.attr('name')] for c in childs]
params[name] = get_nd_op(op_name)(*in_params, **attr)
attr = {'precision': str(get_bit(params[name]))}
op = mx.sym.var(name, shape=params[name].shape, attr=attr)
return op
def _impl(op, params, graph):
name, op_name = op.attr('name'), op.attr('op_name')
_, oshp, _ = op.infer_shape()
if is_params(op, params):
if oshp is None:
oshp = [params[name].shape]
op = mx.sym.var(name, shape=oshp[0])
assert params[name].shape == oshp[0], \
"Parameter %s's shape %s is inconsistent with \
params dict %s" % (name, oshp[0], params[name].shape)
elif is_inputs(op, params):
if input_shape is None:
assert oshp is not None, "It seems that graph doesn't set \
input_shape, please invoke attach_input_shape first."
else:
oshp = [input_shape]
op = mx.sym.var(name, shape=oshp[0])
infer_shapes[name] = oshp
return op
def _pack(inputs, attrs, params):
axis = attrs['axis']
inputs_reshped = []
for s in inputs:
if sutils.is_params(s, params):
name = s.attr('name')
new_name = name + '_const'
if params[name].shape == ():
assert axis == 0
params[new_name] = nd.array([params[name].asnumpy()],
dtype=params[name].dtype)
else:
params[new_name] = nd.expand_dims(params[name], axis=axis)
inputs_reshped.append(
mx.sym.var(new_name, shape=params[new_name].shape))
else:
inputs_reshped.append(mx.sym.expand_dims(s, axis=axis))
# inputs_reshped = [mx.sym.expand_dims(i, axis=axis) for i in inputs]
op = mx.sym.concat(*inputs_reshped, dim=axis)
return mx.sym.cast(op, attrs['T'])
def compile_to_cvm(model,
model_name,
datadir="/data/std_out",
input_shape=None,
target="cuda"):
""" Compile Mxnet model into CVM Accept-JSON&BIN-Format
"""
logger = logging.getLogger("mrt.compile")
symbol, params = model.symbol, model.params
datadir = path.join(datadir, model_name)
os.makedirs(datadir, exist_ok=True)
# transform from mxnet symbol to cvm
logger.info("Transform Mxnet symbol into CVM")
nnvm_sym, _ = to_nnvm(symbol, params)
dtype, nnvm_params = "int32", {}
tvm_ctx = tvm.context(target, 0)
for sym in topo_sort(symbol):
if sutils.is_params(sym, params):
key, value = sym.attr('name'), params[sym.attr('name')]
flat = value.asnumpy()
assert np.abs(flat).max() <= sutils.INT32_MAX, \
"key: {}\nvalue: {}".format(key, value)
assert (flat.astype(dtype).astype("float64") == flat).all(), \
"key: {}\nvalue: {}".format(key, value)
nnvm_params[key] = tvm.nd.array(flat.astype(dtype), tvm_ctx)
# compile to JSON&Bytes format
# graph = nnvm.graph.create(nnvm_sym)
# open("/tmp/tmp.nnvm.json", "w").write(graph.json())
logger.info("Compile into CVM graph")
if input_shape is None:
for sym in topo_sort(symbol):
if sutils.is_inputs(sym, params):
_, oshp, _ = sym.infer_shape()
assert len(oshp) == 1
input_shape = oshp[0]
input_shapes = {'data': input_shape}
with nnvm.compiler.build_config(opt_level=0):
deploy_graph, _, nnvm_params = nnvm.compiler.build(nnvm_sym,
target=target,
shape=input_shapes,
params=nnvm_params,
dtype=dtype)
# tvm parameters reduce
logger.info("Parameters precision reduce")
for sym in topo_sort(nnvm_sym):
if sutils.is_params(sym, nnvm_params):
name, attr = sym.attr('name'), sym.list_attr()
precision = sutils.get_attr(attr, "precision")
dtype = "int32" if precision > 8 else "int8"
nnvm_params[name] = tvm.nd.array(
params[name].asnumpy().astype(dtype), tvm_ctx)
# dump
logger.info("CVM Json&Params dump")
with open(path.join(datadir, "symbol"), "w") as fout:
fout.write(deploy_graph.json())
param_bytes = nnvm.compiler.save_param_dict(nnvm_params)
with open(path.join(datadir, "params"), "wb") as fout:
fout.write(param_bytes)
return deploy_graph, nnvm_params
def params_unique(symbol, params):
new_params = {s.attr('name'):params[s.attr('name')] \
for s in topo_sort(symbol) if is_params(s, params)}
return symbol, new_params
def requant(sym, oprec, oscale=None, **kwargs):
if sutils.is_params(sym, kwargs['params']):
return requant_parameter(sym.attr('name'), oprec, oscale, **kwargs)
return requant_operator(sym, oprec, oscale, **kwargs)