mlech26l / ncps

PyTorch and TensorFlow implementation of NCP, LTC, and CfC wired neural models
https://www.nature.com/articles/s42256-020-00237-3
Apache License 2.0
1.86k stars 298 forks source link

GPU support for NCP model training? #8

Closed bhorsfield closed 3 years ago

bhorsfield commented 3 years ago

Hi there,

Is it possible to train an NCP model in a TensorFlow Docker container with GPU support?

I have been getting good results by training my NCP model with the CPU-only version of TensorFlow, but I would really like to train in the GPU in order to reduce training time.

Is it possible to do this? I have tried it once using a TensorFlow Docker container with GPU support, but I got the following error message: ModuleNotFoundError: No module named 'kerasncp'.

Thanks, Brendan.

mlech26l commented 3 years ago

Hi Brendan,

GPU support is provided implicitly via Tensorflow. However, the achieved speedup depends on various parameters (size of the networks, other layers, batch size, ...).

According to your error message, it seems that the keras-ncp module is not installed in your docker container. Before running your script, you first have to run

pip3 install keras-ncp

inside the docker container.

Best, Mathias

bhorsfield commented 3 years ago

Hi Mathias,

Thanks, that fixed it. I can now train my NCP model with GPU support. The training time has decreased by a factor of nearly 3 compared to CPU-only training.

Best Regards, Brendan.