Skip to content

GitLab

Unverified Commit 4cd5c19e authored by Chao Ma's avatar Chao Ma Committed by GitHub
Browse files

use np.unique instead of loop (#858)

parent a8c28825
Hide whitespace changes
Inline Side-by-side
Showing with 6 additions and 4 deletions
python/dgl/contrib/dis_kvstore.py
View file @ 4cd5c19e
...@@ -367,8 +367,9 @@ class KVClient(object): ...@@ -367,8 +367,9 @@ class KVClient(object):
numpy_id = F.asnumpy(ID) numpy_id = F.asnumpy(ID)
count = math.ceil(self._data_size[name] / self._server_count) count = math.ceil(self._data_size[name] / self._server_count)
server_id = numpy_id / count server_id = numpy_id / count
for id in server_id: id_list, id_count = np.unique(server_id, return_counts=True)
group_size[int(id)] += 1 for idx in range(len(id_list)):
group_size[int(id_list[idx])] += id_count[idx]
min_idx = 0 min_idx = 0
max_idx = 0 max_idx = 0
for idx in range(self._server_count): for idx in range(self._server_count):
...@@ -426,8 +427,9 @@ class KVClient(object): ...@@ -426,8 +427,9 @@ class KVClient(object):
numpy_id = F.asnumpy(ID) numpy_id = F.asnumpy(ID)
count = math.ceil(self._data_size[name] / self._server_count) count = math.ceil(self._data_size[name] / self._server_count)
server_id = numpy_id / count server_id = numpy_id / count
for id in server_id: id_list, id_count = np.unique(server_id, return_counts=True)
group_size[int(id)] += 1 for idx in range(len(id_list)):
group_size[int(id_list[idx])] += id_count[idx]
min_idx = 0 min_idx = 0
max_idx = 0 max_idx = 0
server_count = 0 server_count = 0
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment