search

QUVA-Lab / e2cnn

E(2)-Equivariant CNNs Library for Pytorch
https://quva-lab.github.io/e2cnn/
Other
596 stars 75 forks source link

A question about cuda acceleration #43

Closed kuazhangxiaoai closed 3 years ago

kuazhangxiaoai commented 3 years ago

Is this package accelerated by cuda?

ejnnr commented 3 years ago

The nn subpackage uses Pytorch internally for everything, so it supports CUDA acceleration. You can move GeometricTensors to GPU with GeometricTensor.to. EquivariantModules are just Pytorch Modules so they can also be moved to GPU with .to(device).

The other subpackages use numpy and are therefore not CUDA accelerated (also see https://github.com/QUVA-Lab/e2cnn/issues/39#issuecomment-832546489). But that part of the code is only run once when you initialize your network. So all of the training and inference can be done on GPU.

Gabri95 commented 3 years ago

Hi @kgavrilyuk

Thanks @ejnnr for the faster reply :) Indeed, CUDA is supported in the nn subpackage, so once you have built your neural network, you should be able to use CUDA acceleration as you usually do in PyTorch

Best, Gabriele