Hands-On Deep Learning with Go

GPU acceleration

Convolution and its associated operations tend to do very well on GPU acceleration. You saw earlier that our GPU acceleration had minimal impact, but it is extremely useful for building CNNs. All we need to do is add the magical 'cuda' build tag, as shown here:

go build -tags='cuda'

As we tend to be more memory constrained on GPUs, be aware that the same batch size may not work on your GPU. The model as mentioned previously uses around 4 GB of memory, so you will probably want to reduce the batch size if you have less than 6 GB of GPU memory (because presumably, you will be using about 1 GB for your normal desktop). If your model is running very slowly, or the CUDA version of your executable just fails, it would be prudent ...