TensorFlow supports both CPUs and GPUs. It also supports distributed computation. We can use TensorFlow on multiple devices in one or more computer system. TensorFlow names the supported devices as "/device:CPU:0" (or "/cpu:0") for the CPU devices and "/device:GPU:I" (or "/gpu:I") for the i^th GPU device.

As mentioned earlier, GPUs are much faster than CPUs because they have many small cores. However, it is not always an advantage in terms of computational speed to use GPUs for all types of computations. The overhead associated with GPUs can sometimes be more computationally expensive than the advantage of parallel computation offered by GPUs. To deal with this issue, TensorFlow has provisions to place computations on a particular device. By default, if both CPU and GPU are present, TensorFlow gives priority to GPU.

TensorFlow represents devices as strings. Here, we will show you how one can manually assign a device for matrix multiplication in TensorFlow. To verify that TensorFlow is indeed using the device (CPU or GPU) specified, we create the session with the log_device_placement flag set to True, namely, config=tf.ConfigProto(log_device_placement=True):

1. If you are not sure about the device and want TensorFlow to choose the existing and supported device, you can set the allow_soft_placement flag to True:

config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True) 

2. Manually select CPU for operation:

with tf.device('/cpu:0'): 
    rand_t = tf.random_uniform([50,50], 0, 10, dtype=tf.float32, seed=0) 
    a = tf.Variable(rand_t) 
    b = tf.Variable(rand_t) 
    c = tf.matmul(a,b) 
    init = tf.global_variables_initializer() 
 
sess = tf.Session(config) 
sess.run(init) 
print(sess.run(c))    

3. We get the following output:

We can see that all the devices, in this case, are '/cpu:0'.

4. Manually select a single GPU for operation:

with tf.device('/gpu:0'): 
    rand_t = tf.random_uniform([50,50], 0, 10, dtype=tf.float32, seed=0) 
    a = tf.Variable(rand_t) 
    b = tf.Variable(rand_t) 
    c = tf.matmul(a,b) 
    init = tf.global_variables_initializer() 
 
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True)) 
sess.run(init) 
print(sess.run(c)) 

5. The output now changes to the following:

6. The '/cpu:0' after each operation is now replaced by '/gpu:0'.
7. Manually select multiple GPUs:

c=[]
for d in ['/gpu:1','/gpu:2']:
     with tf.device(d):
         rand_t = tf.random_uniform([50, 50], 0, 10, dtype=tf.float32, seed=0)
         a = tf.Variable(rand_t)
         b = tf.Variable(rand_t)
         c.append(tf.matmul(a,b))
         init = tf.global_variables_initializer()

sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True,log_device_placement=True))
sess.run(init)
print(sess.run(c))
sess.close() 

8. In this case, if the system has three GPU devices, then the first set of multiplication will be carried on by '/gpu:1' and the second set by '/gpu:2'.

The tf.device() argument selects the device (CPU or GPU). The with block ensures the operations for which the device is selected. All the variables, constants, and operations defined within the with block will use the device selected in tf.device(). Session configuration is controlled using tf.ConfigProto. By setting the allow_soft_placement and log_device_placement flags, we tell TensorFlow to automatically choose the available devices in case the specified device is not available, and to give log messages as output describing the allocation of devices while the session is executed.