dmlc--dgl
Hongzhi (Steve), Chen
a9f2acf323
* [Misc] Black auto fix. * sort Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
21 行
626 B
Python
21 行
626 B
Python
"""DataLoader utils."""
|
|
from gluoncv.data.batchify import Pad
|
|
from mxnet import nd
|
|
|
|
import dgl
|
|
|
|
|
|
def dgl_mp_batchify_fn(data):
|
|
if isinstance(data[0], tuple):
|
|
data = zip(*data)
|
|
return [dgl_mp_batchify_fn(i) for i in data]
|
|
|
|
for dt in data:
|
|
if dt is not None:
|
|
if isinstance(dt, dgl.DGLGraph):
|
|
return [d for d in data if isinstance(d, dgl.DGLGraph)]
|
|
elif isinstance(dt, nd.NDArray):
|
|
pad = Pad(axis=(1, 2), num_shards=1, ret_length=False)
|
|
data_list = [dt for dt in data if dt is not None]
|
|
return pad(data_list)
|