def test_resnet_rewrite_tarjan(linearize=False):
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 6 # use n>5 (see test_chain_memory)
nodes = make_resnet(n)
a0 = nodes[0]
a = nodes[-1]
checkpoints = [nodes[3], nodes[5]] # ['a03_add:0', 'a05_add:0']
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
if linearize:
added = linearize_lib.linearize(grad.op)
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
peak_memory = cpu_peak()
expected_peak = 4e6
util.report_memory(peak_memory, expected_peak)
if not REMOVE_ASSERTS:
assert (peak_memory -
expected_peak) < 1.1 * 10**6, "Difference too large."
def test_long_chain_tarjan(linearize=False):
"""Like test_chain, but use automatic rewriting with checkpoints="tarjan"
strategy."""
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 100
nodes = make_chain_tanh_constant(n)
a0 = nodes[0]
a = nodes[-1]
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
if linearize:
added = linearize_lib.linearize()
peak_memory = cpu_peak()
# points picked
# a09:0,19:0,a29:0,a39:0,a49:0,a58:0,a68:0,a78:0,a88:0,a97:0
expected_peak = 18e6
util.report_memory(peak_memory, expected_peak)
# todo: remove "REMOVE_ASSERTS"
if not REMOVE_ASSERTS:
assert (peak_memory - expected_peak) < 1.1e6, "Difference too large."
def test_long_resnet_rewrite_tarjan(linearize=False):
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 100
nodes = make_resnet(n)
a0 = nodes[0]
a = nodes[-1]
start_time = time.time()
with tf.control_dependencies([a]):
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
start_time = time.time()
if linearize:
added = linearize_lib.linearize(grad.op)
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
peak_memory = cpu_peak()
# 20 mem used with following tensors picked automatically
# ['a10_add:0', 'a19_add:0', 'a28_add:0', 'a37_add:0', 'a46_add:0',
# 'a55_add:0', 'a64_add:0', 'a73_add:0', 'a82_add:0', 'a91_add:0']
expected_peak = 18 * 10**6
util.report_memory(peak_memory, expected_peak)
if not REMOVE_ASSERTS:
assert (peak_memory - expected_peak) < 1.1e6, "Difference too large."
def test_chain_tarjan(linearize=False):
"""Like test_chain, but use automatic rewriting with checkpoints="tarjan"
strategy."""
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 6 # for n=5, only choice of a2 saves memory, and alg picks a3
# hence use n>5 to avoid this edge condition
nodes = util.make_chain_tanh_fill(n)
a0 = nodes[0]
a = nodes[-1]
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
if linearize:
linearize_lib.linearize()
peak_memory = cpu_peak()
expected_peak = 5e6 # originally needed 7 units, now a3,a5 are recomputed
util.report_memory(peak_memory, expected_peak)
if not REMOVE_ASSERTS:
assert (peak_memory - expected_peak) < 1e5, "Difference too large."
def test_resnet_rewrite_tarjan(linearize=False):
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 6 # use n>5 (see test_chain_memory)
nodes = make_resnet(n)
a0 = nodes[0]
a = nodes[-1]
checkpoints = [nodes[3], nodes[5]] # ['a03_add:0', 'a05_add:0']
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
if linearize:
added = linearize_lib.linearize(grad.op)
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
peak_memory = cpu_peak()
expected_peak = 4e6
util.report_memory(peak_memory, expected_peak)
if not REMOVE_ASSERTS:
assert (peak_memory - expected_peak) < 1.1*10**6, "Difference too large."
def test_long_resnet_rewrite_tarjan(linearize=False):
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 100
nodes = make_resnet(n)
a0 = nodes[0]
a = nodes[-1]
start_time = time.time()
with tf.control_dependencies([a]):
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
start_time = time.time()
if linearize:
added = linearize_lib.linearize(grad.op)
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
peak_memory = cpu_peak()
# 20 mem used with following tensors picked automatically
# ['a10_add:0', 'a19_add:0', 'a28_add:0', 'a37_add:0', 'a46_add:0',
# 'a55_add:0', 'a64_add:0', 'a73_add:0', 'a82_add:0', 'a91_add:0']
expected_peak = 18 * 10**6
util.report_memory(peak_memory, expected_peak)
if not REMOVE_ASSERTS:
assert (peak_memory - expected_peak) < 1.1e6, "Difference too large."
def test_long_chain_tarjan(linearize=False):
"""Like test_chain, but use automatic rewriting with checkpoints="tarjan"
strategy."""
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 100
nodes = make_chain_tanh_constant(n)
a0 = nodes[0]
a = nodes[-1]
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
if linearize:
added = linearize_lib.linearize()
peak_memory = cpu_peak()
# points picked
# a09:0,19:0,a29:0,a39:0,a49:0,a58:0,a68:0,a78:0,a88:0,a97:0
expected_peak = 18e6
util.report_memory(peak_memory, expected_peak)
# todo: remove "REMOVE_ASSERTS"
if not REMOVE_ASSERTS:
assert (peak_memory - expected_peak) < 1.1e6, "Difference too large."
def test_chain_tarjan(linearize=False):
"""Like test_chain, but use automatic rewriting with checkpoints="tarjan"
strategy."""
tf.reset_default_graph()
tf_dev = tf.device('/cpu:0')
tf_dev.__enter__()
n = 6 # for n=5, only choice of a2 saves memory, and alg picks a3
# hence use n>5 to avoid this edge condition
nodes = util.make_chain_tanh_fill(n)
a0 = nodes[0]
a = nodes[-1]
grad = memory_saving_gradients.gradients_tarjan([a], [a0])[0]
sess = create_session()
sessrun(tf.global_variables_initializer())
sessrun(grad.op)
if linearize:
linearize_lib.linearize()
peak_memory = cpu_peak()
expected_peak = 5e6 # originally needed 7 units, now a3,a5 are recomputed
util.report_memory(peak_memory, expected_peak)
if not REMOVE_ASSERTS:
assert (peak_memory - expected_peak) < 1e5, "Difference too large."