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 == '_greater_scalar':
valid_thresh = sutils.get_attr(attr, 'scalar', 0)
attr['scalar'] = int(valid_thresh * oscales[1])
node = sutils.get_mxnet_op(op_name)(*childs, **attr, name=name)
elif op_name == '_contrib_box_nms':
valid_thresh = sutils.get_attr(attr, 'valid_thresh', 0)
attr['valid_thresh'] = int(valid_thresh * oscales[1])
node = sutils.get_mxnet_op(op_name)(*childs, **attr, name=name)
return node
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 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 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 _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 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)
def _impl(inputs, attrs, params):
assert len(inputs) == 2, \
"{} take 2 inputs, {} given".format(name, len(inputs))
return sutils.get_mxnet_op(name)(*inputs)