Recently, I've trained an AI for the game Snake using deep reinforcement learning. For this I've used the Keras API with tensorflow-gpu 2.
(Not sure if you need to look at the code, but if so, these are the most important parts:
- https://github.com/aeduin/snake_ai/blob/2e935fd124a54ef19f4287d19be9bd88af50f337/train.py
- https://github.com/aeduin/snake_ai/blob/2e935fd124a54ef19f4287d19be9bd88af50f337/ai.py#L125
- https://github.com/aeduin/snake_ai/blob/2e935fd124a54ef19f4287d19be9bd88af50f337/convnet_ai.py#L161
)
I measured that the total time spend in model.predict is 30 times more than in model.fit, while there are only 3 times as many predictions as actual training examples. My hypothesis would be that this is because the simulation of the game runs on the CPU, and latency of transferring data between CPU and GPU is what causes predicting to be so slow. (Fitting happens at the end of a simulated game, so all the training data is transferred in one go)
A secondary question would be: how do I test this hypothesis? But the main question is: how do I solve this problem? I know how to write CUDA programs. I could write a CUDA implementation of Snake. However, I don't know enough about Tensorflow or any other ML framework to know how to let it use this data without letting the data go through the CPU/RAM first. How would I do that?
One thing that helps is increasing the amount of simultaneously simulated games. The 30 times factor was with 512 or 256 (I don't remember which one) simultaneous games, and with lower counts, it gets much worse (50 times with only 2 simultaneous games). But I don't expect more than 512 simultaneous games adds a lot of value.
I've also tried running the convolutional neural network on the CPU. While the ratio is more reasonable (predicting is only 4 times slower than fitting), fitting and predicting both take more time than when using the GPU.