dmlc--dgl
Minjie Wang
27 行
658 B
Python
27 行
658 B
Python
import networkx as nx
|
|
import dgl
|
|
import torch
|
|
import pickle
|
|
import io
|
|
|
|
def _reconstruct_pickle(obj):
|
|
f = io.BytesIO()
|
|
pickle.dump(obj, f)
|
|
f.seek(0)
|
|
obj = pickle.load(f)
|
|
f.close()
|
|
return obj
|
|
|
|
def test_pickling_batched_graph():
|
|
# NOTE: this is a test for a wierd bug mentioned in
|
|
# https://github.com/dmlc/dgl/issues/438
|
|
glist = [nx.path_graph(i + 5) for i in range(5)]
|
|
glist = [dgl.DGLGraph(g) for g in glist]
|
|
bg = dgl.batch(glist)
|
|
bg.ndata['x'] = torch.randn((35, 5))
|
|
bg.edata['y'] = torch.randn((60, 3))
|
|
new_bg = _reconstruct_pickle(bg)
|
|
|
|
if __name__ == '__main__':
|
|
test_pickling_batched_graph()
|