Merge pull request #316 from corochann/relgcn_fix

use mlp for relgcn

Author: mottodora

examples/molnet/train_molnet.py

@@ -134,9 +134,9 @@ def set_up_predictor(method, n_unit, conv_layers, class_num):
     elif method == 'relgcn':
         print('Training an RelGCN predictor...')
         num_edge_type = 4
-        relgcn = RelGCN(out_channels=class_num, num_edge_type=num_edge_type,
+        relgcn = RelGCN(out_channels=n_unit, num_edge_type=num_edge_type,
                         scale_adj=True)
-        return GraphConvPredictor(relgcn, None)
+        return GraphConvPredictor(relgcn, mlp)
     elif method == 'relgat':
         print('Train Relational GAT model...')
         relgat = RelGAT(out_dim=n_unit, hidden_dim=n_unit,

examples/own_dataset/train_own_dataset.py

@@ -169,9 +169,9 @@ def set_up_predictor(method, n_unit, conv_layers, class_num):
     elif method == 'relgcn':
         print('Training an RelGCN predictor...')
         num_edge_type = 4
-        relgcn = RelGCN(out_channels=class_num, num_edge_type=num_edge_type,
+        relgcn = RelGCN(out_channels=n_unit, num_edge_type=num_edge_type,
                         scale_adj=True)
-        predictor = GraphConvPredictor(relgcn, None)
+        predictor = GraphConvPredictor(relgcn, mlp)
     elif method == 'relgat':
         print('Training an RelGAT predictor...')
         relgat = RelGAT(out_dim=n_unit, hidden_dim=n_unit,

examples/qm9/train_qm9.py

@@ -173,9 +173,9 @@ def set_up_predictor(method, n_unit, conv_layers, class_num, scaler):
     elif method == 'relgcn':
         print('Use Relational GCN predictor...')
         num_edge_type = 4
-        relgcn = RelGCN(out_channels=class_num, num_edge_type=num_edge_type,
+        relgcn = RelGCN(out_channels=n_unit, num_edge_type=num_edge_type,
                         scale_adj=True)
-        predictor = GraphConvPredictor(relgcn, None, scaler)
+        predictor = GraphConvPredictor(relgcn, mlp, scaler)
     elif method == 'relgat':
         print('Train Relational GAT predictor...')
         relgat = RelGAT(out_dim=n_unit, hidden_dim=n_unit,

examples/tox21/predictor.py

@@ -46,8 +46,9 @@ def build_predictor(method, n_unit, conv_layers, class_num):
         print('Use Relational GCN predictor...')
         num_edge_type = 4
         predictor = GraphConvPredictor(
-            RelGCN(out_channels=class_num, num_edge_type=num_edge_type,
-                   scale_adj=True))
+            RelGCN(out_channels=n_unit, num_edge_type=num_edge_type,
+                   scale_adj=True),
+            MLP(out_dim=class_num, hidden_dim=n_unit))
     elif method == 'relgat':
         print('Use GAT predictor...')
         predictor = GraphConvPredictor(