def report_embedding_info(self):
# TODO(qijun): only support float32
infos = []
if self._embedding_layers:
for layer in self._embedding_layers:
infos.append(
EmbeddingTableInfo(
layer.embedding_weight_name,
layer.output_dim,
layer.embeddings_initializer,
dtype_numpy_to_tensor(np.dtype("float32")),
))
if self._embedding_columns:
for column in self._embedding_columns:
# TODO(brightcoder01): The initializer in embedding column is
# a variable initializer function. For embedding layer, it's a
# tf.keras.initializers. Keep aligned between these two.
infos.append(
EmbeddingTableInfo(
column.embedding_weight_name,
column.dimension,
Initializer.UNIFORM,
dtype_numpy_to_tensor(np.dtype("float32")),
))
self._ps_client.push_embedding_table_infos(infos)
def report_embedding_info(self):
model = elasticdl_pb2.Model()
if self._embedding_layers:
embedding_infos = model.embedding_table_infos
for layer in self._embedding_layers:
embedding_info = embedding_infos.add()
embedding_info.name = layer.embedding_weight_name
embedding_info.dim = layer.output_dim
embedding_info.initializer = layer.embeddings_initializer
# set to float32
embedding_info.dtype = dtype_numpy_to_tensor(
np.dtype("float32")
)
if self._embedding_columns:
embedding_infos = model.embedding_table_infos
for column in self._embedding_columns:
embedding_info = embedding_infos.add()
embedding_info.name = column.embedding_weight_name
embedding_info.dim = column.dimension
# TODO(brightcoder01): The initializer in embedding column is
# a variable initializer function. For embedding layer, it's a
# tf.keras.initializers. Keep aligned between these two.
embedding_info.initializer = "uniform"
# set to float32
embedding_info.dtype = dtype_numpy_to_tensor(
np.dtype("float32")
)
for ps_id in range(self._ps_num):
self._ps_stubs[ps_id].push_embedding_table_infos(model)
def _create_tensor_pb(self, values, indices=None):
pb = elasticdl_pb2.Tensor()
pb.dim.extend(values.shape)
pb.dtype = dtype_numpy_to_tensor(values.dtype)
pb.content = values.tobytes()
if indices is not None:
pb.indices.extend(tuple(indices))
return pb
def to_embedding_table_info_pb(self):
"""Convert the embedding table information to a protobuf"""
embedding_pb = EmbeddingTableInfo()
embedding_pb.name = self.name
embedding_pb.dim = self.dim
embedding_pb.initializer = str(self.initializer_value)
embedding_pb.dtype = dtype_numpy_to_tensor(self.dtype)
return embedding_pb
def serialize_ndarray(array, pb):
dtype = dtype_numpy_to_tensor(array.dtype)
if not dtype:
raise ValueError("Dtype of ndarray %s is not supported", array.dtype)
pb.dtype = dtype
pb.tensor_content = array.tobytes()
for d in array.shape:
pb_d = pb.tensor_shape.dim.add()
pb_d.size = d
def serialize_tensor(tensor, tensor_pb):
"""Serialize ElasticDL Tensor to tensor protocol buffer."""
dtype = dtype_numpy_to_tensor(tensor.values.dtype)
if not dtype:
raise ValueError("Dtype of ndarray %s is not supported",
tensor.values.dtype)
tensor_pb.dtype = dtype
tensor_pb.dim.extend(tensor.values.shape)
tensor_pb.content = tensor.values.tobytes()
if tensor.is_indexed_slices():
tensor_pb.indices.extend(tuple(tensor.indices))
if tensor.name:
tensor_pb.name = tensor.name
def ndarray_to_tensor(arr, indices=None):
"""Convert ndarray to Tensor PB"""
if not is_numpy_dtype_allowed(arr.dtype):
raise ValueError("Dtype of ndarray %s is not supported", arr.dtype)
tensor = elasticdl_pb2.Tensor()
tensor.dim.extend(arr.shape)
tensor.content = arr.tobytes()
if indices:
tensor.indices.extend(indices)
tensor.dtype = dtype_numpy_to_tensor(arr.dtype)
return tensor
