def testPruning(self):
x = constant_op.constant(1)
tensor_list = list_ops.empty_tensor_list(element_dtype=x.dtype,
element_shape=x.shape)
def Cond(x, tl):
del tl # Unused for Cond.
return x < 5
def Body(x, tl):
return x + 1, list_ops.tensor_list_push_back(tl, x)
outputs = while_loop_v1(Cond, Body, [x, tensor_list])
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(outputs[0])
def GetOptimizedGraph():
mg = meta_graph.create_meta_graph_def(
graph=ops.get_default_graph())
rewriter_config = rewriter_config_pb2.RewriterConfig(
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
memory_optimization=rewriter_config_pb2.RewriterConfig.MANUAL)
return tf_optimizer.OptimizeGraph(rewriter_config, mg)
g = GetOptimizedGraph()
self.assertEqual(len([n for n in g.node if n.op == "Enter"]), 1)
stack = list_ops.tensor_list_stack(outputs[1], element_dtype=x.dtype)
train_op.append(stack)
g = GetOptimizedGraph()
self.assertEqual(len([n for n in g.node if n.op == "Enter"]), 2)
def testAccumulateInsideWhileLoop(self):
with self.cached_session():
random_arrays = [
np.random.rand(16, 16, 16, 16).astype(np.float32)
for _ in range(20)
]
random_tensors = [
ops.convert_to_tensor(x, dtype=dtypes_lib.float32)
for x in random_arrays
]
def cond_fn(i, acc, tensors):
del acc, tensors # unused
return i < 1 # do just one iteration
def body_fn(i, acc, tensors):
return i + 1, acc + math_ops.accumulate_n(tensors), tensors
zeros = np.zeros((16, 16, 16, 16)).astype(np.float32)
_, tf_val, _ = while_loop_v1(cond_fn, body_fn,
(0, zeros, random_tensors))
np_val = random_arrays[0]
for random_array in random_arrays[1:]:
np_val += random_array
self.assertAllClose(np_val, self.evaluate(tf_val))
def testPruning(self):
x = constant_op.constant(1)
tensor_list = list_ops.empty_tensor_list(
element_dtype=x.dtype, element_shape=x.shape)
def Cond(x, tl):
del tl # Unused for Cond.
return x < 5
def Body(x, tl):
return x + 1, list_ops.tensor_list_push_back(tl, x)
outputs = while_loop_v1(Cond, Body, [x, tensor_list])
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(outputs[0])
def GetOptimizedGraph():
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
rewriter_config = rewriter_config_pb2.RewriterConfig(
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
memory_optimization=rewriter_config_pb2.RewriterConfig.MANUAL)
return tf_optimizer.OptimizeGraph(rewriter_config, mg)
g = GetOptimizedGraph()
self.assertEqual(len([n for n in g.node if n.op == "Enter"]), 1)
stack = list_ops.tensor_list_stack(outputs[1], element_dtype=x.dtype)
train_op.append(stack)
g = GetOptimizedGraph()
self.assertEqual(len([n for n in g.node if n.op == "Enter"]), 2)
def testAccumulateInsideWhileLoop(self):
with self.cached_session():
random_arrays = [
np.random.rand(16, 16, 16, 16).astype(np.float32) for _ in range(20)
]
random_tensors = [
ops.convert_to_tensor(x, dtype=dtypes_lib.float32)
for x in random_arrays
]
def cond_fn(i, acc, tensors):
del acc, tensors # unused
return i < 1 # do just one iteration
def body_fn(i, acc, tensors):
return i + 1, acc + math_ops.accumulate_n(tensors), tensors
zeros = np.zeros((16, 16, 16, 16)).astype(np.float32)
_, tf_val, _ = while_loop_v1(cond_fn, body_fn, (0, zeros, random_tensors))
np_val = random_arrays[0]
for random_array in random_arrays[1:]:
np_val += random_array
self.assertAllClose(np_val, self.evaluate(tf_val))