"""Op with a custom checkpointing implementation.

def __init__(self, name):

self._table_ref = gen_data_flow_ops._mutable_hash_table(

key_dtype=dtypes.string, value_dtype=dtypes.float32, name=name)

self._name = name

self._saveable = CheckpointedOp.CustomSaveable(self, name)

ops_lib.add_to_collection(ops_lib.GraphKeys.SAVEABLE_OBJECTS,

self._saveable)

@property

def name(self):

return self._name

@property

def saveable(self):

return self._saveable

def insert(self, keys, values):

return gen_data_flow_ops._lookup_table_insert(self._table_ref, keys, values)

def keys(self):

return self._export()[0]

def values(self):

return self._export()[1]

def _export(self):

return gen_data_flow_ops._lookup_table_export(self._table_ref,

dtypes.string, dtypes.float32)

class CustomSaveable(saver_module.BaseSaverBuilder.SaveableObject):

def __init__(self, table, name):

tensors = table._export()

specs = [

saver_module.BaseSaverBuilder.SaveSpec(tensors[0], "",

name + "-keys"),

saver_module.BaseSaverBuilder.SaveSpec(tensors[1], "",

name + "-values")

]

super(CheckpointedOp.CustomSaveable, self).__init__(table, specs, name)

def restore(self, restore_tensors, shapes):

return gen_data_flow_ops._lookup_table_import(self.op._table_ref,

restore_tensors[0],

def basicSaveRestore(self, variable_op):

save_path = os.path.join(self.get_temp_dir(), "basic_save_restore")

# Build a graph with 2 parameter nodes, and Save and

# Restore nodes for them.

v0 = variable_op(10.0, name="v0")

v1 = variable_op(20.0, name="v1")

v2 = CheckpointedOp(name="v2")

v2_init = v2.insert("k1", 30.0)

save = saver_module.Saver(

{

"v0": v0,

"v1": v1,

"v2": v2.saveable

}, restore_sequentially=True)

init_all_op = [variables.global_variables_initializer(), v2_init]

with self.test_session() as sess:

# Initialize all variables

sess.run(init_all_op)

# Check that the parameter nodes have been initialized.

self.assertEqual(10.0, v0.eval())

self.assertEqual(20.0, v1.eval())

self.assertEqual(b"k1", v2.keys().eval())

self.assertEqual(30.0, v2.values().eval())

# Save the initialized values in the file at "save_path"

val = save.save(sess, save_path)

self.assertTrue(isinstance(val, six.string_types))

self.assertEqual(save_path, val)

# Start a second session. In that session the parameter nodes

# have not been initialized either.

with self.test_session() as sess:

v0 = variable_op(-1.0, name="v0")

v1 = variable_op(-1.0, name="v1")

v2 = CheckpointedOp(name="v2")

save = saver_module.Saver({"v0": v0, "v1": v1, "v2": v2.saveable})

# Assert that the variables are not initialized.

self.assertEqual(

len(variables.report_uninitialized_variables().eval()), 2)

self.assertEqual(0, len(v2.keys().eval()))

self.assertEqual(0, len(v2.values().eval()))

# Restore the saved values in the parameter nodes.

save.restore(sess, save_path)

# Check that the parameter nodes have been restored.

self.assertEqual(10.0, v0.eval())

self.assertEqual(20.0, v1.eval())

self.assertEqual(b"k1", v2.keys().eval())

self.assertEqual(30.0, v2.values().eval())

# Build another graph with 2 nodes, initialized

# differently, and a Restore node for them.

with self.test_session() as sess:

v0_2 = variable_op(1000.0, name="v0")

v1_2 = variable_op(2000.0, name="v1")

v2_2 = CheckpointedOp(name="v2")

save2 = saver_module.Saver({"v0": v0_2, "v1": v1_2, "v2": v2_2.saveable})

v2_2.insert("k1000", 3000.0).run()

variables.global_variables_initializer().run()

# Check that the parameter nodes have been initialized.

self.assertEqual(1000.0, v0_2.eval())

self.assertEqual(2000.0, v1_2.eval())

self.assertEqual(b"k1000", v2_2.keys().eval())

self.assertEqual(3000.0, v2_2.values().eval())

# Restore the values saved earlier in the parameter nodes.

save2.restore(sess, save_path)

# Check that the parameter nodes have been restored.

self.assertEqual(10.0, v0_2.eval())

self.assertEqual(20.0, v1_2.eval())

self.assertEqual(b"k1", v2_2.keys().eval())

self.assertEqual(30.0, v2_2.values().eval())

def testBasic(self):

self.basicSaveRestore(variables.Variable)

def testResourceBasic(self):

self.basicSaveRestore(resource_variable_ops.ResourceVariable)

def testInvalidPath(self):

v0 = variables.Variable(0, name="v0")

for ver in (saver_pb2.SaverDef.V1, saver_pb2.SaverDef.V2):

with self.test_session() as sess:

save = saver_module.Saver({"v0": v0}, write_version=ver)

with self.assertRaisesRegexp(errors.NotFoundError,

"Failed to find any matching files for"):

save.restore(sess, "invalid path")

def testInt64(self):

save_path = os.path.join(self.get_temp_dir(), "int64")

with self.test_session() as sess:

# Build a graph with 1 node, and save and restore for them.

v = variables.Variable(np.int64(15), name="v")

save = saver_module.Saver({"v": v}, restore_sequentially=True)

variables.global_variables_initializer().run()

# Save the initialized values in the file at "save_path"

val = save.save(sess, save_path)

self.assertTrue(isinstance(val, six.string_types))

self.assertEqual(save_path, val)

with self.test_session() as sess:

v = variables.Variable(np.int64(-1), name="v")

save = saver_module.Saver({"v": v})

with self.assertRaisesWithPredicateMatch(

errors_impl.OpError, lambda e: "uninitialized value v" in e.message):

sess.run(v)

# Restore the saved values in the parameter nodes.

save.restore(sess, save_path)

# Check that the parameter nodes have been restored.

self.assertEqual(np.int64(15), v.eval())

def testSomeErrors(self):

with ops_lib.Graph().as_default():

v0 = variables.Variable([10.0], name="v0")

v1 = variables.Variable([20.0], name="v1")

v2 = variables.Variable([20.0], name="v2")

v2._set_save_slice_info(

variables.Variable.SaveSliceInfo("v1", [1], [0], [1]))

# By default the name used for "v2" will be "v1" and raise an error.

with self.assertRaisesRegexp(ValueError, "same name: v1"):

saver_module.Saver([v0, v1, v2])

# The names are different and will work.

saver_module.Saver({"vee1": v1, "other": [v2]})

# Partitioned variables also cause name conflicts.

p_v1 = variable_scope.get_variable(

"p_v1",

shape=[4, 5],

partitioner=partitioned_variables.fixed_size_partitioner(

num_shards=2))

p_v2 = variable_scope.get_variable(

"p_v2",

shape=[4, 5],

partitioner=partitioned_variables.fixed_size_partitioner(

num_shards=2))

p_v2._name = "p_v1"

with self.assertRaisesRegexp(ValueError, "same name: p_v1"):

saver_module.Saver([p_v1, p_v2])

def testSameName(self):

with ops_lib.Graph().as_default():

v0 = variables.Variable([10.0], name="v0")

v2 = CheckpointedOp(name="v2")

# Saving one variable under two names raises an error.

with self.assertRaisesRegexp(

ValueError, "The same saveable will be restored with two names: v0"):

saver_module.Saver({"v0": v0, "v0too": v0})

# Ditto for custom saveables.

with self.assertRaisesRegexp(

ValueError, "The same saveable will be restored with two names: v2"):

saver_module.Saver({"v2": v2.saveable, "v2too": v2.saveable})

# Verify non-duplicate names work.

saver_module.Saver({"v0": v0, "v2": v2.saveable})

def testBasicsWithListOfVariables(self):

save_path = os.path.join(self.get_temp_dir(), "basics_with_list")

with self.test_session(graph=ops_lib.Graph()) as sess:

# Build a graph with 2 parameter nodes, and Save and

# Restore nodes for them.

v0 = variables.Variable(10.0, name="v0")

v1 = variables.Variable(20.0, name="v1")

v2 = CheckpointedOp(name="v2")

v2_init = v2.insert("k1", 30.0)

save = saver_module.Saver([v0, v1, v2.saveable])

variables.global_variables_initializer().run()

v2_init.run()

# Check that the parameter nodes have been initialized.

self.assertEqual(10.0, v0.eval())

self.assertEqual(20.0, v1.eval())

self.assertEqual(b"k1", v2.keys().eval())

self.assertEqual(30.0, v2.values().eval())

# Save the initialized values in the file at "save_path"

val = save.save(sess, save_path)

self.assertTrue(isinstance(val, six.string_types))

self.assertEqual(save_path, val)

# Start a second session. In that session the variables

# have not been initialized either.

with self.test_session(graph=ops_lib.Graph()) as sess:

v0 = variables.Variable(-1.0, name="v0")

v1 = variables.Variable(-1.0, name="v1")

v2 = CheckpointedOp(name="v2")

save = saver_module.Saver([v0, v1, v2.saveable])

with self.assertRaisesWithPredicateMatch(

errors_impl.OpError, lambda e: "uninitialized value v0" in e.message):

sess.run(v0)

with self.assertRaisesWithPredicateMatch(

errors_impl.OpError, lambda e: "uninitialized value v1" in e.message):

sess.run(v1)

self.assertEqual(0, len(v2.keys().eval()))

self.assertEqual(0, len(v2.values().eval()))

# Restore the saved values in the parameter nodes.

save.restore(sess, save_path)

# Check that the parameter nodes have been restored.

self.assertEqual(10.0, v0.eval())

self.assertEqual(20.0, v1.eval())

self.assertEqual(b"k1", v2.keys().eval())

self.assertEqual(30.0, v2.values().eval())

# Build another graph with 2 nodes, initialized

# differently, and a Restore node for them.

with self.test_session(graph=ops_lib.Graph()) as sess:

v0_2 = variables.Variable(1000.0, name="v0")

v1_2 = variables.Variable(2000.0, name="v1")

v2_2 = CheckpointedOp(name="v2")

save2 = saver_module.Saver([v0_2, v1_2, v2_2.saveable])

v2_2.insert("k1000", 3000.0).run()

variables.global_variables_initializer().run()

# Check that the parameter nodes have been initialized.

self.assertEqual(1000.0, v0_2.eval())

self.assertEqual(2000.0, v1_2.eval())

self.assertEqual(b"k1000", v2_2.keys().eval())

self.assertEqual(3000.0, v2_2.values().eval())

# Restore the values saved earlier in the parameter nodes.

save2.restore(sess, save_path)

# Check that the parameter nodes have been restored.

self.assertEqual(10.0, v0_2.eval())

self.assertEqual(20.0, v1_2.eval())

self.assertEqual(b"k1", v2_2.keys().eval())

self.assertEqual(30.0, v2_2.values().eval())

def _SaveAndLoad(self, var_name, var_value, other_value, save_path):

with self.test_session() as sess:

var = variables.Variable(var_value, name=var_name)

save = saver_module.Saver({var_name: var})

var.initializer.run()

val = save.save(sess, save_path)

self.assertEqual(save_path, val)

with self.test_session() as sess:

var = variables.Variable(other_value, name=var_name)

save = saver_module.Saver({var_name: var})

save.restore(sess, save_path)

self.assertAllClose(var_value, var.eval())

def testCacheRereadsFile(self):

save_path = os.path.join(self.get_temp_dir(), "cache_rereads")

# Save and reload one Variable named "var0".

self._SaveAndLoad("var0", 0.0, 1.0, save_path)

# Save and reload one Variable named "var1" in the same file.

# The cached readers should know to re-read the file.

self._SaveAndLoad("var1", 1.1, 2.2, save_path)

def testAllowEmpty(self):

save_path = os.path.join(self.get_temp_dir(), "allow_empty")

with self.test_session() as sess:

_ = constant_op.constant(1)

save = saver_module.Saver(allow_empty=True)

val = save.save(sess, save_path)

self.assertIsNone(val)

with self.test_session() as sess:

save = saver_module.Saver(allow_empty=True)

save.restore(sess, save_path)

def testGPU(self):

if not test.is_gpu_available():

return

save_path = os.path.join(self.get_temp_dir(), "gpu")

with session.Session("", graph=ops_lib.Graph()) as sess:

with sess.graph.device(test.gpu_device_name()):

v0_1 = variables.Variable(123.45)

save = saver_module.Saver({"v0": v0_1})

variables.global_variables_initializer().run()

save.save(sess, save_path)

with session.Session("", graph=ops_lib.Graph()) as sess:

with sess.graph.device(test.gpu_device_name()):

v0_2 = variables.Variable(543.21)

save = saver_module.Saver({"v0": v0_2})

variables.global_variables_initializer().run()

def testVariables(self):

save_path = os.path.join(self.get_temp_dir(), "variables")

with session.Session("", graph=ops_lib.Graph()) as sess:

one = variables.Variable(1.0)

twos = variables.Variable([2.0, 2.0, 2.0])

v2 = CheckpointedOp(name="v2")

init = variables.global_variables_initializer()

save = saver_module.Saver()

init.run()

v2.insert("k1", 3.0).run()

save.save(sess, save_path)

with session.Session("", graph=ops_lib.Graph()) as sess:

one = variables.Variable(0.0)

twos = variables.Variable([0.0, 0.0, 0.0])

v2 = CheckpointedOp(name="v2")

# Saver with no arg, defaults to 'all variables'.

save = saver_module.Saver()

save.restore(sess, save_path)

self.assertAllClose(1.0, one.eval())

self.assertAllClose([2.0, 2.0, 2.0], twos.eval())

self.assertEqual(b"k1", v2.keys().eval())

self.assertEqual(3.0, v2.values().eval())

def testVarListShouldBeEmptyInDeferredBuild(self):

with ops_lib.Graph().as_default():

v = variables.Variable(1.0)

with self.assertRaisesRegexp(ValueError, "defer_build"):

saver_module.Saver([v], defer_build=True)

def testBuildShouldBeCalledBeforeSaveInCaseOfDeferBuild(self):

save_path = os.path.join(self.get_temp_dir(), "error_deferred_build")

with ops_lib.Graph().as_default(), session.Session() as sess:

variables.Variable(1.0)

saver = saver_module.Saver(defer_build=True)

with self.assertRaisesRegexp(RuntimeError, "build"):

saver.save(sess, save_path)

def testDeferredBuild(self):

save_path = os.path.join(self.get_temp_dir(), "deferred_build")

with session.Session("", graph=ops_lib.Graph()) as sess:

one = variables.Variable(1.0)

save = saver_module.Saver(defer_build=True)

# if build is not deferred, saver cannot save the `twos`.

twos = variables.Variable([2.0, 2.0, 2.0])

init = variables.global_variables_initializer()

save.build()

init.run()

save.save(sess, save_path)

with session.Session("", graph=ops_lib.Graph()) as sess:

one = variables.Variable(0.0)

twos = variables.Variable([0.0, 0.0, 0.0])

# Saver with no arg, defaults to 'all variables'.

save = saver_module.Saver()

save.restore(sess, save_path)

self.assertAllClose(1.0, one.eval())

self.assertAllClose([2.0, 2.0, 2.0], twos.eval())

def testReshape(self):

save_path = os.path.join(self.get_temp_dir(), "variables_reshape")

with session.Session("", graph=ops_lib.Graph()) as sess:

var = variables.Variable([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

init = variables.global_variables_initializer()

save = saver_module.Saver()

init.run()

save.save(sess, save_path)

# Error when restoring with default reshape=False

with session.Session("", graph=ops_lib.Graph()) as sess:

var = variables.Variable([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])

save = saver_module.Saver()

with self.assertRaisesRegexp(

errors_impl.InvalidArgumentError,

"Assign requires shapes of both tensors to match."):

save.restore(sess, save_path)

# Restored to new shape with reshape=True

with session.Session("", graph=ops_lib.Graph()) as sess:

var = variables.Variable([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]])

save = saver_module.Saver(reshape=True)

save.restore(sess, save_path)

self.assertAllClose([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], var.eval())

def testSaveWithGlobalStep(self, pad_step_number=False):

save_path = os.path.join(self.get_temp_dir(), "ckpt_with_global_step")

global_step_int = 5

# Save and reload one Variable named "var0".

self._SaveAndLoad("var0", 0.0, 1.0, save_path)

for use_tensor in [True, False]:

with self.test_session() as sess:

var = variables.Variable(1.0, name="var0")

save = saver_module.Saver(

{

var.op.name: var

}, pad_step_number=pad_step_number)

var.initializer.run()

if use_tensor:

global_step = constant_op.constant(global_step_int)

val = save.save(sess, save_path, global_step=global_step)

else:

val = save.save(sess, save_path, global_step=global_step_int)

if pad_step_number:

expected_save_path = "%s-%s" % (save_path,

"{:08d}".format(global_step_int))

else:

expected_save_path = "%s-%d" % (save_path, global_step_int)

self.assertEqual(expected_save_path, val)

def testSaveWithGlobalStepWithPadding(self):

self.testSaveWithGlobalStep(pad_step_number=True)

def testSaveToNonexistingPath(self):

save_path = os.path.join(self.get_temp_dir(), "nonexisting_dir/path")

# Build a graph with 2 parameter nodes, and Save and

# Restore nodes for them.

v0 = variables.Variable(10.0, name="v0")

v1 = variables.Variable(20.0, name="v1")

save = saver_module.Saver({"v0": v0, "v1": v1}, restore_sequentially=True)

init_all_op = variables.global_variables_initializer()

with self.test_session() as sess:

# Initialize all variables

sess.run(init_all_op)

# Check that the parameter nodes have been initialized.

self.assertEqual(10.0, v0.eval())

self.assertEqual(20.0, v1.eval())

error_msg_template = "Parent directory of {} doesn't exist, can't save."

# Assert saving fails when parent dir of save path doesn't exist

with self.assertRaisesWithPredicateMatch(

ValueError,

lambda e: error_msg_template.format(save_path) in str(e)

):

save.save(sess, save_path)

def testSaveToURI(self):

# ParseURI functions don't work on Windows yet.

# TODO(jhseu): Remove this check when it works.

if os.name == "nt":

self.skipTest("Local URI support doesn't work on Windows")

save_path = "file://" + os.path.join(self.get_temp_dir(), "uri")

# Build a graph with 2 parameter nodes, and Save and

# Restore nodes for them.

v0 = variables.Variable(10.0, name="v0")

v1 = variables.Variable(20.0, name="v1")

save = saver_module.Saver({"v0": v0, "v1": v1}, restore_sequentially=True)

init_all_op = variables.global_variables_initializer()

with self.test_session() as sess:

# Initialize all variables

sess.run(init_all_op)

# Check that the parameter nodes have been initialized.

self.assertEqual(10.0, v0.eval())

self.assertEqual(20.0, v1.eval())

def _get_test_dir(self, dirname):

test_dir = os.path.join(self.get_temp_dir(), dirname)

gfile.MakeDirs(test_dir)

return test_dir

def testBasics(self):

save_path = os.path.join(self.get_temp_dir(), "sharded_basics")

# Build a graph with 2 parameter nodes on different devices.

with session.Session(

target="",

config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:

with sess.graph.device("/cpu:0"):

v0 = variables.Variable(10, name="v0")

t0 = CheckpointedOp(name="t0")

with sess.graph.device("/cpu:1"):

v1 = variables.Variable(20, name="v1")

t1 = CheckpointedOp(name="t1")

save = saver_module.Saver(

{

"v0": v0,

"v1": v1,

"t0": t0.saveable,

"t1": t1.saveable

},

sharded=True)

variables.global_variables_initializer().run()

t0.insert("k1", 30.0).run()

t1.insert("k2", 40.0).run()

val = save.save(sess, save_path)

if save._write_version is saver_pb2.SaverDef.V1:

self.assertEqual(save_path + "-?????-of-00002", val)

else:

self.assertEqual(save_path, val)

meta_graph_filename = save._MetaGraphFilename(val)

self.assertEqual(save_path + ".meta", meta_graph_filename)

if save._write_version is saver_pb2.SaverDef.V1:

# Restore different ops from shard 0 of the saved files.

with session.Session(

target="",

config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:

with sess.graph.device("/cpu:0"):

v0 = variables.Variable(111, name="v0")

t0 = CheckpointedOp(name="t0")

save = saver_module.Saver({"v0": v0, "t0": t0.saveable}, sharded=True)

variables.global_variables_initializer().run()

t0.insert("k11", 33.0).run()

self.assertEqual(111, v0.eval())

self.assertEqual(b"k11", t0.keys().eval())

self.assertEqual(33.0, t0.values().eval())

save.restore(sess, save_path + "-00000-of-00002")

self.assertEqual(10, v0.eval())

self.assertEqual(b"k1", t0.keys().eval())

self.assertEqual(30.0, t0.values().eval())

# Restore different ops from shard 1 of the saved files.

with session.Session(

target="",

config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:

with sess.graph.device("/cpu:0"):

v1 = variables.Variable(222)

t1 = CheckpointedOp(name="t1")

save = saver_module.Saver({"v1": v1, "t1": t1.saveable}, sharded=True)

variables.global_variables_initializer().run()

t1.insert("k22", 44.0).run()

self.assertEqual(222, v1.eval())

self.assertEqual(b"k22", t1.keys().eval())

self.assertEqual(44.0, t1.values().eval())

save.restore(sess, save_path + "-00001-of-00002")

self.assertEqual(20, v1.eval())

self.assertEqual(b"k2", t1.keys().eval())

self.assertEqual(40.0, t1.values().eval())

# Now try a restore with the sharded filename.

with session.Session(

target="",

config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:

with sess.graph.device("/cpu:0"):

v0 = variables.Variable(111, name="v0")

t0 = CheckpointedOp(name="t0")

with sess.graph.device("/cpu:1"):

v1 = variables.Variable(222, name="v1")

t1 = CheckpointedOp(name="t1")

save = saver_module.Saver(

{

"v0": v0,

"v1": v1,

"t0": t0.saveable,

"t1": t1.saveable

},

sharded=True)

variables.global_variables_initializer().run()

t0.insert("k11", 33.0).run()

t1.insert("k22", 44.0).run()

self.assertEqual(111, v0.eval())

self.assertEqual(222, v1.eval())

self.assertEqual(b"k11", t0.keys().eval())

self.assertEqual(33.0, t0.values().eval())

self.assertEqual(b"k22", t1.keys().eval())

self.assertEqual(44.0, t1.values().eval())

save_path = os.path.join(self.get_temp_dir(), "sharded_basics")

if save._write_version is saver_pb2.SaverDef.V1:

save.restore(sess, save_path + "-?????-of-?????")

else:

save.restore(sess, save_path)

self.assertEqual(10, v0.eval())

self.assertEqual(20, v1.eval())

self.assertEqual(b"k1", t0.keys().eval())

self.assertEqual(30.0, t0.values().eval())

self.assertEqual(b"k2", t1.keys().eval())

self.assertEqual(40.0, t1.values().eval())

if save._write_version is saver_pb2.SaverDef.V1:

self.assertEqual(

saver_module.latest_checkpoint(self.get_temp_dir()),

os.path.join(self.get_temp_dir(), "sharded_basics-?????-of-00002"))

else:

self.assertEqual(

saver_module.latest_checkpoint(self.get_temp_dir()),

os.path.join(self.get_temp_dir(), "sharded_basics"))

def testSaverDef(self):

with self.test_session():

v0 = variables.Variable(123, name="v0")

save = saver_module.Saver({"v0": v0}, sharded=True)

sd = save.as_saver_def()

self.assertTrue(sd.sharded)

def _testPartitionedVariables(self, use_resource):

var_full_shape = [10, 3]

# Allows save/restore mechanism to work w/ different slicings.

var_name = "my_var"

saved_dir = self._get_test_dir("partitioned_variables")

saved_path = os.path.join(saved_dir, "ckpt")

call_saver_with_dict = False # updated by test loop below

def _save(slices=None, partitioner=None):

with self.test_session(graph=ops_lib.Graph()) as sess:

# Calls .eval() to return the ndarray that makes up the full variable.

rnd = random_ops.random_uniform(var_full_shape).eval()

if slices:

assert not partitioner

# TODO(apassos): make create_partitioned_variables take use_resource

# option to make this test passable without creating a named

# variable_scope.

vs = partitioned_variables.create_partitioned_variables(

var_full_shape, slices, rnd, name=var_name)

elif partitioner:

vs = [

variable_scope.get_variable(

var_name,

shape=var_full_shape,

initializer=rnd,

partitioner=partitioner,

use_resource=use_resource)

]

else:

if use_resource:

vs = [resource_variable_ops.ResourceVariable(rnd, name=var_name)]

else:

vs = [variables.Variable(rnd, name=var_name)]

variables.global_variables_initializer().run()

if call_saver_with_dict:

saver = saver_module.Saver({var_name: (vs if slices else vs[0])})

else:

saver = saver_module.Saver(vs)

actual_path = saver.save(sess, saved_path)

self.assertEqual(saved_path, actual_path)

return rnd

def _restore(slices=None, partitioner=None):

with self.test_session(graph=ops_lib.Graph()) as sess:

if slices:

assert not partitioner

new_vs = partitioned_variables.create_partitioned_variables(

var_full_shape,

slices,

array_ops.zeros(var_full_shape), # != original contents.

name=var_name)

elif partitioner:

new_vs = [

variable_scope.get_variable(

var_name,

shape=var_full_shape,

initializer=array_ops.zeros(var_full_shape),

partitioner=partitioner)

]

else:

new_vs = [

variables.Variable(

array_ops.zeros(

shape=var_full_shape), # != original contents.

name=var_name)

]

variables.global_variables_initializer().run()

if call_saver_with_dict:

saver = saver_module.Saver({

var_name: (new_vs if slices else new_vs[0])

})

else:

saver = saver_module.Saver(new_vs)

saver.restore(sess, saved_path)

if partitioner:

return new_vs[0].as_tensor().eval()

elif slices and slices[0] != 1:

return array_ops.concat(new_vs, 0).eval()

elif slices and slices[1] != 1:

return array_ops.concat(new_vs, 1).eval()

else: # Non-sliced.

return new_vs[0].eval()

for call_saver_with_dict in {False, True}:

# Save PartitionedVariable and restore into full variable.

saved_full = _save(

partitioner=partitioned_variables.fixed_size_partitioner(

num_shards=2))

restored_full = _restore()

self.assertAllEqual(saved_full, restored_full)

# Saves 10 horizontal parts of a partitioned variable.

# Restores into a full variable, non-sliced.

saved_full = _save(slices=[10, 1])

restored_full = _restore()

self.assertAllEqual(saved_full, restored_full)

# Restores into a different number/orientation of slices.

restored_full = _restore(slices=[2, 1]) # 2 horizon parts.

self.assertAllEqual(saved_full, restored_full)

restored_full = _restore(slices=[1, 3]) # 3 vertical parts.

self.assertAllEqual(saved_full, restored_full)

# Restores into a PartitionedVariable

restored_full = _restore(

partitioner=partitioned_variables.fixed_size_partitioner(

num_shards=2))

self.assertAllEqual(saved_full, restored_full)

# Now, saves a full variable and restores in slices.

saved_full = _save()

restored_full = _restore(slices=[1, 3])

self.assertAllEqual(saved_full, restored_full)

def testPartitionedVariable(self):

self._testPartitionedVariables(use_resource=False)

def testPartitionedResourceVariable(self):

def _get_test_dir(self, dirname):

test_dir = os.path.join(self.get_temp_dir(), dirname)

gfile.MakeDirs(test_dir)

return test_dir

def testNonSharded(self):

save_dir = self._get_test_dir("max_to_keep_non_sharded")

with self.test_session() as sess:

v = variables.Variable(10.0, name="v")

save = saver_module.Saver({"v": v}, max_to_keep=2)

variables.global_variables_initializer().run()

self.assertEqual([], save.last_checkpoints)

s1 = save.save(sess, os.path.join(save_dir, "s1"))

self.assertEqual([s1], save.last_checkpoints)

self.assertTrue(saver_module.checkpoint_exists(s1))

s2 = save.save(sess, os.path.join(save_dir, "s2"))

self.assertEqual([s1, s2], save.last_checkpoints)

self.assertTrue(saver_module.checkpoint_exists(s1))

self.assertTrue(saver_module.checkpoint_exists(s2))

s3 = save.save(sess, os.path.join(save_dir, "s3"))

self.assertEqual([s2, s3], save.last_checkpoints)

self.assertFalse(saver_module.checkpoint_exists(s1))

self.assertTrue(saver_module.checkpoint_exists(s2))

self.assertTrue(saver_module.checkpoint_exists(s3))

# Create a second helper, identical to the first.

save2 = saver_module.Saver(saver_def=save.as_saver_def())

save2.set_last_checkpoints(save.last_checkpoints)

# Create a third helper, with the same configuration but no knowledge of

# previous checkpoints.

save3 = saver_module.Saver(saver_def=save.as_saver_def())

# Exercise the first helper.

# Adding s2 again (old s2 is removed first, then new s2 appended)

s2 = save.save(sess, os.path.join(save_dir, "s2"))

self.assertEqual([s3, s2], save.last_checkpoints)

self.assertFalse(saver_module.checkpoint_exists(s1))

self.assertFalse(

saver_module.checkpoint_exists(save._MetaGraphFilename(s1)))

self.assertTrue(saver_module.checkpoint_exists(s3))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s3)))

self.assertTrue(saver_module.checkpoint_exists(s2))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s2)))

# Adding s1 (s3 should now be deleted as oldest in list)

s1 = save.save(sess, os.path.join(save_dir, "s1"))

self.assertEqual([s2, s1], save.last_checkpoints)

self.assertFalse(saver_module.checkpoint_exists(s3))

self.assertFalse(

saver_module.checkpoint_exists(save._MetaGraphFilename(s3)))

self.assertTrue(saver_module.checkpoint_exists(s2))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s2)))

self.assertTrue(saver_module.checkpoint_exists(s1))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s1)))

# Exercise the second helper.

# Adding s2 again (old s2 is removed first, then new s2 appended)

s2 = save2.save(sess, os.path.join(save_dir, "s2"))

self.assertEqual([s3, s2], save2.last_checkpoints)

# Created by the first helper.

self.assertTrue(saver_module.checkpoint_exists(s1))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s1)))

# Deleted by the first helper.

self.assertFalse(saver_module.checkpoint_exists(s3))

self.assertFalse(

saver_module.checkpoint_exists(save._MetaGraphFilename(s3)))

self.assertTrue(saver_module.checkpoint_exists(s2))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s2)))

# Adding s1 (s3 should now be deleted as oldest in list)

s1 = save2.save(sess, os.path.join(save_dir, "s1"))

self.assertEqual([s2, s1], save2.last_checkpoints)

self.assertFalse(saver_module.checkpoint_exists(s3))

self.assertFalse(

saver_module.checkpoint_exists(save._MetaGraphFilename(s3)))

self.assertTrue(saver_module.checkpoint_exists(s2))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s2)))

self.assertTrue(saver_module.checkpoint_exists(s1))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s1)))

# Exercise the third helper.

# Adding s2 again (but helper is unaware of previous s2)

s2 = save3.save(sess, os.path.join(save_dir, "s2"))

self.assertEqual([s2], save3.last_checkpoints)

# Created by the first helper.

self.assertTrue(saver_module.checkpoint_exists(s1))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s1)))

# Deleted by the first helper.

self.assertFalse(saver_module.checkpoint_exists(s3))

self.assertFalse(

saver_module.checkpoint_exists(save._MetaGraphFilename(s3)))

self.assertTrue(saver_module.checkpoint_exists(s2))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s2)))

# Adding s1 (s3 should not be deleted because helper is unaware of it)

s1 = save3.save(sess, os.path.join(save_dir, "s1"))

self.assertEqual([s2, s1], save3.last_checkpoints)

self.assertFalse(saver_module.checkpoint_exists(s3))

self.assertFalse(

saver_module.checkpoint_exists(save._MetaGraphFilename(s3)))

self.assertTrue(saver_module.checkpoint_exists(s2))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s2)))

self.assertTrue(saver_module.checkpoint_exists(s1))

self.assertTrue(

saver_module.checkpoint_exists(save._MetaGraphFilename(s1)))

def testSharded(self):