TensorFlow 1.x Deep Learning Cookbook

How to do it...

We proceed with the recipe as follows:

1. Start by importing a few modules

import sysimport numpy as npimport tensorflow as tffrom datetime import datetime

2. Get from command line the type of processing unit that you desire to use (either "gpu" or "cpu")

device_name = sys.argv[1] # Choose device from cmd line. Options: gpu or cpushape = (int(sys.argv[2]), int(sys.argv[2]))if device_name == "gpu":  device_name = "/gpu:0"else:  device_name = "/cpu:0"

3. Execute the matrix multiplication either on GPU or on CPU. The key instruction is with tf.device(device_name). It creates a new context manager, telling TensorFlow to perform those actions on either the GPU or the CPU

with tf.device(device_name): random_matrix = tf.random_uniform(shape=shape, ...