Improve speed of the StochPool Model

bayesmech / research-gnn-pooling

Research on Heirarchical pooling for Graph Neural Networks

MIT License

0 stars 0 forks source link

Improve speed of the StochPool Model #15

Closed AnimeshSinha1309 closed 1 year ago

AnimeshSinha1309 commented 2 years ago

[x] Use Pytorch profiler to see what is the bottleneck
[ ] Reimplement Stochpool in C++ or similar to speed up.

Jai2500 commented 1 year ago

Replaced for-loop with batched sparse matrices. Considerable speedup over previous times (approx. 30x). Closing for now without implementing the module in CUDA/C++.