def pull_embedding_vectors(self, layer_name, embedding_ids):
"""Pulls and returns embedding vectors ordered by the embedding ids."""
ps_ids = {}
ps_ids_index = {}
for idx, embedding_id in enumerate(embedding_ids):
ps_id = int_to_id(embedding_id, self._ps_num)
ps_ids.setdefault(ps_id, []).append(embedding_id)
ps_ids_index.setdefault(ps_id, []).append(idx)
embeddings = []
index = []
pb_future_and_id_pairs = []
for ps_id, embedding_ids in ps_ids.items():
req = elasticdl_pb2.PullEmbeddingVectorRequest()
req.name = layer_name
req.ids.extend(embedding_ids)
pb_future = self._ps_stubs[ps_id].pull_embedding_vectors.future(
req
)
pb_future_and_id_pairs.append((pb_future, ps_id))
for pb_future, ps_id in pb_future_and_id_pairs:
pb = pb_future.result()
embeddings.append(pb_to_ndarray(pb))
index.extend(ps_ids_index[ps_id])
embeddings = np.concatenate(embeddings)
# adjust the order of embedding vectors
new_embeddings = np.empty_like(embeddings)
new_embeddings[index] = embeddings
return new_embeddings
def get_embedding_vectors(self, name, ids):
pull_req = elasticdl_pb2.PullEmbeddingVectorRequest()
pull_req.name = name
pull_req.ids.extend(ids)
res = self._stub.pull_embedding_vectors(pull_req)
if res.tensor_content:
return pb_to_ndarray(res)
else:
return None
