def _create_elus(g): # type: (NNEFGraph)->None input, gt_out, exp_out, sub_out, mul_out, output, zero, one, alpha = matcher.tensors( 9) _gt = matcher.Operation(name='gt', inputs=(input, zero), outputs=gt_out) _exp = matcher.Operation(name='exp', inputs=input, outputs=exp_out) _sub = matcher.Operation(name='sub', inputs=(exp_out, one), outputs=sub_out) _mul = matcher.Operation(name='mul', inputs=(sub_out, alpha), outputs=mul_out) select = matcher.Operation(name='select', inputs=(gt_out, input, mul_out), outputs=output) matcher.replace( g, select, lambda m: NNEFOperation(graph=g, name="elu", inputs=m[input], outputs=m[output], attribs=dict(_alpha=float(m[alpha].data[0]))), lambda m: m[zero].data == [0.0] and m[one].data == [1.0] and m[ alpha].rank == 0 and m[alpha].data is not None)
def transform_cgf_stb(g): # type: (TFGraph)->None orig_input, output_grad, stb1_output, stb2_output = matcher.tensors(4) stb1 = matcher.Operation( name=["tf.space_to_batch", "tf.space_to_batch_nd"], inputs=orig_input, outputs=stb1_output) _stb2 = matcher.Operation( name=["tf.space_to_batch", "tf.space_to_batch_nd"], inputs=output_grad, outputs=stb2_output) cgf = matcher.Operation(name="_conv_grad_filter", inputs=(stb1_output, stb2_output)) pattern = matcher.SetParams(cgf, allow_multi_consumer_inside=True) def action(m): # type: (matcher.Match)->None block_shape = (m[stb1].attribs["block_shape"] if m[stb1].name.endswith("_nd") else [m[stb1].attribs["block_size"]] * len(m[stb1].attribs["paddings"])) padding = "SAME" if utils.recursive_any(m[stb1].attribs["paddings"], lambda x: x > 0) else "VALID" TFOperation(graph=g, name=m[cgf].name, inputs=(m[orig_input], m[output_grad]), attribs=utils.updated_dict(m[cgf].attribs, dilation=block_shape, padding=padding), outputs=m[cgf].outputs) g.remove_operation(m[cgf], unlink=True) matcher.for_each(g, pattern, action)
def _merge_pads(g): # type: (NNEFGraph)->None t = matcher.Tensor() pad = matcher.Operation(name=['box', 'pad'], outputs=t) sliding = matcher.Operation(name=['argmax_pool', 'max_pool', 'max_pool_with_index', 'avg_pool', 'conv'], inputs={0: t}) def condition(m): # type: (matcher.Match)->bool if not (m[pad].name == 'pad' or (m[pad].name == 'box' and all(s == 1 for s in m[pad].attribs.get('size', [])) and all(s == 1 for s in m[pad].attribs.get('stride', [])) and all(s == 1 for s in m[pad].attribs.get('dilation', [])) and not m[pad].attribs.get('normalize', False))): return False value = m[pad].attribs.get('_value', 0.0) if value not in [0.0, float('-inf')]: return False if value == float('-inf'): if not m[sliding].name in ['argmax_pool', 'max_pool', 'max_pool_with_index']: return False if m[pad].attribs.get('border', 'constant') != 'constant': return False if (m[sliding].attribs.get('border', 'constant') != 'constant' and any(p != 0 or q != 0 for p, q in m[sliding].attribs.get('padding', []))): return False if m[sliding].name in ['conv'] and any(p != 0 or q != 0 for p, q in m[pad].attribs.get('padding', [])[:2]): return False return True def action(m): # type: (matcher.Match)->None value = m[pad].attribs.get('_value', 0.0) pad_padding = m[pad].attribs.get('padding', [(0, 0) * m[t].rank]) sliding_padding = m[sliding].attribs.get('padding', [(0, 0) * m[t].rank]) if m[sliding].name in ['conv']: pad_padding = pad_padding[2:] assert len(pad_padding) == len(sliding_padding) m[sliding].attribs['padding'] = [(p + pp, q + qq) for (p, q), (pp, qq) in zip(pad_padding, sliding_padding)] m[sliding].attribs['border'] = 'ignore' if value == float('-inf') else 'constant' graph_utils.remove_passthrough(g, m[pad]) matcher.for_each(graph=g, pattern=sliding, action=action, condition=condition) for op in g.operations: if op.name in ['box', 'pad'] and '_value' in op.attribs: raise utils.NNEFToolsException('Could not export {} with value={}'.format(op.name, op.attribs['_value']))
def _create_thresholds(g): # type: (NNEFGraph)->None input, threshold, gt_output, one, zero, output = matcher.tensors(6) _gt_op = matcher.Operation(name="gt", inputs=(input, threshold), outputs=gt_output) select_op = matcher.Operation(name="select", inputs=(gt_output, one, zero), outputs=output) def condition(m): return (m[output].shape == m[input].shape and m[threshold].data is not None and m[one].data is not None and m[zero].data is not None and m[threshold].rank == 0 and np.all(m[one].get_numpy_array() == 1.0) and np.all(m[zero].get_numpy_array() == 0.0)) def replacement(m): NNEFOperation( graph=g, name="_threshold", inputs=m[input], outputs=m[output], attribs=dict(threshold=float(m[threshold].get_numpy_array()))) matcher.replace(g, select_op, replacement, condition)
def _unite_powers(g): # type: (CaffeGraph)->None input, scaled, shifted, powered = matcher.tensors(4) scale_op = matcher.Operation(name=['Power'], inputs=input, outputs=scaled) shift_op = matcher.Operation(name=['Power'], inputs=scaled, outputs=shifted) power_op = matcher.Operation(name=['Power', 'Exp'], inputs=shifted, outputs=powered) def condition(m): return (list(m[scale_op].attribs.keys()) == ['scale'] and list(m[shift_op].attribs.keys()) == ['shift'] and list(m[power_op].attribs.keys()) in (['power'], ['base'])) def replacement(m): CaffeOperation(graph=g, name=m[power_op].name, inputs=m[input], outputs=m[powered], attribs=utils.dict_union(m[scale_op].attribs, m[shift_op].attribs, m[power_op].attribs)) matcher.replace(g, power_op, replacement, condition) input, output1, output2 = matcher.tensors(3) op1 = matcher.Operation(name='Power', inputs=input, outputs=output1) op2 = matcher.Operation(name=['Power', 'Exp'], inputs=output1, outputs=output2) def condition(m): return ((list(m[op2].attribs.keys()) in (['power'], ['base']) and list(m[op1].attribs.keys()) == ['shift']) or (list(m[op2].attribs.keys()) in (['power'], ['base']) and list(m[op1].attribs.keys()) == ['scale']) or (list(m[op2].attribs.keys()) == ['shift'] and list(m[op1].attribs.keys()) == ['scale'])) def replacement(m): CaffeOperation(graph=g, name=m[op2].name, inputs=m[input], outputs=m[output2], attribs=utils.dict_union(m[op2].attribs, m[op1].attribs)) matcher.replace(g, op2, replacement, condition)
def transform_bts_conv_stb(g): # type: (TFGraph)->None input, filter, stb_output, conv_output = matcher.tensors(4) stb = matcher.Operation(name=["tf.space_to_batch", "tf.space_to_batch_nd"], inputs=input, outputs=stb_output) conv = matcher.Operation(name=["_conv", "_deconv"], inputs=(stb_output, filter), outputs=conv_output) bts = matcher.Operation(name=["tf.batch_to_space", "tf.batch_to_space_nd"], inputs=conv_output) def replacement(m): # type: (matcher.Match)->TFOperation block_shape = (m[stb].attribs["block_shape"] if m[stb].name.endswith("_nd") else [m[stb].attribs["block_size"]] * len(m[stb].attribs["paddings"])) if m[conv].name == "_conv": padding = "SAME" if utils.recursive_any( m[stb].attribs["paddings"], lambda x: x > 0) else "VALID" return TFOperation(graph=g, name=m[conv].name, inputs=(m[input], m[filter]), attribs=utils.updated_dict(m[conv].attribs, dilation=block_shape, padding=padding), outputs=m[bts].outputs) else: padding = "SAME" if utils.recursive_any( m[bts].attribs["crops"], lambda x: x > 0) else "VALID" output_shape = _apply_block_shape( shape=m[conv].attribs["output_shape"], block_shape=block_shape, data_format=m[conv].attribs["data_format"], crops=m[bts].attribs["crops"]) return TFOperation(graph=g, name=m[conv].name, inputs=(m[input], m[filter]), attribs=utils.updated_dict( m[conv].attribs, dilation=block_shape, padding=padding, output_shape=output_shape), outputs=m[bts].outputs) matcher.replace(g, bts, replacement)
def transform_fuse_activations(tf_graph): # type: (TFGraph)->None fuse_to = [ "tf.add", "tf.subtract", "tf.multiply", "tf.divide", "tf.nn.conv2d", "tf.nn.depthwise_conv2d", "tf.nn.max_pool", "tf.nn.avg_pool", # "tf.nn.conv2d_transpose", (not working yet) "tf.matmul", "tf.nn.l2_normalize", # "tf.concat" (not working yet) ] conv_output = matcher.Tensor() convlike = matcher.Operation(name=fuse_to, outputs=conv_output) activation = matcher.Operation(name="tf.nn.relu", inputs={0: conv_output}) matcher.replace( tf_graph, activation, lambda m: TFOperation( graph=tf_graph, name=m[convlike].name, attribs=utils.dict_union(m[convlike].attribs, dict(fused_activation_function='RELU')), inputs=m[convlike].inputs, outputs=m[activation].outputs), lambda m: not m[convlike].attribs.get('fused_activation_function')) conv_output = matcher.Tensor() convlike = matcher.Operation(name=fuse_to, outputs=conv_output) activation = matcher.Operation(name="tf.clip_by_value", inputs={0: conv_output}) matcher.replace( graph=tf_graph, pattern=activation, replacement=lambda m: TFOperation( graph=tf_graph, name=m[convlike].name, attribs=utils.dict_union(m[convlike].attribs, dict(fused_activation_function='RELU6')), inputs=m[convlike].inputs, outputs=m[activation].outputs), condition=lambda m: (m[activation].inputs[1].data == [0] and m[activation].inputs[2].data == [6] and not m[convlike].attribs.get('fused_activation_function')))
def transform_fuse_add_to_matmul(tf_graph): # type: (TFGraph)->None matmul_output, bias = matcher.tensors(2) matmul = matcher.Operation(name="tf.matmul", outputs=matmul_output) add = matcher.Operation(name="tf.add", inputs={matmul_output, bias}) matcher.replace(tf_graph, add, lambda m: TFOperation(graph=tf_graph, name=m[matmul].name, attribs=m[matmul].attribs, inputs=tuple(m[matmul].inputs) + (m[bias],), outputs=m[add].outputs), lambda m: len(m[matmul].inputs) == 2)
def transform_fuse_bias_add_to_conv(tf_graph): # type: (TFGraph)->None conv_output = matcher.Tensor() conv = matcher.Operation(name=["tf.nn.conv2d", "tf.nn.depthwise_conv2d"], outputs=conv_output) add = matcher.Operation(name="tf.nn.bias_add", inputs={0: conv_output}) matcher.replace(tf_graph, add, lambda m: TFOperation(graph=tf_graph, name=m[conv].name, attribs=m[conv].attribs, inputs=tuple(m[conv].inputs) + (m[add].inputs[1],), outputs=m[add].outputs), lambda m: len(m[conv].inputs) == 2)
def transform_bias_add_conv(g): # type: (TFGraph)->None conv_output = matcher.Tensor() conv = matcher.Operation(name=[ "_conv", "_planewise_conv", "_separable_conv", "_deconv", "_planewise_deconv" ], inputs={2: None}, outputs={0: conv_output}) add = matcher.Operation(name="tf.nn.bias_add", inputs={0: conv_output}) matcher.replace( g, add, lambda m: TFOperation(graph=g, name=m[conv].name, inputs=(m[conv].inputs[0], m[conv]. inputs[1], m[add].inputs[1]), attribs=m[conv].attribs, outputs=m[add].outputs))
def _merge_up_bias(g): # type: (CaffeGraph)->None input, scale, scaled, bias, output = matcher.tensors(5) scale_op = matcher.Operation(name=['Scale', 'InnerProduct'], inputs=(input, scale), outputs=scaled) bias_op = matcher.Operation(name='Bias', inputs=(scaled, bias), outputs=output) matcher.replace( g, bias_op, lambda m: CaffeOperation( graph=g, name=m[scale_op].name, inputs=(m[input], m[scale], m[bias]), outputs=m[output], attribs=utils.updated_dict(m[scale_op].attribs, bias_term=True)))
def transform_pad(g): # type: (TFGraph)->None input, filter, pad_output = matcher.tensors(3) pad = matcher.Operation(name="tf.pad", inputs=input, outputs=pad_output) conv = matcher.Operation(name=[ "_conv", "_planewise_conv", "_separable_conv", "_max_pool", "_max_pool_with_index", "_avg_pool" ], inputs=(pad_output, filter)) matcher.replace( g, conv, lambda m: TFOperation( graph=g, name=m[conv].name, inputs=(m[input], m[filter]), attribs=utils.updated_dict( m[conv].attribs, padding=[tuple(p) for p in m[pad].attribs["paddings"]], _border=m[pad].attribs["mode"]), outputs=m[conv].outputs), lambda m: m[conv].attribs["padding"].upper() == 'VALID')
def _merge_batch_norm_and_scale(g): # type: (CaffeGraph)->None input, mean, variance, scale_factor, offset, scale, normed, output = matcher.tensors(8) batch_norm_op = matcher.Operation(name='BatchNorm', inputs=(input, mean, variance, scale_factor), outputs=normed) scale_op = matcher.Operation(name='Scale', inputs=(normed, scale, offset), outputs=output, attribs=dict(axis=1, num_axes=1)) matcher.replace(g, scale_op, lambda m: CaffeOperation( graph=g, name='BatchNorm+Scale', inputs=(m[input], m[mean], m[variance], m[scale_factor], m[offset], m[scale]), outputs=m[output], attribs=dict(eps=m[batch_norm_op].attribs['eps'])))