def test_transform(self):
"""Test transform graph."""
transformer = ge.Transformer()
def my_transform_op_handler(info, op, new_inputs):
add_noise = op.name.startswith("Add")
op_, op_outputs_ = ge.transform.copy_op_handler(info, op, new_inputs)
if not add_noise:
return op_, op_outputs_
# add some noise to op
with info.graph_.as_default():
t_ = math_ops.add(
constant_op.constant(1.0, shape=[10], name="Noise"),
op_.outputs[0],
name="AddNoise")
# return the "noisy" op
return op_, [t_]
transformer.transform_op_handler = my_transform_op_handler
graph = ops.Graph()
transformer(self.graph, graph, "", "")
matcher0 = match.OpMatcher("AddNoise").input_ops(
"Noise",
match.OpMatcher("Add").input_ops("Const", "Input"))
matcher1 = match.OpMatcher("AddNoise_1").input_ops(
"Noise_1",
match.OpMatcher("Add_1").input_ops("Const_1", matcher0))
matcher2 = match.OpMatcher("AddNoise_2").input_ops(
"Noise_2",
match.OpMatcher("Add_2").input_ops("Const_2", matcher1))
top = ge.select_ops("^AddNoise_2$", graph=graph)[0]
self.assertTrue(matcher2(top))
def test_select_ops(self):
"""Test select ops."""
parameters = (
(("^foo/", ), 7),
(("^foo/bar/", ), 4),
(("^foo/bar/", "a"), 5),
)
for param, length in parameters:
ops = ge.select_ops(*param, graph=self.graph)
self.assertEqual(len(ops), length)
def capture_ops():
"""Capture ops.
Decorator to capture ops created in the block.
with capture_ops() as ops:
# create some ops
print(ops) # => prints ops created.
"""
micros = int(time.time() * 10**6)
scope_name = str(micros)
op_list = []
with tf.name_scope(scope_name):
yield op_list
g = tf.get_default_graph()
op_list.extend(ge.select_ops(scope_name + "/.*", graph=g))