def _copy_source(s, graph, op_map, handle_captures, inverse_captures,
base_graph):
"""Create a source in a graph based on a Tensor from a different graph.
This function creates a placeholder analog of `s` in a graph with the
following behavior:
1) If s is a captured Tensor or Variable and handle_captures is set to True,
simply capture it in the new graph as well.
2) If s is a PlaceholderWithDefault whose default is a constant, preserve
said default in the new graph.
3) When applicable, copy resource variable metadata from `s` to the newly
created placeholder.
Args:
s: The source of interest.
graph: The destination graph.
op_map: A dict mapping ops and tensors in the old graph to the new one.
handle_captures: A boolean indicating whether to re-capture s in the new
graph or simply create a vanilla placeholder.
inverse_captures: A dict mapping s back to the Tensor or Variable that it
captures.
base_graph: The graph being copied from.
"""
if handle_captures and s in inverse_captures:
copied_placeholder = graph.capture(inverse_captures[s], name=s.op.name)
elif s.op.type == "PlaceholderWithDefault" and _constant_inputs(s):
# Copy the default value to the graph.
default_value = s.op.inputs[0]
unavailable_inputs, unavailable_control_inputs = _copy_non_source(
op=default_value.op, graph=graph, op_map=op_map,
base_graph=base_graph)
if unavailable_inputs or unavailable_control_inputs:
raise AssertionError(
"Could not copy source node {} because it has inputs."
.format(default_value))
with ops.device(s.op.device):
copied_placeholder = array_ops.placeholder_with_default(
input=op_map[default_value], shape=s.shape, name=s.op.name)
else:
with ops.device(s.op.device):
copied_placeholder = array_ops.placeholder(
dtype=s.dtype, shape=s.shape, name=s.op.name)
base_handle = resource_variable_ops.get_resource_handle_data(s)
if base_handle.shape_and_type:
resource_variable_ops._set_handle_shapes_and_types( # pylint: disable=protected-access
copied_placeholder,
base_handle,
graph_mode=True)
op_map[s] = copied_placeholder
# Add an entry for the op of the source tensor so that if there are any nodes
# depending on that op via control dependencies it can work correctly.
op_map[s.op] = copied_placeholder.op
def _copy_source(s, graph, op_map, handle_captures, inverse_captures):
"""Create a source in a graph based on a Tensor from a different graph.
This function creates a placeholder analog of `s` in a graph with the
following behavior:
1) If s is a captured Tensor or Variable and handle_captures is set to True,
simply capture it in the new graph as well.
2) If s is a PlaceholderWithDefault whose default is a constant, preserve
said default in the new graph.
3) When applicable, copy resource variable metadata from `s` to the newly
created placeholder.
Args:
s: The source of interest.
graph: The destination graph.
op_map: A dict mapping ops and tensors in the old graph to the new one.
handle_captures: A boolean indicating whether to re-capture s in the new
graph or simply create a vanilla placeholder.
inverse_captures: A dict mapping s back to the Tensor or Variable that it
captures.
"""
if handle_captures and s in inverse_captures:
copied_placeholder = graph.capture(inverse_captures[s], name=s.op.name)
elif s.op.type == "PlaceholderWithDefault" and _constant_inputs(s):
# Copy the default value to the graph.
default_value = s.op.inputs[0]
unavailable_inputs, unavailable_control_inputs = _copy_non_source(
op=default_value.op, graph=graph, op_map=op_map)
if unavailable_inputs or unavailable_control_inputs:
raise AssertionError(
"Could not copy source node {} because it has inputs."
.format(default_value))
with ops.device(s.op.device):
copied_placeholder = array_ops.placeholder_with_default(
input=op_map[default_value], shape=s.shape, name=s.op.name)
else:
with ops.device(s.op.device):
copied_placeholder = array_ops.placeholder(
dtype=s.dtype, shape=s.shape, name=s.op.name)
base_handle = resource_variable_ops.get_resource_handle_data(s)
if base_handle.shape_and_type:
resource_variable_ops._set_handle_shapes_and_types( # pylint: disable=protected-access
copied_placeholder,
base_handle,
graph_mode=True)
op_map[s] = copied_placeholder
# Add an entry for the op of the source tensor so that if there are any nodes
# depending on that op via control dependencies it can work correctly.
op_map[s.op] = copied_placeholder.op
def _set_handle_data(func_graph, fdef):
"""Adds handle data for resource type inputs and outputs."""
for tensor, arg_def in itertools.chain(
zip(func_graph.inputs, fdef.signature.input_arg),
zip(func_graph.outputs, fdef.signature.output_arg)):
if arg_def.handle_data:
shape_and_dtype = arg_def.handle_data[0]
handle_data = cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
handle_data.is_set = True
handle_data.shape_and_type.append(
cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
shape=shape_and_dtype.shape, dtype=shape_and_dtype.dtype))
resource_variable_ops._set_handle_shapes_and_types( # pylint: disable=protected-access
tensor, handle_data, True)
def _set_handle_data(func_graph, fdef):
"""Adds handle data for resource type inputs and outputs."""
# The shape of the handle itself is [], while the variable shape is
# saved in `handle_data`. Previously, the shape of the resource handle
# was set to `None`. Correct both shapes here.
for tensor, arg_def in itertools.chain(
zip(func_graph.inputs, fdef.signature.input_arg),
zip(func_graph.outputs, fdef.signature.output_arg)):
if arg_def.handle_data:
tensor.set_shape([])
shape_and_dtype = arg_def.handle_data[0]
handle_data = cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
handle_data.is_set = True
handle_data.shape_and_type.append(
cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
shape=shape_and_dtype.shape, dtype=shape_and_dtype.dtype))
resource_variable_ops._set_handle_shapes_and_types( # pylint: disable=protected-access
tensor, handle_data, True)
def __init__(self,
shape,
local_replica_id,
initializer=None,
trainable=True,
use_hashtable=True,
name="EmbeddingVariable",
dtype=None,
key_dtype=None,
*args,
**kwargs):
if (not isinstance(shape, list)) or (len(shape) != 2):
raise ValueError("shape_per_gpu must be a list which represents: "+\
"[vocabulary_size_per_gpu, embedding_vector_size].")
self.m_shape_per_gpu = TensorShape(shape)
self.m_local_replica_id = local_replica_id
self.m_initializer = initializer or InPlaceInitializer(
name="random_uniform")
if not isinstance(self.m_initializer, InPlaceInitializer):
self.m_initializer = tf_initializers.get(self.m_initializer)
self.m_trainable = trainable
self.m_use_hashtable = use_hashtable
self.m_embedding_layer = None
self.m_dtype = dtype or dtypes.float32
self.m_key_dtype = key_dtype or dtypes.int64
# produce intial_value
if isinstance(self.m_initializer, InPlaceInitializer):
# TODO: serialize it
self.m_initial_value = self.m_initializer.name
else:
self.m_initial_value = self.m_initializer(
shape=self.m_shape_per_gpu, dtype=self.m_dtype)
collections = [ops.GraphKeys.GLOBAL_VARIABLES]
if trainable and ops.GraphKeys.TRAINABLE_VARIABLES not in collections:
collections = list(collections) + [
ops.GraphKeys.TRAINABLE_VARIABLES
]
with ops.init_scope():
self._in_graph_mode = not context.executing_eagerly()
with ops.name_scope(name) as var_name_scope:
# TODO: use regulare expression
while var_name_scope[-1] == r"/":
var_name_scope = var_name_scope[:-1]
var_name = var_name_scope
self.m_var_name = var_name
self.m_unique_id = "%s_%d" % (var_name, ops.uid())
# attr = resource_variable_ops.attr_value_pb2.AttrValue(
# list=resource_variable_ops.attr_value_pb2.AttrValue.ListValue(
# s=[resource_variable_ops.compat.as_bytes("loc:@%s" % self.m_var_name)]))
# with ops.get_default_graph()._attr_scope({"_class": attr}):
with ops.NullContextmanager():
# m_handle is the handle to EmbeddingVariable, tf_handle is the handle to TF Var.
self.m_handle, self.tf_handle = kit_lib.create_var(
var_name=var_name,
dtype=self.m_dtype,
shape=self.m_shape_per_gpu)
if self._in_graph_mode:
with ops.name_scope("IsInitialized"):
self._is_initialized_op = ops.convert_to_tensor(
True) # TODO: should not hard-writing???
if (isinstance(self.m_initial_value, ops.Tensor)
and not self.m_initial_value.shape.
is_compatible_with(self.m_shape_per_gpu)):
raise ValueError(
"The initial value's shape (%s) is not compatible with "
"the explicitly supplied `shape` argument (%s)."
% (initial_value.shape,
self.m_shape_per_gpu))
_init_op = kit_lib.assign_embedding_variable(
emb_var_handle=self.m_handle,
tf_var_handle=self.tf_handle,
var_name=var_name,
initial_value=self.m_initial_value,
local_replica_id=self.m_local_replica_id,
trainable=self.m_trainable,
shape=self.m_shape_per_gpu,
use_hashtable=self.m_use_hashtable,
dtype=self.m_dtype,
key_dtype=self.m_key_dtype)
self._initializer_op = control_flow_ops.group(
(_init_op))
else:
raise RuntimeError(
"Currently, EmbeddingVariable does not support Eager mode."
)
if not context.executing_eagerly():
ops.add_to_collections(collections, self)
super(EmbeddingVariable, self).__init__(
trainable=self.m_trainable,
shape=self.m_shape_per_gpu,
dtype=self.m_dtype,
handle=self.m_handle,
handle_name=var_name,
distribute_strategy=get_strategy() if has_strategy() else None,
synchronization=VariableSynchronization.NONE,
aggregation=VariableAggregation.ONLY_FIRST_REPLICA,
unique_id=self.m_unique_id,
initializer_op=self._initializer_op,
is_initialized_op=self._is_initialized_op,
*args,
**kwargs)
handle_data = resource_variable_ops.cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData(
)
handle_data.is_set = True
handle_data.shape_and_type.append(
resource_variable_ops.cpp_shape_inference_pb2.
CppShapeInferenceResult.HandleShapeAndType(
shape=self.shape.as_proto(),
dtype=self.dtype.as_datatype_enum))
resource_variable_ops._set_handle_shapes_and_types(
self.m_handle,
handle_data,
graph_mode=False if context.executing_eagerly() else True)
resource_variable_ops._set_handle_shapes_and_types(
self.tf_handle,
handle_data,
graph_mode=False if context.executing_eagerly() else True)
def __init__(self,
shape,
local_replica_id,
initializer=None,
trainable=True,
use_hashtable=True,
name="EmbeddingVariable",
dtype=None,
key_dtype=None,
*args,
**kwargs):
if (not isinstance(shape, list)) or (len(shape) != 2):
raise ValueError("shape_per_gpu must be a list which represents: "+\
"[vocabulary_size_per_gpu, embedding_vector_size].")
self.m_shape_per_gpu = TensorShape(shape)
self.m_local_replica_id = local_replica_id
self.m_initializer = initializer or InPlaceInitializer(name="random_uniform")
if not isinstance(self.m_initializer, InPlaceInitializer):
self.m_initializer = tf_initializers.get(self.m_initializer)
self.m_trainable = trainable
self.m_use_hashtable = use_hashtable
self.m_embedding_layer = None
self.m_dtype = dtype or dtypes.float32
self.m_key_dtype = key_dtype or dtypes.int64
# produce intial_value
if isinstance(self.m_initializer, InPlaceInitializer):
# TODO: serialize it
self.m_initial_value = self.m_initializer.name
else:
self.m_initial_value = self.m_initializer(shape=self.m_shape_per_gpu, dtype=self.m_dtype)
with ops.init_scope():
with ops.name_scope(name):
self.m_var_name = self._gen_unique_name(name)
self.m_unique_id = "%s_%d" %(self.m_var_name, ops.uid())
# m_handle is the handle to EmbeddingVariable, tf_handle is the handle to TF Var.
self.m_handle, self.tf_handle = kit_lib.create_var(
var_name=self.m_var_name,
dtype=self.m_dtype,
shape=self.m_shape_per_gpu)
with ops.name_scope("IsInitialized"):
self._is_initialized_op = ops.convert_to_tensor(True)
if (isinstance(self.m_initial_value, ops.Tensor) and
not self.m_initial_value.shape.is_compatible_with(self.m_shape_per_gpu)):
raise ValueError("The initial value's shape (%s) is not compatible with "
"the explicitly supplied `shape` argument (%s)." %
(self.m_initial_value.shape, self.m_shape_per_gpu))
_init_op = kit_lib.assign_embedding_variable(emb_var_handle=self.m_handle,
tf_var_handle=self.tf_handle,
var_name=self.m_var_name,
initial_value=self.m_initial_value,
local_replica_id=self.m_local_replica_id,
trainable=self.m_trainable,
shape=self.m_shape_per_gpu,
use_hashtable=self.m_use_hashtable,
dtype=self.m_dtype,
key_dtype=self.m_key_dtype)
self._initializer_op = control_flow_ops.group((_init_op))
super(EmbeddingVariable, self).__init__(trainable=self.m_trainable,
shape=self.m_shape_per_gpu,
dtype=self.m_dtype,
handle=self.m_handle,
handle_name=self.m_var_name,
distribute_strategy=get_strategy() if has_strategy() else None,
synchronization=VariableSynchronization.NONE,
aggregation=VariableAggregation.ONLY_FIRST_REPLICA,
unique_id=self.m_unique_id,
*args, **kwargs)
handle_data = resource_variable_ops.cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
handle_data.is_set = True
handle_data.shape_and_type.append(
resource_variable_ops.cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
shape=self.shape.as_proto(), dtype=self.dtype.as_datatype_enum))
resource_variable_ops._set_handle_shapes_and_types(self.m_handle, handle_data,
graph_mode=False if context.executing_eagerly() else True)
resource_variable_ops._set_handle_shapes_and_types(self.tf_handle, handle_data,
graph_mode=False if context.executing_eagerly() else True)
