GPU Memory is too high when using knn

[semchan]

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]

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.