dmlc--dgl
Hongzhi (Steve), Chen
74c9d27d16
* [Misc] Auto-format tests. * more --------- Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
32 行
663 B
Python
32 行
663 B
Python
import io
|
|
import pickle
|
|
|
|
import dgl
|
|
|
|
import networkx as nx
|
|
import torch
|
|
|
|
|
|
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()
|