loss.backward() is too slow

[S-atan]

🐛 Describe the bug

Hello, I replaced the ordinary convolution operation with a graph convolution operation for the skeleton joint point data. The graph convolution operation takes up to 20 seconds in the loss.backward() step, while the ordinary convolution operation only requires 0.2 seconds. Is there any method to accelerate this process?

Versions

PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 专业版
GCC version: (i686-posix-sjlj, built by strawberryperl.com project) 4.9.2
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: N/A

Python version: 3.12.2 | packaged by Anaconda, Inc. | (main, Feb 27 2024, 17:28:07) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: 10.1.243
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 2080 Ti

Nvidia driver version: 536.23
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Revision=21764

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.2.1+cu121
[pip3] torch_cluster==1.6.3+pt22cu121
[pip3] torch_geometric==2.5.2
[pip3] torch_scatter==2.1.2+pt22cu121
[pip3] torch_sparse==0.6.18+pt22cu121
[pip3] torch_spline_conv==1.2.2+pt22cu121
[pip3] torchaudio==2.2.1+cu121
[pip3] torchvision==0.17.1+cu121
[conda] numpy 1.26.4 pypi_0 pypi
[conda] torch 2.2.1+cu121 pypi_0 pypi
[conda] torch-cluster 1.6.3+pt22cu121 pypi_0 pypi
[conda] torch-geometric 2.5.2 pypi_0 pypi
[conda] torch-scatter 2.1.2+pt22cu121 pypi_0 pypi
[conda] torch-sparse 0.6.18+pt22cu121 pypi_0 pypi
[conda] torch-spline-conv 1.2.2+pt22cu121 pypi_0 pypi
[conda] torchaudio 2.2.1+cu121 pypi_0 pypi
[conda] torchvision 0.17.1+cu121 pypi_0 pypi

[rusty1s]

Do you mind sharing a small code snippet to understand this better? If your input is a regular grid, there is no real point in using a GNN.