[KG] ComplEx score func for MXNet (#918)

* upd

* fig edgebatch edges

* add test

* trigger

* Update README.md for pytorch PinSage example.

Add noting that the PinSage model example under
example/pytorch/recommendation only work with Python 3.6+
as its dataset loader depends on stanfordnlp package
which work only with Python 3.6+.

* Provid a frame agnostic API to test nn modules on both CPU and CUDA side.

1. make dgl.nn.xxx frame agnostic
2. make test.backend include dgl.nn modules
3. modify test_edge_softmax of test/mxnet/test_nn.py and
    test/pytorch/test_nn.py work on both CPU and GPU

* Fix style

* Delete unused code

* Make agnostic test only related to tests/backend

1. clear all agnostic related code in dgl.nn
2. make test_graph_conv agnostic to cpu/gpu

* Fix code style

* fix

* doc

* Make all test code under tests.mxnet/pytorch.test_nn.py
work on both CPU and GPU.

* Fix syntex

* Remove rand

* Add TAGCN nn.module and example

* Now tagcn can run on CPU.

* Add unitest for TGConv

* Fix style

* For pubmed dataset, using --lr=0.005 can achieve better acc

* Fix style

* Fix some descriptions

* trigger

* Fix doc

* Add nn.TGConv and example

* Fix bug

* Update data in mxnet.tagcn test acc.

* Fix some comments and code

* delete useless code

* Fix namming

* Fix bug

* Fix bug

* Add test for mxnet TAGCov

* Add test code for mxnet TAGCov

* Update some docs

* Fix some code

* Update docs dgl.nn.mxnet

* Update weight init

* Fix

* reproduce the bug

* Fix concurrency bug reported at #755.
Also make test_shared_mem_store.py more deterministic.

* Update test_shared_mem_store.py

* Update dmlc/core

* Add complEx for mxnet

* ComplEx is ready for MXNet

apps/kg/config/best_config.sh

@@ -8,7 +8,7 @@ DGLBACKEND=pytorch python3 train.py --model DistMult --dataset FB15k --batch_siz
 
 DGLBACKEND=pytorch python3 train.py --model ComplEx --dataset FB15k --batch_size 1024 \
     --neg_sample_size 256 --hidden_dim 2000 --gamma 500.0 --lr 0.2 --max_step 100000 \
-    --batch_size_eval 16 --gpu 1 --valid --test -adv
+    --batch_size_eval 16 --gpu 0 --valid --test -adv
 
 DGLBACKEND=pytorch python3 train.py --model TransE --dataset FB15k --batch_size 1024 \
     --neg_sample_size 256 --hidden_dim 2000 --gamma 24.0 --lr 0.01 --max_step 20000 \

apps/kg/models/mxnet/score_fun.py

@@ -87,3 +87,60 @@ class DistMultScore(nn.Block):
                 tmp = (heads * relations).reshape(num_chunks, chunk_size, hidden_dim)
                 return nd.linalg_gemm2(tmp, tails)
             return fn
+
+class ComplExScore(nn.Block):
+    def __init__(self):
+        super(ComplExScore, self).__init__()
+
+    def edge_func(self, edges):
+        real_head, img_head = nd.split(edges.src['emb'], num_outputs=2, axis=-1)
+        real_tail, img_tail = nd.split(edges.dst['emb'], num_outputs=2, axis=-1)
+        real_rel, img_rel = nd.split(edges.data['emb'], num_outputs=2, axis=-1)
+
+        score = real_head * real_tail * real_rel \
+                + img_head * img_tail * real_rel \
+                + real_head * img_tail * img_rel \
+                - img_head * real_tail * img_rel
+        # TODO: check if there exists minus sign and if gamma should be used here(jin)
+        return {'score': nd.sum(score, -1)}
+
+    def reset_parameters(self):
+        pass
+
+    def save(self, path, name):
+        pass
+
+    def load(self, path, name):
+        pass
+
+    def forward(self, g):
+        g.apply_edges(lambda edges: self.edge_func(edges))
+
+    def create_neg(self, neg_head):
+        if neg_head:
+            def fn(heads, relations, tails, num_chunks, chunk_size, neg_sample_size):
+                hidden_dim = heads.shape[1]
+                emb_real, emb_img = nd.split(tails, num_outputs=2, axis=-1)
+                rel_real, rel_img = nd.split(relations, num_outputs=2, axis=-1)
+                real = emb_real * rel_real + emb_img * rel_img
+                img = -emb_real * rel_img + emb_img * rel_real
+                emb_complex = nd.concat(real, img, dim=-1)
+                tmp = emb_complex.reshape(num_chunks, chunk_size, hidden_dim)
+                heads = heads.reshape(num_chunks, neg_sample_size, hidden_dim)
+                heads = nd.transpose(heads, axes=(0, 2, 1))
+                return nd.linalg_gemm2(tmp, heads)
+            return fn
+        else:
+            def fn(heads, relations, tails, num_chunks, chunk_size, neg_sample_size):
+                hidden_dim = heads.shape[1]
+                emb_real, emb_img = nd.split(heads, num_outputs=2, axis=-1)
+                rel_real, rel_img = nd.split(relations, num_outputs=2, axis=-1)
+                real = emb_real * rel_real - emb_img * rel_img
+                img = emb_real * rel_img + emb_img * rel_real
+                emb_complex = nd.concat(real, img, dim=-1)
+                tmp = emb_complex.reshape(num_chunks, chunk_size, hidden_dim)
+
+                tails = tails.reshape(num_chunks, neg_sample_size, hidden_dim)
+                tails = nd.transpose(tails, axes=(0, 2, 1))
+                return nd.linalg_gemm2(tmp, tails)
+            return fn

apps/kg/tests/test_score.py

@@ -25,7 +25,8 @@ def generate_rand_graph(n):
     return g, entity_emb, rel_emb
 
 ke_score_funcs = {'TransE': TransEScore(12.0),
-                  'DistMult': DistMultScore()}
+                  'DistMult': DistMultScore(),
+                  'ComplEx': ComplExScore()}
 
 class BaseKEModel:
     def __init__(self, score_func, entity_emb, rel_emb):
@@ -98,6 +99,6 @@ def check_score_func(func_name):
 def test_score_func():
     for key in ke_score_funcs:
         check_score_func(key)
-    
+
 if __name__ == '__main__':
     test_score_func()