Repository: dmlc/dgl

Issue:
A potential performance issue of `ChebConv`.
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
https://github.com/dmlc/dgl/blob/24b7468fcce5d00d205d2b3e12a43a0a181b85e5/python/dgl/nn/pytorch/conv/chebconv.py#L134-L138
At Line 138 of chebconv.py, `Xt = th.cat((Xt, X_i), 1)` is executed `k` - 1 times to concatenate the `k` `x_i` tensors.
However, if the size of `DglGraph` or `k` is too large, the backend may need to allocate and copy memory multiple times, which may results in significant overhead.
If we collect `X_i` into a list, and call `th.cat` once after the loop, the backend would only need to allocate the memory once.
Since I am unsure if this is indeed a performance issue or a special design for certain corner cases, I have not created a Pull Request for this.

## To Reproduce

![image](https://user-images.githubusercontent.com/32944191/235428233-762fbf15-fc8c-4d0d-a8aa-bea0d50ee617.png)

From the chart above, the orange line stands for the current implementation of `ChebConv`, which is slower than the improved version.
My implementation is:
```diff
@@ -123,23 +123,25 @@ class ChebConv(nn.Module):
             re_norm = 2.0 / lambda_max

             # X_0 is the raw feature, Xt refers to the concatenation of X_0, X_1, ... X_t
-            Xt = X_0 = feat
+            X_0 = feat
+            Xt = [X_0]

             # X_1(f)
             if self._k > 1:
                 h = unnLaplacian(X_0, D_invsqrt, graph)
                 X_1 = -re_norm * h + X_0 * (re_norm - 1)
                 # Concatenate Xt and X_1
-                Xt = th.cat((Xt, X_1), 1)
+                Xt.append(X_1)

             # Xi(x), i = 2...k
             for _ in range(2, self._k):
                 h = unnLaplacian(X_1, D_invsqrt, graph)
                 X_i = -2 * re_norm * h + X_1 * 2 * (re_norm - 1) - X_0
                 # Concatenate Xt and X_i
-                Xt = th.cat((Xt, X_i), 1)
+                Xt.append(X_i)
                 X_1, X_0 = X_i, X_1

+            Xt = th.cat(Xt, dim=1)
             # linear projection
             h = self.linear(Xt)
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

 - DGL Version (e.g., 1.0): 0.8.2
     According to the latest source code, this issue still exists. https://github.com/dmlc/dgl/blob/24b7468fcce5d00d205d2b3e12a43a0a181b85e5/python/dgl/nn/pytorch/conv/chebconv.py#L138
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.8.1
 - OS (e.g., Linux): Ubuntu 18.04
 - How you installed DGL (`conda`, `pip`, source): conda
 - ~~Build command you used (if compiling from source):~~
 - Python version: 3.8.5
 - ~~CUDA/cuDNN version (if applicable):~~
 - ~~GPU models and configuration (e.g. V100):~~
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Similar issue exists for `TensorFlow` version of `ChebConv`
https://github.com/dmlc/dgl/blob/24b7468fcce5d00d205d2b3e12a43a0a181b85e5/python/dgl/nn/tensorflow/conv/chebconv.py#L141-L145
Make sense! Thanks for catching this. Please go ahead and create a PR for this.

Resolving patch:
diff --git a/python/dgl/nn/pytorch/conv/chebconv.py b/python/dgl/nn/pytorch/conv/chebconv.py
--- a/python/dgl/nn/pytorch/conv/chebconv.py
+++ b/python/dgl/nn/pytorch/conv/chebconv.py
@@ -120,24 +120,28 @@ def unnLaplacian(feat, D_invsqrt, graph):
             lambda_max = broadcast_nodes(graph, lambda_max)
             re_norm = 2.0 / lambda_max
 
-            # X_0 is the raw feature, Xt refers to the concatenation of X_0, X_1, ... X_t
-            Xt = X_0 = feat
+            # X_0 is the raw feature, Xt is the list of X_0, X_1, ... X_t
+            X_0 = feat
+            Xt = [X_0]
 
             # X_1(f)
             if self._k > 1:
                 h = unnLaplacian(X_0, D_invsqrt, graph)
                 X_1 = -re_norm * h + X_0 * (re_norm - 1)
-                # Concatenate Xt and X_1
-                Xt = th.cat((Xt, X_1), 1)
+                # Append X_1 to Xt
+                Xt.append(X_1)
 
             # Xi(x), i = 2...k
             for _ in range(2, self._k):
                 h = unnLaplacian(X_1, D_invsqrt, graph)
                 X_i = -2 * re_norm * h + X_1 * 2 * (re_norm - 1) - X_0
-                # Concatenate Xt and X_i
-                Xt = th.cat((Xt, X_i), 1)
+                # Add X_1 to Xt
+                Xt.append(X_i)
                 X_1, X_0 = X_i, X_1
 
+            # Create the concatenation
+            Xt = th.cat(Xt, dim=1)
+
             # linear projection
             h = self.linear(Xt)
 
diff --git a/python/dgl/nn/tensorflow/conv/chebconv.py b/python/dgl/nn/tensorflow/conv/chebconv.py
--- a/python/dgl/nn/tensorflow/conv/chebconv.py
+++ b/python/dgl/nn/tensorflow/conv/chebconv.py
@@ -127,24 +127,28 @@ def unnLaplacian(feat, D_invsqrt, graph):
             lambda_max = broadcast_nodes(graph, lambda_max)
             re_norm = 2.0 / lambda_max
 
-            # X_0 is the raw feature, Xt refers to the concatenation of X_0, X_1, ... X_t
-            Xt = X_0 = feat
+            # X_0 is the raw feature, Xt is the list of X_0, X_1, ... X_t
+            X_0 = feat
+            Xt = [X_0]
 
             # X_1(f)
             if self._k > 1:
                 h = unnLaplacian(X_0, D_invsqrt, graph)
                 X_1 = -re_norm * h + X_0 * (re_norm - 1)
-                # Concatenate Xt and X_1
-                Xt = tf.concat((Xt, X_1), 1)
+                # Append X_1 to Xt
+                Xt.append(X_1)
 
             # Xi(x), i = 2...k
             for _ in range(2, self._k):
                 h = unnLaplacian(X_1, D_invsqrt, graph)
                 X_i = -2 * re_norm * h + X_1 * 2 * (re_norm - 1) - X_0
-                # Concatenate Xt and X_i
-                Xt = tf.concat((Xt, X_i), 1)
+                # Append X_i to Xt
+                Xt.append(X_i)
                 X_1, X_0 = X_i, X_1
 
+            # Create the concatenation
+            Xt = tf.concat(Xt, 1)
+
             # linear projection
             h = self.linear(Xt)
 


Test patch: