Repository: dmlc/dgl

Issue:
[Bug] g.to(gpu) break the g._node_frames
## 🐛 Bug
How to reproduce:

```
>>> import dgl
>>> import torch as th
>>> g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays', 'game')
>>> g.nodes['user'].data['h'] = th.tensor([0,1])
Success
```
```
>>> import dgl
>>> import torch as th
>>> g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays', 'game')
>>> g = g.to(th.device('cuda:0'))
>>> g.nodes['user'].data['h'] = th.tensor([0,1]).cuda()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/ubuntu/work/xiangsx/dgl/python/dgl/view.py", line 317, in __setitem__
    self._graph._set_n_repr(self._ntid, self._nodes, {key : val})
  File "/home/ubuntu/work/xiangsx/dgl/python/dgl/heterograph.py", line 2750, in _set_n_repr
    self._node_frames[ntid][key] = val
  File "/home/ubuntu/work/xiangsx/dgl/python/dgl/frame.py", line 728, in __setitem__
    self.update_data(key, val, inplace=False)
  File "/home/ubuntu/work/xiangsx/dgl/python/dgl/frame.py", line 755, in update_data
    self.update_column(key, val, inplace=inplace)
  File "/home/ubuntu/work/xiangsx/dgl/python/dgl/frame.py", line 786, in update_column
    self._frame[name] = data
  File "/home/ubuntu/work/xiangsx/dgl/python/dgl/frame.py", line 343, in __setitem__
    self.update_column(name, data)
  File "/home/ubuntu/work/xiangsx/dgl/python/dgl/frame.py", line 430, in update_column
    (self.num_rows, len(col)))
dgl._ffi.base.DGLError: Expected data to have 0 rows, got 2.
```

## Expected behavior

No crash happens.

## Environment

 - DGL Version (e.g., 1.0): master
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): pytorch 1.5.1
 - OS (e.g., Linux): 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:

## Additional context

<!-- Add any other context about the problem here. -->



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
@@ -284,11 +284,11 @@ def to(self, ctx, **kwargs):  # pylint: disable=invalid-name
         new_nframes = []
         for nframe in self._node_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in nframe.items()}
-            new_nframes.append(FrameRef(Frame(new_feats)))
+            new_nframes.append(FrameRef(Frame(new_feats, num_rows=nframe.num_rows)))
         new_eframes = []
         for eframe in self._edge_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in eframe.items()}
-            new_eframes.append(FrameRef(Frame(new_feats)))
+            new_eframes.append(FrameRef(Frame(new_feats, num_rows=eframe.num_rows)))
         # TODO(minjie): replace the following line with the commented one to enable GPU graph.
         new_gidx = self._graph
         #new_gidx = self._graph.copy_to(utils.to_dgl_context(ctx))
diff --git a/python/dgl/data/heterograph_serialize.py b/python/dgl/data/heterograph_serialize.py
--- a/python/dgl/data/heterograph_serialize.py
+++ b/python/dgl/data/heterograph_serialize.py
@@ -49,12 +49,12 @@ def get_graph(self):
         gidx = _CAPI_GetGindexFromHeteroGraphData(self)
         nframes = []
         eframes = []
-        for ntensor in ntensor_list:
+        for ntid, ntensor in enumerate(ntensor_list):
             ndict = {ntensor[i]: F.zerocopy_from_dgl_ndarray(ntensor[i+1]) for i in range(0, len(ntensor), 2)}
-            nframes.append(FrameRef(Frame(ndict)))
+            nframes.append(FrameRef(Frame(ndict, num_rows=gidx.number_of_nodes(ntid))))
         
-        for etensor in etensor_list:
+        for etid, etensor in enumerate(etensor_list):
             edict = {etensor[i]: F.zerocopy_from_dgl_ndarray(etensor[i+1]) for i in range(0, len(etensor), 2)}
-            eframes.append(FrameRef(Frame(edict)))
+            eframes.append(FrameRef(Frame(edict, num_rows=gidx.number_of_edges(etid))))
         
         return DGLHeteroGraph(gidx, ntype_names, etype_names, nframes, eframes)
diff --git a/python/dgl/heterograph.py b/python/dgl/heterograph.py
--- a/python/dgl/heterograph.py
+++ b/python/dgl/heterograph.py
@@ -4316,11 +4316,11 @@ def to(self, ctx, **kwargs):  # pylint: disable=invalid-name
         new_nframes = []
         for nframe in self._node_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in nframe.items()}
-            new_nframes.append(FrameRef(Frame(new_feats)))
+            new_nframes.append(FrameRef(Frame(new_feats, num_rows=nframe.num_rows)))
         new_eframes = []
         for eframe in self._edge_frames:
             new_feats = {k : F.copy_to(feat, ctx) for k, feat in eframe.items()}
-            new_eframes.append(FrameRef(Frame(new_feats)))
+            new_eframes.append(FrameRef(Frame(new_feats, num_rows=eframe.num_rows)))
         # TODO(minjie): replace the following line with the commented one to enable GPU graph.
         new_gidx = self._graph
         #new_gidx = self._graph.copy_to(utils.to_dgl_context(ctx))


Test patch:
diff --git a/tests/compute/test_batched_heterograph.py b/tests/compute/test_batched_heterograph.py
--- a/tests/compute/test_batched_heterograph.py
+++ b/tests/compute/test_batched_heterograph.py
@@ -402,9 +402,22 @@ def test_to_device(index_dtype):
         assert bg.batch_num_nodes('user') == bg1.batch_num_nodes('user')
         assert bg.batch_num_edges('plays') == bg1.batch_num_edges('plays')
 
+    # set feature
+    g1 = dgl.heterograph({
+        ('user', 'plays', 'game'): [(0, 0), (1, 1)]
+    }, index_dtype=index_dtype)
+    g2 = dgl.heterograph({
+        ('user', 'plays', 'game'): [(0, 0), (1, 0)]
+    }, index_dtype=index_dtype)
+    bg = dgl.batch_hetero([g1, g2])
+    if F.is_cuda_available():
+        bg1 = bg.to(F.cuda())
+        bg1.nodes['user'].data['test'] = F.copy_to(F.tensor([0,1,2,3]), F.cuda())
+        bg1.edata['test'] = F.copy_to(F.tensor([0,1,2,3]), F.cuda())
+
 if __name__ == '__main__':
     test_batching_hetero_topology('int32')
     test_batching_hetero_and_batched_hetero_topology('int32')
     test_batched_features('int32')
     test_batching_with_zero_nodes_edges('int32')
-    # test_to_device()
+    test_to_device('int32')
diff --git a/tests/compute/test_heterograph.py b/tests/compute/test_heterograph.py
--- a/tests/compute/test_heterograph.py
+++ b/tests/compute/test_heterograph.py
@@ -885,6 +885,15 @@ def test_to_device(index_dtype):
         g1 = g.to(F.cuda())
         assert g1 is not None
 
+    # set feature after g.to
+    g = create_test_heterograph(index_dtype)
+    if F.is_cuda_available():
+        g1 = g.to(F.cuda())
+        assert g1 is not None
+        g1.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cuda())
+        g1.nodes['game'].data['i'] = F.copy_to(F.ones((2, 5)), F.cuda())
+        g1.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cuda())
+
 @parametrize_dtype
 def test_convert_bound(index_dtype):
     def _test_bipartite_bound(data, card):
@@ -2068,7 +2077,7 @@ def test_reverse(index_dtype):
     # test_flatten()
     # test_convert_bound()
     # test_convert()
-    # test_to_device()
+    # test_to_device("int32")
     # test_transform("int32")
     # test_subgraph("int32")
     # test_subgraph_mask("int32")
diff --git a/tests/compute/test_serialize.py b/tests/compute/test_serialize.py
--- a/tests/compute/test_serialize.py
+++ b/tests/compute/test_serialize.py
@@ -233,7 +233,7 @@ def create_heterographs2(index_dtype):
     g_y.nodes['user'].data['hh'] = F.ones((4, 5))
     g_y.edges['knows'].data['ww'] = F.randn((2, 10))
     g = dgl.hetero_from_relations([g_x, g_y, g_z])
-    return [g, g_x, g_y]
+    return [g, g_x, g_y, g_z]
 
 def test_deserialize_old_heterograph_file():
     path = os.path.join(
@@ -273,11 +273,11 @@ def test_serialize_heterograph():
             assert g_list[i].canonical_etypes[j] == etypes
     assert g_list[1].restrict_format() == 'any'
     assert g_list[2].restrict_format() == 'csr'
-    assert g_list[3].idtype == F.int32
+    assert g_list[4].idtype == F.int32
     assert np.allclose(
         F.asnumpy(g_list[2].nodes['user'].data['hh']), np.ones((4, 5)))
     assert np.allclose(
-        F.asnumpy(g_list[5].nodes['user'].data['hh']), np.ones((4, 5)))
+        F.asnumpy(g_list[6].nodes['user'].data['hh']), np.ones((4, 5)))
     edges = g_list[0]['follows'].edges()
     assert np.allclose(F.asnumpy(edges[0]), np.array([0, 1, 2]))
     assert np.allclose(F.asnumpy(edges[1]), np.array([1, 2, 3]))
@@ -285,6 +285,10 @@ def test_serialize_heterograph():
         assert g_list[i].ntypes == g_list0[i].ntypes
         assert g_list[i].etypes == g_list0[i].etypes
 
+    # test set feature after load_graph
+    g_list[3].nodes['user'].data['test'] = F.tensor([0, 1, 2, 4])
+    g_list[3].edata['test'] = F.tensor([0, 1, 2])
+
     os.unlink(path)
 
 @pytest.mark.skip(reason="lack of permission on CI")
@@ -308,17 +312,17 @@ def test_serialize_heterograph_s3():
 
 if __name__ == "__main__":
     pass
-    test_graph_serialize_with_feature(True)
-    test_graph_serialize_with_feature(False)
-    test_graph_serialize_without_feature(True)
-    test_graph_serialize_without_feature(False)
-    test_graph_serialize_with_labels(True)
-    test_graph_serialize_with_labels(False)
-    test_serialize_tensors()
-    test_serialize_empty_dict()
-    test_load_old_files1()
-    test_load_old_files2()
-    test_serialize_heterograph()
-    # test_serialize_heterograph_s3()
-    test_deserialize_old_heterograph_file()
-    # create_old_heterograph_files()
+    #test_graph_serialize_with_feature(True)
+    #test_graph_serialize_with_feature(False)
+    #test_graph_serialize_without_feature(True)
+    #test_graph_serialize_without_feature(False)
+    #test_graph_serialize_with_labels(True)
+    #test_graph_serialize_with_labels(False)
+    #test_serialize_tensors()
+    #test_serialize_empty_dict()
+    #test_load_old_files1()
+    #test_load_old_files2()
+    #test_serialize_heterograph()
+    #test_serialize_heterograph_s3()
+    #test_deserialize_old_heterograph_file()
+    #create_old_heterograph_files()