update layers.py

Author: cszhangzhen

README.md

@@ -36,12 +36,12 @@ More detailed information can be found [here](https://github.com/runningoat/hgmn
 ## Run
 Just execuate the following command for graph-graph classification task:
 ```
-python main_classification.py --datasets openssl_min50
+python main_classification.py --dataset openssl_min50
 ```
 
 Similarly, execuate the following command for graph-graph regression task:
 ```
-python main_regression.py --datasets AIDS700nef
+python main_regression.py --dataset AIDS700nef
 ```
 
 ## Citing

layers.py

@@ -5,7 +5,6 @@
 from torch.nn.parameter import Parameter
 from torch_geometric.nn import global_add_pool, global_mean_pool, HypergraphConv
 from torch_geometric.nn.pool.topk_pool import topk
-from torch_geometric.utils import dense_to_sparse
 from torch_scatter import scatter_add
 from torch_scatter import scatter
 from torch_geometric.utils import degree
@@ -272,7 +271,7 @@ def forward(self, x_left, batch_left, x_right, batch_right):
         # Construct batch fully connected graph in block diagonal matirx format
         for idx_i, idx_j, idx_x, idx_y in zip(shift_cum_num_nodes_x_left, cum_num_nodes_x_left, shift_cum_num_nodes_x_right, cum_num_nodes_x_right):
             adj[idx_i:idx_j, idx_x:idx_y] = 1.0
-        new_edge_index, _ = dense_to_sparse(adj)
+        new_edge_index, _ = self.dense_to_sparse(adj)
         row, col = new_edge_index
 
         assign_index1 = torch.stack([col, row], dim=0)
@@ -283,6 +282,12 @@ def forward(self, x_left, batch_left, x_right, batch_right):
 
         return out1, out2
 
+    def dense_to_sparse(self, adj):
+        assert adj.dim() == 2
+        index = adj.nonzero(as_tuple=False).t().contiguous()
+        value = adj[index[0], index[1]]
+        return index, value
+
 
 class ReadoutModule(torch.nn.Module):
     def __init__(self, args):