CODEWITHSUNDEEP
tensorflow
Nitish Gupta
Nitish Gupta

Reputation: 11

GPU + CPU Tensorflow Training

Setup

I have a network, one whose parameter is a large-embedding matrix (3Million X 300 sized), say embed_mat.

During training, for each mini-batch, I only update a small subset of the vectors from embed_mat (max 15000 vectors) which are chosen using the embedding_lookup op. I am using the Adam optimizer to train my model.

As I cannot store this embed_mat in the GPU, due to its size, I define it under CPU (say /cpu:0) device, but the rest of the parameters of the model, the optimizer etc. are defined under a GPU (say, gpu:/0) device.

Questions

  1. I see that my GPU usage is very minimal (200 MB), which suggests all my training is happening on the CPU. What I expected was that the result of the embedding_lookup is copied to the GPU and all my training happens there. Am I doing something wrong.

  2. The training time is very largely affected by the size (num_vectors) of the embedding matrix which doesn't seem correct to me. In any mini-batch, I only update my network parameters and the vectors I looked up (~15000), so the training time should, if at all, grow sub-linearly with the size of the embedding matrix.

  3. Is there a way to automatically and seamlessly split up my embed_mat to multiple GPUs for faster training?

  4. I suspect the Adam Optimizer for this. Looks like because the embed_mat is on the CPU, all training is happening on the CPU. Is this correct?

Upvotes: 1

Views: 1171

Answers (1)

Alexandre Passos
Alexandre Passos

Reputation: 5206

Try visualizing on tensorboard where each of your ops is placed. In the "graph" tab you can color by "device". Ideally the embedding variable, the embedding lookup, and the embedding gradient update should be in the CPU, while most other things should be in the GPU.

Upvotes: 1

Related Questions