"VariableV2",

"AutoReloadVariable",

"ReadVariableOp",

"""Creates a new device string based on `device_string` but using /CPU:0.

parsed_device = pydev.DeviceSpec.from_string(device_string)

parsed_device.device_type = "CPU"

parsed_device.device_index = 0

"""Base class for Savers.

class SaveSpec(object):

"""Class used to describe tensor slices that need to be saved."""

def __init__(self, tensor, slice_spec, name):

"""Creates a `SaveSpec` object.

self.tensor = tensor

self.slice_spec = slice_spec

self.name = name

class SaveableObject(object):

"""Base class for saving and restoring saveable objects."""

def __init__(self, op, specs, name):

"""Creates a `SaveableObject` object.

self.op = op

self.specs = specs

self.name = name

# The device of this saveable. All tensors must be on the same device.

self.device = specs[0].tensor.device

def restore(self, restored_tensors, restored_shapes):

"""Restores this object from 'restored_tensors'.

# pylint: disable=unused-argument

raise ValueError("Calling an abstract method.")

class VariableSaveable(SaveableObject):

"""SaveableObject implementation that handles Variables."""

def __init__(self, var, slice_spec, name):

spec = BaseSaverBuilder.SaveSpec(var, slice_spec, name)

super(BaseSaverBuilder.VariableSaveable, self).__init__(var, [spec], name)

def restore(self, restored_tensors, restored_shapes):

restored_tensor = restored_tensors[0]

if restored_shapes is not None:

restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])

return state_ops.assign(

self.op,

restored_tensor,

validate_shape=restored_shapes is None and

self.op.get_shape().is_fully_defined())

class ResourceVariableSaveable(SaveableObject):

"""SaveableObject implementation that handles ResourceVariables."""

def __init__(self, var, slice_spec, name):

self.read_op = var

spec = BaseSaverBuilder.SaveSpec(var, slice_spec, name)

super(BaseSaverBuilder.ResourceVariableSaveable, self).__init__(

var, [spec], name)

def restore(self, restored_tensors, restored_shapes):

restored_tensor = restored_tensors[0]

if restored_shapes is not None:

restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])

return resource_variable_ops.assign_variable_op(

self.read_op.op.inputs[0],

restored_tensor)

def __init__(self, write_version=saver_pb2.SaverDef.V2):

self._write_version = write_version

def save_op(self, filename_tensor, saveables):

"""Create an Op to save 'saveables'.

# pylint: disable=protected-access

tensor_names = []

tensors = []

tensor_slices = []

for saveable in saveables:

for spec in saveable.specs:

tensor_names.append(spec.name)

tensors.append(spec.tensor)

tensor_slices.append(spec.slice_spec)

if self._write_version == saver_pb2.SaverDef.V1:

return io_ops._save(

filename=filename_tensor,

tensor_names=tensor_names,

tensors=tensors,

tensor_slices=tensor_slices)

elif self._write_version == saver_pb2.SaverDef.V2:

# "filename_tensor" is interpreted *NOT AS A FILENAME*, but as a prefix

# of a V2 checkpoint: e.g. "/fs/train/ckpt-<step>/tmp/worker<i>-<step>".

return io_ops.save_v2(filename_tensor, tensor_names, tensor_slices,

tensors)

else:

raise RuntimeError("Unexpected write_version: " + self._write_version)

# pylint: disable=unused-argument

def restore_op(self, filename_tensor, saveable, preferred_shard):

"""Create ops to restore 'saveable'.

# pylint: disable=protected-access

tensors = []

for spec in saveable.specs:

tensors.append(

io_ops.restore_v2(

filename_tensor,

[spec.name],

[spec.slice_spec],

[spec.tensor.dtype])[0])

return tensors

# pylint: enable=unused-argument

def sharded_filename(self, filename_tensor, shard, num_shards):

"""Append sharding information to a filename.

# pylint: disable=protected-access

return gen_io_ops._sharded_filename(filename_tensor, shard, num_shards)

def _AddSaveOps(self, filename_tensor, saveables):

"""Add ops to save variables that are on the same shard.

save = self.save_op(filename_tensor, saveables)

return control_flow_ops.with_dependencies([save], filename_tensor)

def _AddShardedSaveOpsForV2(self, checkpoint_prefix, per_device):

"""Add ops to save the params per shard, for the V2 format.

# IMPLEMENTATION DETAILS: most clients should skip.

#

# Suffix for any well-formed "checkpoint_prefix", when sharded.

# Transformations:

# * Users pass in "save_path" in save() and restore(). Say "myckpt".

# * checkpoint_prefix gets fed <save_path><_SHARDED_SUFFIX>.

#

# Example:

# During runtime, a temporary directory is first created, which contains

# files

#

# <train dir>/myckpt_temp/

# part-?????-of-?????{.index, .data-00000-of-00001}

#

# Before .save() finishes, they will be (hopefully, atomically) renamed to

#

# <train dir>/

# myckpt{.index, .data-?????-of-?????}

#

# Users only need to interact with the user-specified prefix, which is

# "<train dir>/myckpt" in this case. Save() and Restore() work with the

# prefix directly, instead of any physical pathname. (On failure and

# subsequent restore, an outdated and orphaned temporary directory can be

# safely removed.)

_SHARDED_SUFFIX = "_temp_%s/part" % uuid.uuid4().hex

tmp_checkpoint_prefix = string_ops.string_join(

[checkpoint_prefix, _SHARDED_SUFFIX])

num_shards = len(per_device)

sharded_saves = []

sharded_prefixes = []

num_shards_tensor = constant_op.constant(num_shards, name="num_shards")

last_device = None

for shard, (device, saveables) in enumerate(per_device):

last_device = device

with ops.device(device):

sharded_filename = self.sharded_filename(tmp_checkpoint_prefix, shard,

num_shards_tensor)

sharded_prefixes.append(sharded_filename)

sharded_saves.append(self._AddSaveOps(sharded_filename, saveables))

with ops.control_dependencies([x.op for x in sharded_saves]):

# Co-locates the merge step with the last device.

with ops.device(last_device):

# V2 format write path consists of a metadata merge step. Once merged,

# attempts to delete the temporary directory, "<user-fed prefix>_temp".

merge_step = gen_io_ops.merge_v2_checkpoints(

sharded_prefixes, checkpoint_prefix, delete_old_dirs=True)

with ops.control_dependencies([merge_step]):

# Returns the prefix "<user-fed prefix>" only. DOES NOT include the

# sharded spec suffix.

return array_ops.identity(checkpoint_prefix)

def _AddShardedSaveOps(self, filename_tensor, per_device):

"""Add ops to save the params per shard.

if self._write_version == saver_pb2.SaverDef.V2:

return self._AddShardedSaveOpsForV2(filename_tensor, per_device)

num_shards = len(per_device)

sharded_saves = []

num_shards_tensor = constant_op.constant(num_shards, name="num_shards")

for shard, (device, saveables) in enumerate(per_device):

with ops.device(device):

sharded_filename = self.sharded_filename(filename_tensor, shard,

num_shards_tensor)

sharded_saves.append(self._AddSaveOps(sharded_filename, saveables))

# Return the sharded name for the save path.

with ops.control_dependencies([x.op for x in sharded_saves]):

# pylint: disable=protected-access

return gen_io_ops._sharded_filespec(filename_tensor, num_shards_tensor)

def _AddRestoreOps(self,

filename_tensor,

saveables,

restore_sequentially,

reshape,

preferred_shard=-1,

name="restore_all"):

"""Add operations to restore saveables.

assign_ops = []

for saveable in saveables:

restore_control_inputs = assign_ops[-1:] if restore_sequentially else []

# Load and optionally reshape on the CPU, as string tensors are not

# available on the GPU.

# TODO(touts): Re-enable restore on GPU when we can support annotating

# string tensors as "HostMemory" inputs.

with ops.device(_set_cpu0(saveable.device) if saveable.device else None):

with ops.control_dependencies(restore_control_inputs):

tensors = self.restore_op(filename_tensor, saveable, preferred_shard)

shapes = None

if reshape:

# Compute the shapes, let the restore op decide if and how to do

# the reshape.

shapes = []

for spec in saveable.specs:

v = spec.tensor

shape = v.get_shape()

if not shape.is_fully_defined():

shape = array_ops.shape(v)

shapes.append(shape)

assign_ops.append(saveable.restore(tensors, shapes))

# Create a Noop that has control dependencies from all the updates.

return control_flow_ops.group(*assign_ops, name=name)

def _AddShardedRestoreOps(self, filename_tensor, per_device,

restore_sequentially, reshape):

"""Add Ops to restore variables from multiple devices.

sharded_restores = []

for shard, (device, saveables) in enumerate(per_device):

with ops.device(device):

sharded_restores.append(

self._AddRestoreOps(

filename_tensor,

saveables,

restore_sequentially,

reshape,

preferred_shard=shard,

name="restore_shard"))

return control_flow_ops.group(*sharded_restores, name="restore_all")

@staticmethod

def _IsVariable(v):

return isinstance(v, ops.Tensor) and v.op.type in _VARIABLE_OPS

def _GroupByDevices(self, saveables):

"""Group Variable tensor slices per device.

per_device = collections.defaultdict(lambda: [])

for saveable in saveables:

canonical_device = set(

pydev.canonical_name(spec.tensor.device) for spec in saveable.specs)

if len(canonical_device) != 1:

raise ValueError("All tensors of a saveable object must be "

"on the same device: %s" % saveable.name)

per_device[canonical_device.pop()].append(saveable)

return sorted(per_device.items(), key=lambda t: t[0])

@staticmethod

def OpListToDict(op_list):

"""Create a dictionary of names to operation lists.

if not isinstance(op_list, (list, tuple, set)):

raise TypeError("Variables to save should be passed in a dict or a "

"list: %s" % op_list)

op_list = set(op_list)

names_to_saveables = {}

# pylint: disable=protected-access

for var in op_list:

if isinstance(var, BaseSaverBuilder.SaveableObject):

names_to_saveables[var.name] = var

elif isinstance(var, variables.PartitionedVariable):

if var.name in names_to_saveables:

raise ValueError("At least two variables have the same name: %s" %

var.name)

names_to_saveables[var.name] = var

elif isinstance(var, variables.Variable) and var._save_slice_info:

name = var._save_slice_info.full_name

if name in names_to_saveables:

if not isinstance(names_to_saveables[name], list):

raise ValueError("Mixing slices and non-slices with the same name: "

"%s" % name)

names_to_saveables[name].append(var)

else:

names_to_saveables[name] = [var]

else:

var = ops.internal_convert_to_tensor(var, as_ref=True)

if not BaseSaverBuilder._IsVariable(var):

raise TypeError("Variable to save is not a Variable: %s" % var)

name = var.op.name

if name in names_to_saveables:

raise ValueError("At least two variables have the same name: %s" %

name)

names_to_saveables[name] = var

# pylint: enable=protected-access

return names_to_saveables

def _ValidateAndSliceInputs(self, names_to_saveables):

"""Returns the variables and names that will be used for a Saver.

if not isinstance(names_to_saveables, dict):

names_to_saveables = BaseSaverBuilder.OpListToDict(names_to_saveables)

saveables = []

seen_ops = set()

for name in sorted(names_to_saveables.keys()):

if not isinstance(name, six.string_types):

raise TypeError(

"names_to_saveables must be a dict mapping string names to "

"checkpointable operations. Name is not a string: %s" % name)

op = names_to_saveables[name]

if isinstance(op, BaseSaverBuilder.SaveableObject):

self._AddSaveable(saveables, seen_ops, op)

elif isinstance(op, (list, tuple, variables.PartitionedVariable)):

if isinstance(op, variables.PartitionedVariable):

op = list(op)

# A set of slices.

slice_name = None

# pylint: disable=protected-access

for variable in op:

if not isinstance(variable, variables.Variable):

raise ValueError("Slices must all be Variables: %s" % variable)

if not variable._save_slice_info:

raise ValueError("Slices must all be slices: %s" % variable)

if slice_name is None:

slice_name = variable._save_slice_info.full_name

elif slice_name != variable._save_slice_info.full_name:

raise ValueError(

"Slices must all be from the same tensor: %s != %s" %

(slice_name, variable._save_slice_info.full_name))

saveable = BaseSaverBuilder.VariableSaveable(

variable, variable._save_slice_info.spec, name)

self._AddSaveable(saveables, seen_ops, saveable)

# pylint: enable=protected-access

else:

# A variable or tensor.

variable = ops.internal_convert_to_tensor(op, as_ref=True)

if not BaseSaverBuilder._IsVariable(variable):

raise TypeError("names_to_saveables must be a dict mapping string "

"names to Tensors/Variables. Not a variable: %s" %

variable)

if variable.op.type in ["Variable", "VariableV2", "AutoReloadVariable"]:

saveable = BaseSaverBuilder.VariableSaveable(variable, "", name)

else:

saveable = BaseSaverBuilder.ResourceVariableSaveable(

variable, "", name)

self._AddSaveable(saveables, seen_ops, saveable)

return saveables

def _AddSaveable(self, saveables, seen_ops, saveable):

"""Adds the saveable to the saveables list.

if saveable.op in seen_ops:

raise ValueError("The same saveable will be restored with two names: %s" %

saveable.name)

saveables.append(saveable)

seen_ops.add(saveable.op)

def build(self,

names_to_saveables,

reshape=False,

sharded=False,

max_to_keep=5,

keep_checkpoint_every_n_hours=10000.0,

name=None,

restore_sequentially=False,

filename="model"):

"""Adds save/restore nodes to the graph and creates a SaverDef proto.

saveables = self._ValidateAndSliceInputs(names_to_saveables)

if max_to_keep is None:

max_to_keep = 0

with ops.name_scope(name, "save",

[saveable.op for saveable in saveables]) as name:

# Add the Constant string tensor for the filename.

filename_tensor = constant_op.constant(filename)

# Add the save ops.

if sharded:

per_device = self._GroupByDevices(saveables)

save_tensor = self._AddShardedSaveOps(filename_tensor, per_device)

restore_op = self._AddShardedRestoreOps(filename_tensor, per_device,

restore_sequentially, reshape)

else:

save_tensor = self._AddSaveOps(filename_tensor, saveables)

restore_op = self._AddRestoreOps(filename_tensor, saveables,

restore_sequentially, reshape)

# In the following use case, it's possible to have restore_ops be called

# something else:

# - Build inference graph and export a meta_graph.

# - Import the inference meta_graph

# - Extend the inference graph to a train graph.

# - Export a new meta_graph.

# Now the second restore_op will be called "restore_all_1".

# As such, comment out the assert for now until we know whether supporting

# such usage model makes sense.

#

# assert restore_op.name.endswith("restore_all"), restore_op.name

return saver_pb2.SaverDef(

filename_tensor_name=filename_tensor.name,

save_tensor_name=save_tensor.name,

restore_op_name=restore_op.name,

max_to_keep=max_to_keep,

sharded=sharded,

keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours,

"""Returns a filename for storing the CheckpointState.

if latest_filename is None:

latest_filename = "checkpoint"

model_checkpoint_path,

all_model_checkpoint_paths=None):

"""Generates a checkpoint state proto.

if all_model_checkpoint_paths is None:

all_model_checkpoint_paths = []

if (not all_model_checkpoint_paths or

all_model_checkpoint_paths[-1] != model_checkpoint_path):

logging.info("%s is not in all_model_checkpoint_paths. Manually adding it.",

model_checkpoint_path)

all_model_checkpoint_paths.append(model_checkpoint_path)

# Relative paths need to be rewritten to be relative to the "save_dir"

# if model_checkpoint_path already contains "save_dir".

if not os.path.isabs(save_dir):

if not os.path.isabs(model_checkpoint_path):

model_checkpoint_path = os.path.relpath(model_checkpoint_path, save_dir)

for i in range(len(all_model_checkpoint_paths)):

p = all_model_checkpoint_paths[i]

if not os.path.isabs(p):

all_model_checkpoint_paths[i] = os.path.relpath(p, save_dir)

coord_checkpoint_proto = CheckpointState(

model_checkpoint_path=model_checkpoint_path,

all_model_checkpoint_paths=all_model_checkpoint_paths)

model_checkpoint_path,

all_model_checkpoint_paths=None,

latest_filename=None):

"""Updates the content of the 'checkpoint' file.

# Writes the "checkpoint" file for the coordinator for later restoration.

coord_checkpoint_filename = _GetCheckpointFilename(save_dir, latest_filename)

ckpt = generate_checkpoint_state_proto(

save_dir,

model_checkpoint_path,

all_model_checkpoint_paths=all_model_checkpoint_paths)

if coord_checkpoint_filename == ckpt.model_checkpoint_path:

raise RuntimeError("Save path '%s' conflicts with path used for "

"checkpoint state. Please use a different save path." %

model_checkpoint_path)

# Preventing potential read/write race condition by *atomically* writing to a

# file.

file_io.atomic_write_string_to_file(coord_checkpoint_filename,

"""Returns CheckpointState proto from the "checkpoint" file.

ckpt = None

coord_checkpoint_filename = _GetCheckpointFilename(checkpoint_dir,

latest_filename)

f = None

try:

# Check that the file exists before opening it to avoid

# many lines of errors from colossus in the logs.

if file_io.file_exists(coord_checkpoint_filename):

file_content = file_io.read_file_to_string(

coord_checkpoint_filename).decode("utf-8")

ckpt = CheckpointState()

text_format.Merge(file_content, ckpt)

if not ckpt.model_checkpoint_path:

raise ValueError("Invalid checkpoint state loaded from %s",

checkpoint_dir)

# For relative model_checkpoint_path and all_model_checkpoint_paths,

# prepend checkpoint_dir.

if not os.path.isabs(ckpt.model_checkpoint_path):

ckpt.model_checkpoint_path = os.path.join(checkpoint_dir,

ckpt.model_checkpoint_path)

for i in range(len(ckpt.all_model_checkpoint_paths)):

p = ckpt.all_model_checkpoint_paths[i]

if not os.path.isabs(p):

ckpt.all_model_checkpoint_paths[i] = os.path.join(checkpoint_dir, p)

except errors.OpError as e:

# It's ok if the file cannot be read

logging.warning(str(e))

logging.warning("%s: Checkpoint ignored", coord_checkpoint_filename)

return None

except text_format.ParseError as e:

logging.warning(str(e))

logging.warning("%s: Checkpoint ignored", coord_checkpoint_filename)

return None

finally:

if f:

f.close()

"""Saves and restores variables.

def __init__(self,

var_list=None,

reshape=False,

sharded=False,

max_to_keep=5,

keep_checkpoint_every_n_hours=10000.0,

name=None,

restore_sequentially=False,

saver_def=None,

builder=None,

defer_build=False,

allow_empty=False,

write_version=saver_pb2.SaverDef.V2,

pad_step_number=False):

"""Creates a `Saver`.