def graph_equal(src, des):
if isinstance(src, mx.sym.Symbol) and len(src) > 1:
for i, op in enumerate(src):
if i >= len(des):
return op, des[-1]
r1, r2 = graph_equal(op, des[i])
if r1 is not None:
return r1, r2
return None, None
if src.attr('op_name') != des.attr('op_name'):
return src, des
if sutils.get_entry_id(src) != sutils.get_entry_id(des):
return src, des
if src.list_attr() != des.list_attr():
return src, des
src_childs = sutils.sym_iter(src.get_children())
des_childs = sutils.sym_iter(des.get_children())
if src_childs is None:
if des_childs is None:
return None, None
else:
return src, des
if len(src_childs) != len(des_childs):
return src, des
for i, op in enumerate(src_childs):
r1, r2 = graph_equal(op, des_childs[i])
if r1 is not None:
return r1, r2
return None, None
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 get_conv_names(modelname):
conv_name_dct_old = {}
weight_thresh = {}
myfile = path.expanduser('~/tvm-cvm/cvm/models/' + modelname + '_mid.ini')
for k, v in [v.split(':') for v in load_file(myfile)]:
wname = k.strip()
thresh = float(v.replace(',', '').strip())
weight_thresh[wname] = thresh
myfile = path.expanduser('~/tvm-cvm/cvm/models/' + modelname +
'_normal.ini')
weight_normal = [v.strip() for v in load_file(myfile)]
sym_file = path.expanduser('~/tvm-cvm/data/' + modelname + '.prepare.json')
params_file = path.expanduser('~/tvm-cvm/data/' + modelname +
'.prepare.params')
sym = mx.sym.load(sym_file)
params = nd.load(params_file)
weight_2_conv = []
for sym in sutils.topo_sort(sym, params):
if sym.attr('op_name') == 'Convolution':
name = sym.attr('name')
wname = sutils.sym_iter(sym.get_children())[1].attr('name')
if wname in weight_thresh or wname in weight_normal:
continue
weight_2_conv.append((wname, name))
return weight_2_conv
def _impl(op, params, graph, **kwargs):
deps = kwargs['deps']
name, op_name = op.attr('name'), op.attr('op_name')
childs, attr = sym_iter(op.get_children()), op.list_attr()
if op_name == 'null':
start_time = None
out = data if is_inputs(op, params) else params[name]
elif childs is None:
start_time= time.time()
out = get_nd_op(op_name)(**attr)
if gpu_flag:
nd.waitall()
end_time = time.time()
else:
cinfos = [(c.attr('name'), get_entry_id(c)) for c in childs]
nd_inputs = [out_cache[n[0]][n[1]] for n in cinfos]
start_time = time.time()
out = get_nd_op(op_name)(*nd_inputs, **attr)
if gpu_flag:
nd.waitall()
end_time = time.time()
for n, _ in cinfos:
assert n in deps
deps[n].remove(name)
if len(deps[n]) == 0:
del out_cache[n]
if start_time is not None:
if op_name not in times:
times[op_name] = {}
times[op_name][name] = end_time - start_time
out = [out] if len(op) == 1 else out
out_cache[name] = [o.as_in_context(ctx) for o in out]
def _impl(op, params, graph, **kwargs):
deps = kwargs['deps']
name, op_name = op.attr('name'), op.attr('op_name')
childs, attr = sutils.sym_iter(op.get_children()), op.list_attr()
if op_name == 'null':
out = data if sutils.is_inputs(op, params) else params[name]
elif childs is None:
out = sutils.get_nd_op(op_name)(**attr)
else:
cinfos = [(c.attr('name'), sutils.get_entry_id(c)) for c in childs]
nd_inputs = [out_cache[n[0]][n[1]] for n in cinfos]
out = sutils.get_nd_op(op_name)(*nd_inputs, **attr)
for n, _ in cinfos:
assert n in deps
if name not in deps[n]:
# for op like: op = broadcast_mul(X, X)
# `cinfos` will have duplicate entries
# avoid removing more than once
continue
deps[n].remove(name)
if len(deps[n]) == 0:
del out_cache[n]
if name == check_point:
ans[check_point] = out
out = [out] if len(op) == 1 else out
out_cache[name] = [o.as_in_context(ctx) for o in out]
def split_model(symbol, params, inputs_ext, keys, logger=logging):
infer_shapes = spass.sym_infer_shape(symbol, params, inputs_ext)
bases = [s for s in sutils.topo_sort(symbol) if s.attr('name') in keys]
base = mx.sym.Group(bases)
base_params = {k: params[k] for k in base.list_inputs() if k in params}
base_inputs_ext = inputs_ext
graph = {}
inputs = {k: v for k, v in inputs_ext.items()}
for sym in sutils.topo_sort(symbol):
name, op_name = sym.attr('name'), sym.attr('op_name')
childs, attr = sutils.sym_iter(sym.get_children()), sym.list_attr()
node = sym
if childs is not None:
childs = [graph[c.attr('name')] for c in childs]
node = sutils.get_mxnet_op(op_name)(*childs, **attr, name=name)
if name in keys:
node = mx.sym.var(name)
inputs[name] = {'shape': infer_shapes[name]}
graph[name] = node
nodes = [graph[sym.attr('name')] for sym in symbol]
top = nodes[0] if len(nodes) == 1 else mx.sym.Group(nodes)
top_params = {k: params[k] for k in top.list_inputs() if k in params}
top_inputs_ext = {k: v for k, v in inputs.items() if k not in inputs_ext}
return base, base_params, base_inputs_ext, top, top_params, top_inputs_ext
def box_nms(node, params, graph):
name, op_name = node.attr('name'), node.attr('op_name')
childs, attr = sutils.sym_iter(node.get_children()), node.list_attr()
if op_name == '_contrib_box_nms':
valid_thresh = sutils.get_attr(attr, 'valid_thresh', 0)
attr['valid_thresh'] = int(valid_thresh * oscales[3])
node = sutils.get_mxnet_op(op_name)(*childs, **attr, name=name)
return node
def _fuse_custom_pad_transpose(sym, params, **kwargs):
name, op_name = sym.attr('name'), sym.attr('op_name')
attr, childs = sym.list_attr(), sutils.sym_iter(sym.get_children())
ret = sym
if op_name != 'transpose' or not is_pad_op(childs[0]):
return ret, params
cattr = childs[0].list_attr()
padding = sutils.get_attr(cattr, 'padding')
axes = sutils.get_attr(attr, 'axes')
cchilds = sutils.sym_iter(childs[0].get_children())
X = mx.sym.transpose(*cchilds, axes=axes)
ret = mx.sym.Custom(X,
padding=[padding[r] for r in axes],
op_type="cvm_pad")
return ret, params
def _fuse_custom_pad_transpose(sym, params, **kwargs):
name, op_name = sym.attr('name'), sym.attr('op_name')
attr, childs = sym.list_attr(), sutils.sym_iter(sym.get_children())
cattr = childs[0].list_attr() if childs else None
ret = sym
if op_name != 'transpose' or childs[0].attr('op_name') != 'Custom' or \
not cattr or 'op_type' not in cattr or \
cattr['op_type'] != 'cvm_pad':
return ret, params
padding = sutils.get_attr(cattr, 'padding')
axes = sutils.get_attr(attr, 'axes')
cchilds = sutils.sym_iter(childs[0].get_children())
X = mx.sym.transpose(*cchilds, axes=axes)
ret = mx.sym.Custom(X,
padding=[padding[r] for r in axes],
op_type="cvm_pad")
return ret, params
def _fuse_pad_eq(sym, params, **kwargs):
name, op_name = sym.attr('name'), sym.attr('op_name')
attr, childs = sym.list_attr(), sutils.sym_iter(sym.get_children())
ret = sym
if op_name not in ['Convolution', 'Pooling'] or \
childs[0].attr('op_name') != 'Pad':
return ret, params
if 'pad' in attr:
assert sutils.get_attr(attr, 'pad') == (0, 0)
cattr = childs[0].list_attr()
pad_width = sutils.get_attr(cattr, 'pad_width')
if len(pad_width) != 8 or pad_width[4] != pad_width[5] or \
pad_width[6] != pad_width[7]:
return ret, params
attr['pad'] = (pad_width[4], pad_width[6])
X = sutils.sym_iter(childs[0].get_children()) + childs[1:]
ret = sutils.get_mxnet_op(op_name)(*X, **attr)
return ret, params
def _fuse_custom_pad(sym, params, **kwargs):
name, op_name = sym.attr('name'), sym.attr('op_name')
attr, childs = sym.list_attr(), sutils.sym_iter(sym.get_children())
ret = sym
if op_name != 'Custom' or 'op_type' not in attr or \
attr['op_type'] != 'cvm_pad':
return ret, params
padding = nd.array(sutils.get_attr(attr, 'padding'))
padding = padding.reshape((-1, )).asnumpy().astype(np.int32).tolist()
ret = mx.sym.pad(*childs, mode='constant', pad_width=tuple(padding))
return ret, params
def op_scales(node, params, graph):
name, op_name = node.attr('name'), node.attr('op_name')
childs, attr = sutils.sym_iter(node.get_children()), node.list_attr()
if name in attr_scales:
scales = attr_scales[name]
elif op_name in attr_scales:
scales = attr_scales[op_name]
else:
return node
for k, v in scales.items():
assert k in attr, "attribute %s not in %s(%s) with %s" \
% (k, op_name, name, attr.keys())
attr[k] = int(float(attr[k]) * oscales_dict[v])
node = sutils.get_mxnet_op(op_name)(*childs, **attr, name=name)
return node
def summary(sym, err_op=None):
_s = ""
for op in sutils.topo_sort(sym):
name, op_name = op.attr('name'), op.attr('op_name')
childs, attr = sutils.sym_iter(op.get_children()), op.list_attr()
prefix = "Op " if op_name == "null" else "Var"
if (err_op is not None) and (name == err_op.attr('name')):
prefix = "> " + prefix
_s += "%5s:%-10s, Name=%-15s, Attr=%-40s" \
% (prefix, op_name, name, attr)
if childs is not None:
cinfos = ["%s(%d)" % (c.attr('name'), sutils.get_entry_id(c)) \
for c in childs]
_s += ", Inputs=%s" % ", ".join(cinfos)
_s += "\n"
return _s
def merge_model(base, base_params, base_inputs_ext, top, top_params, maps):
graph = {maps[c.attr('name')]: c for c in base}
for sym in sutils.topo_sort(top):
name, op_name = sym.attr('name'), sym.attr('op_name')
childs, attr = sutils.sym_iter(sym.get_children()), sym.list_attr()
node = sym
if childs is not None:
childs = [graph[c.attr('name')] for c in childs]
node = sutils.get_mxnet_op(op_name)(*childs, **attr, name=name)
if name in graph:
node = graph[name]
graph[name] = node
symbols = [graph[s.attr('name')] for s in top]
symbol = symbols[0] if len(symbols) == 1 else mx.sym.Group(symbols)
params = base_params
params.update(top_params)
params = {k: params[k] for k in symbol.list_inputs() if k in params}
return symbol, params
def merge_model(base_model, top_model, base_name_maps=None, callback=None):
base_name_maps = {} if base_name_maps is None else base_name_maps
graph = {base_name_maps.get(c.attr('name'), c.attr('name')):c \
for c in base_model.symbol}
for sym in topo_sort(top_model.symbol):
name, op_name = sym.attr('name'), sym.attr('op_name')
childs, attr = sym_iter(sym.get_children()), sym.list_attr()
node = sym
if childs is not None:
childs = [sutils.get_node(c, graph) for c in childs]
node = get_mxnet_op(op_name)(*childs, **attr, name=name)
if name in graph:
node = graph[name]
if callback is not None:
node = callback(node, top_model.params, graph)
graph[name] = node
nodes = [sutils.get_node(s, graph) for s in top_model.symbol]
symbol = nodes[0] if len(nodes) == 1 else mx.sym.Group(nodes)
params = base_model.params
params.update(top_model.params)
return Model(symbol, params)
def split_model(model, keys):
symbol, params = model.symbol, model.params
nodes = [s for s in topo_sort(symbol) if s.attr('name') in keys]
base = nodes[0] if len(nodes) == 1 else mx.sym.Group(nodes)
base_params = {k: params[k] for k in base.list_inputs() if k in params}
graph = {}
infer_shapes = tpass.infer_shape(symbol, params)
for sym in topo_sort(symbol):
name, op_name = sym.attr('name'), sym.attr('op_name')
childs, attr = sym_iter(sym.get_children()), sym.list_attr()
node = sym
if childs is not None:
childs = [sutils.get_node(c, graph) for c in childs]
node = get_mxnet_op(op_name)(*childs, **attr, name=name)
if name in keys:
node = mx.sym.var(name, \
shape=infer_shapes[name][sutils.get_entry_id(sym)])
graph[name] = node
nodes = [sutils.get_node(c, graph) for c in symbol]
top = nodes[0] if len(nodes) == 1 else mx.sym.Group(nodes)
top_params = {k: params[k] for k in top.list_inputs() if k in params}
return Model(base, base_params), Model(top, top_params)
