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