def testFuseDepthwiseConv2dNativeWithBiasAndActivation(self):
layers = [
tf.keras.layers.DepthwiseConv2D(
1, bias_initializer=tf.initializers.constant(0.25)),
tf.keras.layers.ReLU()
]
model = tf.keras.Sequential(layers)
tf.keras.backend.set_learning_phase(0)
input_tensor = tf.constant([1.0, 1.0], shape=[1, 1, 1, 2])
@tf.function
def execute_model(tensor):
return model(tensor)
graph = tf_saved_model_conversion_v2._freeze_saved_model_v2(
execute_model.get_concrete_function(input_tensor))
graph_def = graph.as_graph_def()
optimized_graph_def = fuse_depthwise_conv2d.fuse_depthwise_conv2d(
graph_def)
depthwise_conv2d_count = 0
depthwise_conv2d = None
for node in optimized_graph_def.node:
self.assertNotEqual("BiasAdd", node.op)
self.assertNotEqual("DepthwiseConv2dNative", node.op)
self.assertNotEqual("Relu", node.op)
if node.op == graph_rewrite_util.FUSED_DEPTHWISE_CONV2D:
depthwise_conv2d_count += 1
depthwise_conv2d = node
self.assertEqual(depthwise_conv2d_count, 1)
self.assertEqual(depthwise_conv2d.attr['fused_ops'].list.s,
[b'BiasAdd', b'Relu'])
self.assertEqual(depthwise_conv2d.attr['num_args'].i, 1)
def testFusePreluWithConv2d(self):
layers = [
tf.keras.layers.Conv2D(
16, [3, 3],
padding='same',
use_bias=True,
bias_initializer=tf.initializers.constant(0.25)),
tf.keras.layers.PReLU()
]
model = tf.keras.Sequential(layers)
tf.keras.backend.set_learning_phase(0)
input_tensor = tf.constant([1.0, 1.0], shape=[1, 2, 1, 1])
@tf.function
def execute_model(tensor):
return model(tensor)
graph = tf_saved_model_conversion_v2._freeze_saved_model_v2(
execute_model.get_concrete_function(input_tensor))
graph_def = graph.as_graph_def()
for node in graph_def.node:
if node.op == 'Conv2D':
node.device = "/CPU:0"
config = config_pb2.ConfigProto()
rewriter_config = config.graph_options.rewrite_options
rewriter_config.optimizers[:] = [
'pruning', 'constfold', 'arithmetic', 'dependency', 'pruning',
'remap', 'constfold', 'arithmetic', 'dependency'
]
for output in ['Identity']:
graph.add_to_collection('train_op',
graph.get_operation_by_name(output))
signature = meta_graph_pb2.SignatureDef()
graph_def = tf_saved_model_conversion_v2._run_grappler(
config, graph_def, graph, signature)
graph_def = fuse_prelu.fuse_ops_for_prelu(graph_def)
optimized_graph_def = fuse_prelu.fuse_prelu_with_fused_conv2d(
graph_def)
conv2d_op = None
for node in optimized_graph_def.node:
self.assertNotEqual("Prelu", node.op)
if node.op == '_FusedConv2D':
conv2d_op = node
self.assertNotEqual(conv2d_op, None)
self.assertEqual(conv2d_op.attr['fused_ops'].list.s,
[b'BiasAdd', b'Prelu'])
self.assertEqual(conv2d_op.attr['num_args'].i, 2)
def testFusePrelu(self):
layers = [
tf.keras.layers.PReLU(
alpha_initializer=tf.initializers.constant(0.25)),
tf.keras.layers.PReLU(
alpha_initializer=tf.initializers.constant(0.25))
]
model = tf.keras.Sequential(layers)
tf.keras.backend.set_learning_phase(0)
input_tensor = tf.constant([1.0, 1.0])
@tf.function
def execute_model(tensor):
return model(tensor)
graph = tf_saved_model_conversion_v2._freeze_saved_model_v2(
execute_model.get_concrete_function(input_tensor))
graph_def = graph.as_graph_def()
for node in graph_def.node:
if node.op == 'Conv2D':
node.device = "/CPU:0"
config = config_pb2.ConfigProto()
rewriter_config = config.graph_options.rewrite_options
rewriter_config.optimizers[:] = [
'pruning', 'constfold', 'arithmetic', 'dependency', 'pruning',
'remap', 'constfold', 'arithmetic', 'dependency'
]
for output in ['Identity']:
graph.add_to_collection('train_op',
graph.get_operation_by_name(output))
signature = meta_graph_pb2.SignatureDef()
graph_def = tf_saved_model_conversion_v2._run_grappler(
config, graph_def, graph, signature)
optimized_graph_def = fuse_prelu.fuse_ops_for_prelu(graph_def)
prelu_op_count = 0
value = None
for node in optimized_graph_def.node:
self.assertNotEqual("Relu", node.op)
if node.op == 'Prelu':
prelu_op_count += 1
if node.op == 'Const':
value = common.values_from_const(node)
self.assertEqual(prelu_op_count, 2)
self.assertEqual(value, [0.25])