I have collected the TFLearn DCGAN example code and put it into my local Jupyter environment. Furthermore, I have changed some comments and added with tf.device('/gpu:0'): right before calling gan.fit(...), resulting in the following code:
# coding: utf-8 # In[1]: get_ipython().magic('matplotlib inline') from __future__ import division, print_function, absolute_import import matplotlib.pyplot as plt import numpy as np import tensorflow as tf import tflearn # In[2]: # Data loading and preprocessing import tflearn.datasets.mnist as mnist X, Y, testX, testY = mnist.load_data() X = np.reshape(X, newshape=[-1, 28, 28, 1]) # In[3]: # Noise data input z_dim = 200 total_samples = len(X) # In[4]: # Generator def generator(x, reuse=False): with tf.variable_scope('Generator', reuse=reuse): x = tflearn.fully_connected(x, n_units=7 * 7 * 128) x = tflearn.batch_normalization(x) x = tf.nn.tanh(x) x = tf.reshape(x, shape=[-1, 7, 7, 128]) x = tflearn.upsample_2d(x, 2) x = tflearn.conv_2d(x, 64, 5, activation='tanh') x = tflearn.upsample_2d(x, 2) x = tflearn.conv_2d(x, 1, 5, activation='sigmoid') return x # In[5]: # Discriminator def discriminator(x, reuse=False): with tf.variable_scope('Discriminator', reuse=reuse): x = tflearn.conv_2d(x, 64, 5, activation='tanh') x = tflearn.avg_pool_2d(x, 2) x = tflearn.conv_2d(x, 128, 5, activation='tanh') x = tflearn.avg_pool_2d(x, 2) x = tflearn.fully_connected(x, 1028, activation='tanh') x = tflearn.fully_connected(x, 2) x = tf.nn.softmax(x) return x # In[6]: # Input data gen_input = tflearn.input_data(shape=[None, z_dim], name='input_gen_noise') input_disc_noise = tflearn.input_data(shape=[None, z_dim], name='input_disc_noise') input_disc_real = tflearn.input_data(shape=[None, 28, 28, 1], name='input_disc_real') # In[7]: # Build discriminator disc_fake = discriminator(generator(input_disc_noise)) disc_real = discriminator(input_disc_real, reuse=True) disc_net = tf.concat([disc_fake, disc_real], axis=0) # In[8]: # Build stacked Generator/Discriminator gen_net = generator(gen_input, reuse=True) stacked_gan_net = discriminator(gen_net, reuse=True) # In[9]: # Build training ops for Discriminator # Each network optimization should only update its own variable, thus we need # to retrieve each network variable (with get_layer_variables_by_name) disc_vars = tflearn.get_layer_variables_by_name('Discriminator') # We need 2 target placeholders, for both the real and fake image target disc_target = tflearn.multi_target_data(['target_disc_fake', 'target_disc_real'], shape=[None, 2]) disc_model = tflearn.regression(disc_net, optimizer='adam', placeholder=disc_target, loss='categorical_crossentropy', trainable_vars=disc_vars, batch_size=64, name='target_disc', op_name='DISC') # In[10]: # Build training ops for Generator gen_vars = tflearn.get_layer_variables_by_name('Generator') gan_model = tflearn.regression(stacked_gan_net, optimizer='adam', loss='categorical_crossentropy', trainable_vars=gen_vars, batch_size=64, name='target_gen', op_name='GEN') # In[11]: # Define GAN model, that outputs the generated images gan = tflearn.DNN(gan_model, tensorboard_verbose=3) # In[12]: # Training # Prepare input data to feed to the discriminator disc_noise = np.random.uniform(-1., 1., size=[total_samples, z_dim]) # Prepare target data to feed to the discriminator (0: fake image, 1: real image) y_disc_fake = np.zeros(shape=[total_samples]) y_disc_real = np.ones(shape=[total_samples]) y_disc_fake = tflearn.data_utils.to_categorical(y_disc_fake, 2) y_disc_real = tflearn.data_utils.to_categorical(y_disc_real, 2) # In[13]: # Prepare input data to feed to the stacked Generator/Discriminator gen_noise = np.random.uniform(-1., 1., size=[total_samples, z_dim]) # Prepare target data to feed to the Discriminator # The Generator tries to fool the Discriminator, thus target is 1 (real images) y_gen = np.ones(shape=[total_samples]) y_gen = tflearn.data_utils.to_categorical(y_gen, 2) # In[14]: # Start training, feed both noise and real images with tf.device('/gpu:0'): gan.fit(X_inputs={'input_gen_noise': gen_noise, 'input_disc_noise': disc_noise, 'input_disc_real': X}, Y_targets={'target_gen': y_gen, 'target_disc_fake': y_disc_fake, 'target_disc_real': y_disc_real}, n_epoch=10) # In[15]: # Create another model from the Generator graph to generate some samples # for testing (re-using the same session to re-use the weights learnt) gen = tflearn.DNN(gen_net, session=gan.session) # In[16]: f, a = plt.subplots(4, 10, figsize=(10, 4)) for i in range(10): # Noise input z = np.random.uniform(-1., 1., size=[4, z_dim]) g = np.array(gen.predict({'input_gen_noise': z})) for j in range(4): # Generate image from noise. Extend to 3 channels for matplot figure. img = np.reshape(np.repeat(g[j][:, :, np.newaxis], 3, axis=2), newshape=(28, 28, 3)) a[j][i].imshow(img) f.show() plt.draw() I want to run this code on my NVIDIA GPU. I already have CUDA and cuDNN installed on my machine. Upon examining Windows Task Manager during training, I see that my CPU is stressed and my GPU lies dorment.
Can anyone give advice on how to properly implement with tf.device('/gpu:0'): as it's clear that the above code does not run on my NVIDIA GPU?
