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)