Repository: dmlc/dgl

Issue:
[BatchedHeteroGraph] Attributes Missing After Batching with num_workers > 1
## 🐛 Bug

When we set `num_workers` to be greater than 1 in batching DGLHeteroGraphs, attributes `_batch_size`, `_batch_num_nodes` and `_batch_num_edges` are lost.

## To Reproduce

```python
import dgl
from torch.utils.data import DataLoader

class ToyData(object):
    def __init__(self):
        self.graphs = [dgl.graph([(0, 1), (1, 2)], 'user', 'follows') for _ in range(2)]

    def __getitem__(self, index):
        return self.graphs[index]

    def __len__(self):
        return len(self.graphs)

def collate(data):
    return dgl.batch_hetero(data)

dataset = ToyData()
# Case when num_workers=0
data_loader = DataLoader(dataset, batch_size=2, collate_fn=collate, num_workers=0)
for bg in data_loader:
    print(bg.batch_size)
    print(bg.batch_num_nodes('user'))
    print(bg.batch_num_edges('follows'))
# Case when num_workers=2
data_loader = DataLoader(dataset, batch_size=2, collate_fn=collate, num_workers=2)
for bg in data_loader:
    print(bg.batch_size)
    print(bg.batch_num_nodes('user'))
    print(bg.batch_num_edges('follows'))
```

## Environment

 - DGL Version (e.g., 1.0): DGL Nightly
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
 - OS (e.g., Linux):
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:



Discussion:

Resolving patch:
diff --git a/python/dgl/batched_heterograph.py b/python/dgl/batched_heterograph.py
--- a/python/dgl/batched_heterograph.py
+++ b/python/dgl/batched_heterograph.py
@@ -302,6 +302,14 @@ def to(self, ctx, **kwargs):  # pylint: disable=invalid-name
                                      batch_num_nodes=self._batch_num_nodes,
                                      batch_num_edges=self._batch_num_edges)
 
+    def __getstate__(self):
+        state = super().__getstate__()
+        return state, self._batch_size, self._batch_num_nodes, self._batch_num_edges
+
+    def __setstate__(self, state):
+        state, self._batch_size, self._batch_num_nodes, self._batch_num_edges = state
+        super().__setstate__(state)
+
 def unbatch_hetero(graph):
     """Return the list of heterographs in this batch.
 


Test patch:
diff --git a/tests/compute/test_pickle.py b/tests/compute/test_pickle.py
--- a/tests/compute/test_pickle.py
+++ b/tests/compute/test_pickle.py
@@ -78,6 +78,13 @@ def _assert_is_identical_batchedgraph(bg1, bg2):
     assert bg1.batch_num_nodes == bg2.batch_num_nodes
     assert bg1.batch_num_edges == bg2.batch_num_edges
 
+def _assert_is_identical_batchedhetero(bg1, bg2):
+    _assert_is_identical_hetero(bg1, bg2)
+    for ntype in bg1.ntypes:
+        assert bg1.batch_num_nodes(ntype) == bg2.batch_num_nodes(ntype)
+    for canonical_etype in bg1.canonical_etypes:
+        assert bg1.batch_num_edges(canonical_etype) == bg2.batch_num_edges(canonical_etype)
+
 def _assert_is_identical_index(i1, i2):
     assert i1.slice_data() == i2.slice_data()
     assert F.array_equal(i1.tousertensor(), i2.tousertensor())
@@ -258,6 +265,33 @@ def test_pickling_heterograph():
     new_g = _reconstruct_pickle(g)
     _assert_is_identical_hetero(g, new_g)
 
+def test_pickling_batched_heterograph():
+    # copied from test_heterograph.create_test_heterograph()
+    plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
+    wishes_nx = nx.DiGraph()
+    wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
+    wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
+    wishes_nx.add_edge('u0', 'g1', id=0)
+    wishes_nx.add_edge('u2', 'g0', id=1)
+
+    follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
+    plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game')
+    wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game')
+    develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')
+    g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
+    g2 = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
+
+    g.nodes['user'].data['u_h'] = F.randn((3, 4))
+    g.nodes['game'].data['g_h'] = F.randn((2, 5))
+    g.edges['plays'].data['p_h'] = F.randn((4, 6))
+    g2.nodes['user'].data['u_h'] = F.randn((3, 4))
+    g2.nodes['game'].data['g_h'] = F.randn((2, 5))
+    g2.edges['plays'].data['p_h'] = F.randn((4, 6))
+
+    bg = dgl.batch_hetero([g, g2])
+    new_bg = _reconstruct_pickle(bg)
+    _assert_is_identical_batchedhetero(bg, new_bg)
+
 @unittest.skipIf(dgl.backend.backend_name != "pytorch", reason="Only test for pytorch format file")
 def test_pickling_heterograph_index_compatibility():
     plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
@@ -287,3 +321,4 @@ def test_pickling_heterograph_index_compatibility():
     test_pickling_nodeflow()
     test_pickling_batched_graph()
     test_pickling_heterograph()
+    test_pickling_batched_heterograph()