dmlc--dgl
Hongzhi (Steve), Chen
704bcaf6dd
Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
18 行
358 B
Python
18 行
358 B
Python
from time import time
|
|
|
|
import dgl
|
|
|
|
import numpy as np
|
|
from dgl import backend as F
|
|
from dgl.transforms import metis_partition
|
|
|
|
|
|
def get_partition_list(g, psize):
|
|
p_gs = metis_partition(g, psize)
|
|
graphs = []
|
|
for k, val in p_gs.items():
|
|
nids = val.ndata[dgl.NID]
|
|
nids = F.asnumpy(nids)
|
|
graphs.append(nids)
|
|
return graphs
|