Graph partition based on balance_edge

[ZhengHongming888]

This pull request add balance_edge as one argument for graph partition based on metis.py. The original metis is just based on node_type for partitioning and not considered for edge_type specially for link prediction case with ddp. Here we used the in-degree for the edge_weight to replace the node_weight when we choose balance_edge as True.

The test results based on no balance_edge argument (same as old version) and balance_edge in ClusterData of cluster.py is

No balance_edge and based on default node_type:
----- cluster_data[0] = Data(x=[1104230], node_type=[1104230], edge_type=[21259372], edge_index=[2, 21259372]) ----
----- cluster_data[1] = Data(x=[1018995], node_type=[1018995], edge_type=[21806130], edge_index=[2, 21806130]) ----
----- cluster_data[2] = Data(x=[1086336], node_type=[1086336], edge_type=[24065724], edge_index=[2, 24065724]) ----
----- cluster_data[3] = Data(x=[1939568], node_type=[1939568], edge_type=[26776812], edge_index=[2, 26776812]) ----

balance_edge=True:
----- cluster_data[0] = Data(x=[1249728], node_type=[1249728], edge_type=[10912882], edge_index=[2, 10912882]) ----
----- cluster_data[1] = Data(x=[1249728], node_type=[1249728], edge_type=[14286552], edge_index=[2, 14286552]) ----
----- cluster_data[2] = Data(x=[1325624], node_type=[1325624], edge_type=[40316592], edge_index=[2, 40316592]) ----
----- cluster_data[3] = Data(x=[1324049], node_type=[1324049], edge_type=[33047668], edge_index=[2, 33047668]) ----

From the test results above, the edge_type column number is more balanced with balance_edge=True which is good for DDP link prediction like cases. The above test is based on Taobao.py dataset.

Any comments please let me know. thanks.

[codecov-commenter]

Codecov Report

Merging #309 (ff21933) into master (a980efd) will decrease coverage by 0.06%.
The diff coverage is 75.00%.

📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more

@@            Coverage Diff             @@
##           master     #309      +/-   ##
==========================================
- Coverage   72.50%   72.44%   -0.06%     
==========================================
  Files          29       29              
  Lines        1131     1136       +5     
==========================================
+ Hits          820      823       +3     
- Misses        311      313       +2     

Impacted Files	Coverage Δ
torch_sparse/metis.py	95.83% <75.00%> (-4.17%)	⬇️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more