solopiu
solopiu

Reputation: 756

Force python script on GPU

Is there a way to force a Python script on GPU? In my code I use tensorflow and keras, and I have already the tensorflow-gpu version, but my code runs on CPU anyway. I'd like to know if there is a way to force the running on GPU independently on Tensorflow, Numpy or others.

Upvotes: 0

Views: 3002

Answers (2)

R. Klein
R. Klein

Reputation: 58

For TensorFlow (but not python in general) there is a good description of how to do this here: https://www.tensorflow.org/guide/gpu

To force a function to be performed on a specific processor (CPU or GPU) use the TensorFlow call to tf.device() as follows:

import tensorflow as tf
with tf.device('/GPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
  c = tf.matmul(a, b)

in this case the data for a and b will be stored on GPU0, and the operation "matmul" will be performed on the GPU0 as well.

Upvotes: 1

Ankur Datta
Ankur Datta

Reputation: 199

There maybe an issue with your tensorflow-gpu installation. Try creating a separate environment and reinstall only tensorflow-gpu.

Similar post

Upvotes: 0

Related Questions