Repository: pyg-team/pytorch_geometric
Pull request #7942 (2023-08-28 08:47:29)

Filter out empty tensors inside `trim_to_layer`
filter out empty tensor for `x`, `edge_index`, `edge_attr` after calling `trim_to_layer` function. This can avoid unnecessary computation when some node/edge types get empty output. For example: when I train `igbh-tiny` dataset with 3 hops sampler and use `trim_to_layer` function, I get a lot of empty edge_index tensor for edge type '('author', 'affiliated_to', 'institute')', but the feature tensor for 'author' node type is still sent to compute in `HeteroConv` implementation.

rusty1s: Can we add an additional test for this? Do you have some benchmark on the speed-up?
kaixuanliu: > Can we add an additional test for this? Do you have some benchmark on the speed-up?

Sure, will add related test case and benchmark data afterwards.
kaixuanliu: > Can we add an additional test for this? Do you have some benchmark on the speed-up?

I tried to use examples/hetero/hierarchical_sage.py to do benchmark, but as the dataset(ogbn-mags) is not big enough, so there is no obvious speed up. However, when I changed to larger dataset(igbh-medium), which is called from GLTorch, the time consumption drops from 7:23:39 to 6:43:24.

andreazanetti: Hi @kaixuanliu @rusty1s , I have some questions about this. 
I have been working on the `--trim` support for the homo case, whereas on its hetero-case implications I am still not 100% crystal clear on all cases, and I am developing my understanding :)
So I would really appreciate if I could get clarified about the following:

a) Are the two cases tested (with and without filtering-out) numerically equivalent? I suppose they should.

b) Is there any chance of corner cases for which, after some 'trimming' only one kind of tensor is empty (say `edge_index)` but not another (say `x`) is not? Would that situation be problematic once the removal happens only for one of the 2?

c) Is it reasonable to think that the performance gain may be connected to the size AND the topology of the graph? 

d) Are any drop in performance possible (in some scenarios) for the additional calls (and for loops) you added?

Thanks very much!


kaixuanliu: Hi @andreazanetti , I try to answer your questions. And @rusty1s maybe you can help correct my answer if I have some mis-understandings.
> a) Are the two cases tested (with and without filtering-out) numerically equivalent? I suppose they should.
> Yes. I tried my best to keep them numerically equivalent when do the testing, except that there is still some difference because of the randomness during sampling. However, I think the randomness cannot cause above 8% perf difference for igbh-medium dataset. 
> b) Is there any chance of corner cases for which, after some 'trimming' only one kind of tensor is empty (say `edge_index)` but not another (say `x` is not)? Would that situation be problematic once the removal happens only for one of the 2?
> There may be the case that some edge type's edge index is empty, but one of the two related node side feature(src or dst ) is not empty. But I think it should not be a problem to remove them.
> c) Is it reasonable to think that the performance gain may be connected to the size AND the topology of the graph?
> Yes. It is related to the size and topology of the graph, as well as the num_neighbors and hops of sampler.
> d) Are any drop in performance possible (in some scenarios) for the additional calls (and for loops) added here?
> The additional calls is negligible, they are just called for each layer and the for loop is for the node/edge type. I think we can ignore it.



andreazanetti: Apologies for the delay in my reply. I was off sick. 

_a) Are the two cases tested (with and without filtering-out) numerically equivalent? I suppose they should.
Yes. I tried my best to keep them numerically equivalent when do the testing, except that there is still some difference because of the randomness during sampling. However, I think the randomness cannot cause above 8% perf difference for igbh-medium dataset._
a) AZ: Could you clarify how do you estimate the 8%? thanks


_b) Is there any chance of corner cases for which, after some 'trimming' only one kind of tensor is empty (say edge_index) but not another (say x is not)? Would that situation be problematic once the removal happens only for one of the 2?
There may be the case that some edge type's edge index is empty, but one of the two related node side feature(src or dst ) is not empty. But I think it should not be a problem to remove them._
b) AZ: Could you clarify how this removal would happen? I am not sure it is crystal clear to me, apologies. Also, I was wondering if those cases and their removal should be handled in this PR.

_c) Is it reasonable to think that the performance gain may be connected to the size AND the topology of the graph?
Yes. It is related to the size and topology of the graph, as well as the num_neighbors and hops of sampler._
c) AZ: Makes sense to me, thanks

_d) Are any drop in performance possible (in some scenarios) for the additional calls (and for loops) added here?
The additional calls is negligible, they are just called for each layer and the for loop is for the node/edge type. I think we can ignore it._
d) AZ: Intuitively makes sense to me. Some data from tests would be great to show this however. Also the removal of the corner cases mentioned above might add some additional checks, I guess.
kaixuanliu: > Apologies for the delay in my reply. I was off sick.
> 
> _a) Are the two cases tested (with and without filtering-out) numerically equivalent? I suppose they should. Yes. I tried my best to keep them numerically equivalent when do the testing, except that there is still some difference because of the randomness during sampling. However, I think the randomness cannot cause above 8% perf difference for igbh-medium dataset._ a) AZ: Could you clarify how do you estimate the 8%? thanks
> **Yes. I did performance benchmark using this [igbh example](https://github.com/alibaba/graphlearn-for-pytorch/tree/main/examples/igbh) from gltorch, which utilize pyg's `trim_to_layer` function. The command line is like this: `numactl -C 0-55 -m 0,1,2,3 python train_rgnn.py --dataset_size='medium' --path /lfs/lfs14/kaixuan/igbh/ --with_trim=True --epochs 2`, and the original time consumption is 7:23:39, after applying this patch, time consumption drops to 6:45:32, hence get ~8% perf improvement.**

> _b) Is there any chance of corner cases for which, after some 'trimming' only one kind of tensor is empty (say edge_index) but not another (say x is not)? Would that situation be problematic once the removal happens only for one of the 2? There may be the case that some edge type's edge index is empty, but one of the two related node side feature(src or dst ) is not empty. But I think it should not be a problem to remove them._ b) AZ: Could you clarify how this removal would happen? I am not sure it is crystal clear to me, apologies. Also, I was wondering if those cases and their removal should be handled in this PR.
> **For example, before applying this patch, one output after `trim_to_layer` function is `edge type:('institute', 'rev_affiliated_to', 'author'), edge shape:torch.Size([2, 0]), src shape:torch.Size([0, 128]), dst shape:torch.Size([19723, 128])`, and this output will still be sent to [HeteroConv class](https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/hetero_conv.py#L153) to do convolution computation, which is not necessary. After applying this patch, this edge type will be removed directly, this can avoid unnecessary computation.**

> _c) Is it reasonable to think that the performance gain may be connected to the size AND the topology of the graph? Yes. It is related to the size and topology of the graph, as well as the num_neighbors and hops of sampler._ c) AZ: Makes sense to me, thanks
> 
> _d) Are any drop in performance possible (in some scenarios) for the additional calls (and for loops) added here? The additional calls is negligible, they are just called for each layer and the for loop is for the node/edge type. I think we can ignore it._ d) AZ: Intuitively makes sense to me. Some data from tests would be great to show this however. Also the removal of the corner cases mentioned above might add some additional checks, I guess.
**I also tested some other dataset include igbh-tiny, igbh-small ,ogbn-mags, there are no performance drop in all these datasets. For example, if we have a 3-layer model, the additional calls are invoked only 3 times for each batch, which can be ignored in the performance perspective as I think. For the corner cases, I think if the edge_index is empty, it means there is no link for specified edge type, and we don't have to process the node feature related to it.**


kaixuanliu: @rusty1s @andreazanetti @rfdavid , Hi, I added related unit test, but seems there is sth wrong with test/loader/test_imbalanced_sampler.py file, so cannot get all test cases passed.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7942?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7942](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7942?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (cf9d282) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/60a2948ba1745e35ee5209d5e3d0b861aca54495?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (60a2948) will **increase** coverage by `0.00%`.
> Report is 2 commits behind head on master.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #7942   +/-   ##
=======================================
  Coverage   89.57%   89.58%           
=======================================
  Files         460      460           
  Lines       26967    26981   +14     
=======================================
+ Hits        24156    24170   +14     
  Misses       2811     2811           
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7942?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/transforms/random\_link\_split.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7942?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvcmFuZG9tX2xpbmtfc3BsaXQucHk=) | `97.45% <ø> (ø)` | |
| [torch\_geometric/nn/aggr/base.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7942?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2FnZ3IvYmFzZS5weQ==) | `94.54% <100.00%> (ø)` | |
| [torch\_geometric/sampler/neighbor\_sampler.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7942?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvbmVpZ2hib3Jfc2FtcGxlci5weQ==) | `93.00% <100.00%> (+0.09%)` | :arrow_up: |
| [torch\_geometric/utils/trim\_to\_layer.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7942?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3RyaW1fdG9fbGF5ZXIucHk=) | `90.76% <100.00%> (+1.67%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

kaixuanliu: @rusty1s @wsad1 , Hi, do you have other comments or suggestions for this PR? Seems it needs another approval before merge.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 18f974b42..6b0caa744 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -96,6 +96,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Changed the `trim_to_layer` function to filter out non-reachable node and edge types when operating on heterogeneous graphs ([#7942](https://github.com/pyg-team/pytorch_geometric/pull/7942))
 - Accelerated and simplified `top_k` computation in `TopKPooling` ([#7737](https://github.com/pyg-team/pytorch_geometric/pull/7737))
 - Updated `GIN` implementation in kernel benchmarks to have sequential batchnorms ([#7955](https://github.com/pyg-team/pytorch_geometric/pull/7955))
 - Fixed bugs in benchmarks caused by a lack of the device conditions for CPU and unexpected `cache` argument in heterogeneous  models ([#7956](https://github.com/pyg-team/pytorch_geometric/pull/7956)
diff --git a/torch_geometric/utils/trim_to_layer.py b/torch_geometric/utils/trim_to_layer.py
index bd6a2452f..1efdd96d1 100644
--- a/torch_geometric/utils/trim_to_layer.py
+++ b/torch_geometric/utils/trim_to_layer.py
@@ -1,4 +1,5 @@
-from typing import Dict, List, Optional, Tuple, Union
+import copy
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -14,6 +15,17 @@ from torch_geometric.typing import (
 )
 
 
+def filter_empty_entries(
+        input_dict: Dict[Union[Any], Tensor]) -> Dict[Any, Tensor]:
+    r"""Removes empty tensors from a dictionary. This avoids unnecessary
+    computation when some node/edge types are non-reachable after trimming."""
+    out_dict = copy.copy(input_dict)
+    for key, value in input_dict.items():
+        if value.numel() == 0:
+            del out_dict[key]
+    return out_dict
+
+
 def trim_to_layer(
     layer: int,
     num_sampled_nodes_per_hop: Union[List[int], Dict[NodeType, List[int]]],
@@ -53,6 +65,8 @@ def trim_to_layer(
             k: trim_feat(v, layer, num_sampled_nodes_per_hop[k])
             for k, v in x.items()
         }
+        x = filter_empty_entries(x)
+
         edge_index = {
             k:
             trim_adj(
@@ -64,11 +78,15 @@ def trim_to_layer(
             )
             for k, v in edge_index.items()
         }
+        edge_index = filter_empty_entries(edge_index)
+
         if edge_attr is not None:
             edge_attr = {
                 k: trim_feat(v, layer, num_sampled_edges_per_hop[k])
                 for k, v in edge_attr.items()
             }
+            edge_attr = filter_empty_entries(edge_attr)
+
         return x, edge_index, edge_attr
 
     x = trim_feat(x, layer, num_sampled_nodes_per_hop)


Test patch:
diff --git a/test/utils/test_trim_to_layer.py b/test/utils/test_trim_to_layer.py
index 5ee789feb..c82960084 100644
--- a/test/utils/test_trim_to_layer.py
+++ b/test/utils/test_trim_to_layer.py
@@ -197,3 +197,39 @@ def test_trim_to_layer_with_neighbor_loader():
     assert out2.size() == (2, 16)
 
     assert torch.allclose(out1, out2)
+
+
+def test_trim_to_layer_filtering():
+    x_dict = {
+        'paper': torch.rand((13, 128)),
+        'author': torch.rand((5, 128)),
+        'field_of_study': torch.rand((6, 128))
+    }
+    edge_index_dict = {
+        ('author', 'writes', 'paper'):
+        torch.tensor([[0, 1, 2, 3, 4], [0, 0, 1, 2, 2]]),
+        ('paper', 'has_topic', 'field_of_study'):
+        torch.tensor([[6, 7, 8, 9], [0, 0, 1, 1]])
+    }
+    num_sampled_nodes_dict = {
+        'paper': [1, 2, 10],
+        'author': [0, 2, 3],
+        'field_of_study': [0, 2, 4]
+    }
+    num_sampled_edges_dict = {
+        ('author', 'writes', 'paper'): [2, 3],
+        ('paper', 'has_topic', 'field_of_study'): [0, 4]
+    }
+    x_dict, edge_index_dict, _ = trim_to_layer(
+        layer=1,
+        num_sampled_nodes_per_hop=num_sampled_nodes_dict,
+        num_sampled_edges_per_hop=num_sampled_edges_dict,
+        x=x_dict,
+        edge_index=edge_index_dict,
+    )
+    assert list(edge_index_dict.keys()) == [('author', 'writes', 'paper')]
+    assert torch.equal(edge_index_dict[('author', 'writes', 'paper')],
+                       torch.tensor([[0, 1], [0, 0]]))
+    assert x_dict['paper'].size() == (3, 128)
+    assert x_dict['author'].size() == (2, 128)
+    assert x_dict['field_of_study'].size() == (2, 128)