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))