lightaime / deep_gcns_torch

Pytorch Repo for DeepGCNs (ICCV'2019 Oral, TPAMI'2021), DeeperGCN (arXiv'2020) and GNN1000(ICML'2021): https://www.deepgcns.org
MIT License
1.13k stars 155 forks source link

GPU Memory is too high when using knn #61

Closed semchan closed 3 years ago

semchan commented 3 years ago

thanks for your great job. But I use the GCNs for inference, it is too high when using knn compared with the traditional convolution operation. Do you have some suggestion or potential methods to reduce the GPU usage? thanks a lot

def dense_knn_matrix(x, k=16):

lightaime commented 3 years ago

Hi @semchan. Thanks for your interest. It is possible to reduce the memory footprint by using an efficient implementation of knn such as https://github.com/facebookresearch/pytorch3d/blob/master/pytorch3d/ops/knn.py. I have not tried it yet. But it may be helpful.