TF Notes (2), Speedup and Benchmark with Two GPU Cards

---

這篇文章實作了官網的 同步式 data Parallelism 方法 ref，並且與原本只用一張GPU做個比較。實驗只使用兩張卡，多張卡方法一樣。主要架構如下圖 by TF 官網:

兩張卡等於是把一次要計算的 mini-batch 拆成兩半給兩個 (相同的) models 去並行計算 gradients，然後再交由 cpu 統一更新 model。詳細請自行參考官網。下面直接秀 Codes 和結果。

---

Machine Spec.

GPU 卡為 Tesla K40c

CPU 為 Intel(R) Xeon(R) CPU E5-2680 v3 @ 2.50GHz

---

單 GPU 跑 MNIST

直接上 Codes

import gzip
import os
import tempfile
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.contrib.layers import flatten
from tensorflow.examples.tutorials.mnist import input_data
from sklearn.utils import shuffle

"""
Data Loading
"""
dataPath='../dataset/MNIST_data/'
mnist = input_data.read_data_sets(dataPath, one_hot=True)
# read the images and reformat the image shape from [img_num,img_height,img_width] to [img_num,img_height,img_width,1]
img_width = 28
img_height = 28
images = mnist.train.images
img_num, _ = images.shape
images = np.reshape(images,(img_num,img_height,img_width))
images = images[...,np.newaxis]
print('(Input to CNN) Images with shape {}'.format(images.shape))
# read the labels
labels1Hot = mnist.train.labels
print('(Input to CNN) labels1Hot.shape = {}'.format(labels1Hot.shape))
labels = np.argmax(labels1Hot,axis=1)
labels = labels[...,np.newaxis]
print('labels.shape = {}'.format(labels.shape))
n_classes = len(np.unique(labels))
# load the validation set
images_valid = mnist.validation.images
img_num_valid = len(images_valid)
images_valid = np.reshape(images_valid,(img_num_valid,img_height,img_width))
images_valid = images_valid[...,np.newaxis]
labels1Hot_valid = mnist.validation.labels
print('Having %d number of validation images' % img_num_valid)
# plotting sample images
plt.figure(figsize=(15,5))
for i in np.arange(2*7):
    random_idx = np.random.randint(0,img_num)
    plt.subplot(2,7,i+1)
    plt.imshow(images[random_idx][...,0],cmap='gray')
    plt.title(labels[random_idx][0])

"""
First define the hyper-parameters
"""
# Hyper-parameters
EPOCHS = 30
BATCH_SIZE = 512
rate = 0.001
drop_out_keep_prob = 0.5
ksize = 5
cnn_depth_list = [16, 32]
mlp_depth_list = [256, 128]
cNum = 1

"""
Define the input output tensors
"""
# using one-hot decoding
x = tf.placeholder(tf.float32, (None, img_height, img_width, cNum))
one_hot_y = tf.placeholder(tf.int32, (None, n_classes))
#one_hot_y = tf.one_hot(y, n_classes)
keep_prob = tf.placeholder(tf.float32) # probability to keep units

"""
Define the graph and construct it
"""
class MNISTCNN:
    def __init__(self, ksize, cnn_depth_list, mlp_depth_list, img_height, img_width, cNum, n_classes):
        self._ksize = ksize
        self._cnn_depth_list = cnn_depth_list
        self._mlp_depth_list = mlp_depth_list
        self._img_height = img_height
        self._img_width = img_width
        self._cNum = cNum
        self._n_classes = n_classes
        
        self._mu = 0
        self._sigma = 0.1
        
    def create(self,x,keep_prob):
        conv = self._conv(x, self._cNum, self._ksize, self._cnn_depth_list[0], 'conv1')
        # Pooling. Input = 24x24xlayer_depth['layer_1']. Output = 12x12xlayer_depth['layer_1'].
        conv = tf.nn.max_pool(conv, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
        for lidx in range(1,len(self._cnn_depth_list)):
            conv = self._conv(conv, self._cnn_depth_list[lidx-1], self._ksize, self._cnn_depth_list[lidx], 'conv{}'.format(lidx+1))
            conv = tf.nn.max_pool(conv, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
        
        fc = flatten(conv)
        bsize, fc_in_dim = fc.shape
        fc = self._fc(fc,fc_in_dim,self._mlp_depth_list[0],'fc1')
        fc = tf.nn.dropout(fc, keep_prob)    # dropout
        for lidx in range(1,len(self._mlp_depth_list)):
            fc = self._fc(fc,self._mlp_depth_list[lidx-1],self._mlp_depth_list[lidx],'fc{}'.format(lidx+1))
            fc = tf.nn.dropout(fc, keep_prob)    # dropout
        
        with tf.variable_scope('logits') as scope:
            logits_w = tf.get_variable('logits_w', shape=[self._mlp_depth_list[-1],self._n_classes],\
                                       initializer=tf.random_uniform_initializer(-0.1,0.1))
            logits_b = tf.get_variable('logits_b', shape=[self._n_classes],\
                                       initializer=tf.zeros_initializer)
            logits = tf.nn.xw_plus_b(fc, logits_w, logits_b, name=scope.name)
        print(logits.shape)
        return logits
        
    def _conv(self, x, in_depth, ksize, out_depth, scope_name, relu=True):
        bsize,h,w,cNum = x.shape
        assert(h-(ksize-1)>=1)
        assert(w-(ksize-1)>=1)
        with tf.variable_scope(scope_name) as scope:
            # Create tf variables for the weights and biases
            weights = tf.get_variable('weights', shape=(ksize, ksize, in_depth, out_depth),\
                                      initializer=tf.random_normal_initializer(self._mu,self._sigma))
            biases = tf.get_variable('biases', shape=(out_depth),initializer=tf.zeros_initializer)
            out = tf.nn.conv2d(x, weights, strides=[1, 1, 1, 1], padding='VALID',name=scope.name) + biases
        if relu:
            # Apply ReLu non linearity
            relu = tf.nn.relu(out)
            return relu
        else:
            return out
    
    def _fc(self, x, num_in, num_out, scope_name, relu=True):
        """Create a fully connected layer."""
        with tf.variable_scope(scope_name) as scope:
            # Create tf variables for the weights and biases
            weights = tf.get_variable('weights', shape=[num_in, num_out],\
                                      initializer=tf.random_uniform_initializer(-0.1,0.1))
            biases = tf.get_variable('biases', [num_out],\
                                     initializer=tf.zeros_initializer)
            # Matrix multiply weights and inputs and add bias
            out = tf.nn.xw_plus_b(x, weights, biases, name=scope.name)
        if relu:
            # Apply ReLu non linearity
            relu = tf.nn.relu(out)
            return relu
        else:
            return out
            
mnistCNN = MNISTCNN(ksize, cnn_depth_list, mlp_depth_list, img_height, img_width, cNum, n_classes)
logits = mnistCNN.create(x,keep_prob)

"""
Define loss and optimizer
"""
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_y,logits=logits)
loss_operation = tf.reduce_mean(cross_entropy)
optimizer = tf.train.AdamOptimizer(learning_rate = rate)
training_operation = optimizer.minimize(loss_operation)

# Define accuracy evaluation
# calculate the average accuracy by calling evaluate(X_data, y_data)
correct_prediction = tf.equal(tf.argmax(logits, axis=1), tf.argmax(one_hot_y, axis=1))
accuracy_operation = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
def evaluate(X_data, y_data):
    num_examples = len(X_data)
    total_accuracy = 0
    sess = tf.get_default_session()
    for offset in range(0, num_examples, BATCH_SIZE):
        batch_x, batch_y = X_data[offset:offset+BATCH_SIZE], y_data[offset:offset+BATCH_SIZE]
        accuracy = sess.run(accuracy_operation, feed_dict={x: batch_x, one_hot_y: batch_y, keep_prob: 1.0})
        total_accuracy += (accuracy * len(batch_x))
    return total_accuracy / num_examples
    
"""
Run Session
"""
### Train your model here.
import time
if not os.path.isdir('./models'):
    os.makedirs('./models')
#saver = tf.train.Saver()
accumulate_time = 0.0
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    num_examples = img_num
    
    print("Training...")
    print()
    train_accuracy = np.zeros(EPOCHS)
    validation_accuracy = np.zeros(EPOCHS)
    for i in range(EPOCHS):
        stime = time.time()
        acc_train_accuracy = 0
        X_train, y_train = shuffle(images, labels1Hot)
        for offset in range(0, num_examples, BATCH_SIZE):
            end = offset + BATCH_SIZE
            batch_x, batch_y = X_train[offset:end], y_train[offset:end]
            sess.run(training_operation, feed_dict={x: batch_x, one_hot_y: batch_y, keep_prob: drop_out_keep_prob})

        etime = time.time()
        accumulate_time += etime - stime
        validation_accuracy[i] = evaluate(images_valid, labels1Hot_valid)
        print("EPOCH {} ...".format(i+1))
        print("Validation Accuracy = {:.3f}".format(validation_accuracy[i]))
        print()

    print('Cost time: ' + str(accumulate_time) + ' sec.')

---

同步式 data Parallelism 在兩張 GPU 跑 MNIST

一樣直接上 Codes

import gzip
import os
import tempfile
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.contrib.layers import flatten
from tensorflow.examples.tutorials.mnist import input_data
from sklearn.utils import shuffle

"""
Data Loading
"""
dataPath='../dataset/MNIST_data/'
mnist = input_data.read_data_sets(dataPath, one_hot=True)
# read the images and reformat the image shape from [img_num,img_height,img_width] to [img_num,img_height,img_width,1]
img_width = 28
img_height = 28
images = mnist.train.images
img_num, _ = images.shape
images = np.reshape(images,(img_num,img_height,img_width))
images = images[...,np.newaxis]
print('(Input to CNN) Images with shape {}'.format(images.shape))
# read the labels
labels1Hot = mnist.train.labels
print('(Input to CNN) labels1Hot.shape = {}'.format(labels1Hot.shape))
labels = np.argmax(labels1Hot,axis=1)
labels = labels[...,np.newaxis]
print('labels.shape = {}'.format(labels.shape))
n_classes = len(np.unique(labels))# read the labels
labels1Hot = mnist.train.labels
print('(Input to CNN) labels1Hot.shape = {}'.format(labels1Hot.shape))
labels = np.argmax(labels1Hot,axis=1)
labels = labels[...,np.newaxis]
print('labels.shape = {}'.format(labels.shape))
n_classes = len(np.unique(labels))
# load the validation set
images_valid = mnist.validation.images
img_num_valid = len(images_valid)
images_valid = np.reshape(images_valid,(img_num_valid,img_height,img_width))
images_valid = images_valid[...,np.newaxis]
labels1Hot_valid = mnist.validation.labels
print('Having %d number of validation images' % img_num_valid)
plt.figure(figsize=(15,5))
for i in np.arange(2*7):
    random_idx = np.random.randint(0,img_num)
    plt.subplot(2,7,i+1)
    plt.imshow(images[random_idx][...,0],cmap='gray')
    plt.title(labels[random_idx][0])

"""
First define the hyper-parameters
"""
# Hyper-parameters
EPOCHS = 30
BATCH_SIZE = 512
rate = 0.001
drop_out_keep_prob = 0.5
ksize = 5
cnn_depth_list = [16, 32]
mlp_depth_list = [256, 128]
cNum = 1

"""
Define the graph and construct it
"""
class MNISTCNN:
    def __init__(self, ksize, cnn_depth_list, mlp_depth_list, img_height, img_width, cNum, n_classes):
        self._ksize = ksize
        self._cnn_depth_list = cnn_depth_list
        self._mlp_depth_list = mlp_depth_list
        self._img_height = img_height
        self._img_width = img_width
        self._cNum = cNum
        self._n_classes = n_classes
        
        self._mu = 0
        self._sigma = 0.1
        
    def create(self,x,keep_prob):
        conv = self._conv(x, self._cNum, self._ksize, self._cnn_depth_list[0], 'conv1')
        # Pooling. Input = 24x24xlayer_depth['layer_1']. Output = 12x12xlayer_depth['layer_1'].
        conv = tf.nn.max_pool(conv, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
        for lidx in range(1,len(self._cnn_depth_list)):
            conv = self._conv(conv, self._cnn_depth_list[lidx-1], self._ksize, self._cnn_depth_list[lidx], 'conv{}'.format(lidx+1))
            conv = tf.nn.max_pool(conv, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
        
        fc = flatten(conv)
        bsize, fc_in_dim = fc.shape
        fc = self._fc(fc,fc_in_dim,self._mlp_depth_list[0],'fc1')
        fc = tf.nn.dropout(fc, keep_prob)    # dropout
        for lidx in range(1,len(self._mlp_depth_list)):
            fc = self._fc(fc,self._mlp_depth_list[lidx-1],self._mlp_depth_list[lidx],'fc{}'.format(lidx+1))
            fc = tf.nn.dropout(fc, keep_prob)    # dropout
        
        with tf.variable_scope('logits') as scope:
            with tf.device('/cpu:0'):
                logits_w = tf.get_variable('logits_w', shape=[self._mlp_depth_list[-1],self._n_classes],\
                                           initializer=tf.random_uniform_initializer(-0.1,0.1))
                logits_b = tf.get_variable('logits_b', shape=[self._n_classes],\
                                           initializer=tf.zeros_initializer)
            logits = tf.nn.xw_plus_b(fc, logits_w, logits_b, name=scope.name)
        print(logits.shape)
        return logits
        
    def _conv(self, x, in_depth, ksize, out_depth, scope_name, relu=True):
        bsize,h,w,cNum = x.shape
        assert(h-(ksize-1)>=1)
        assert(w-(ksize-1)>=1)
        with tf.variable_scope(scope_name) as scope:
            with tf.device('/cpu:0'):
                # Create tf variables for the weights and biases
                weights = tf.get_variable('weights', shape=(ksize, ksize, in_depth, out_depth),\
                                          initializer=tf.random_normal_initializer(self._mu,self._sigma))
                biases = tf.get_variable('biases', shape=(out_depth),initializer=tf.zeros_initializer)
            out = tf.nn.conv2d(x, weights, strides=[1, 1, 1, 1], padding='VALID',name=scope.name) + biases
        if relu:
            # Apply ReLu non linearity
            relu = tf.nn.relu(out)
            return relu
        else:
            return out
    
    def _fc(self, x, num_in, num_out, scope_name, relu=True):
        """Create a fully connected layer."""
        with tf.variable_scope(scope_name) as scope:
            with tf.device('/cpu:0'):
                # Create tf variables for the weights and biases
                weights = tf.get_variable('weights', shape=[num_in, num_out],\
                                          initializer=tf.random_uniform_initializer(-0.1,0.1))
                biases = tf.get_variable('biases', [num_out],\
                                         initializer=tf.zeros_initializer)
            # Matrix multiply weights and inputs and add bias
            out = tf.nn.xw_plus_b(x, weights, biases, name=scope.name)
        if relu:
            # Apply ReLu non linearity
            relu = tf.nn.relu(out)
            return relu
        else:
            return out
            
mnistCNN = MNISTCNN(ksize, cnn_depth_list, mlp_depth_list, img_height, img_width, cNum, n_classes)

# Averaging gradients for all tower models on GPU
def average_gradients(tower_grads):
    """Calculate the average gradient for each shared variable across all towers.
    Note that this function provides a synchronization point across all towers.
    Args:
    tower_grads: List of lists of (gradient, variable) tuples. The outer list
      is over individual gradients. The inner list is over the gradient
      calculation for each tower.
    Returns:
     List of pairs of (gradient, variable) where the gradient has been averaged
     across all towers.
    """
    average_grads = []
    for grad_and_vars in zip(*tower_grads):
        # Note that each grad_and_vars looks like the following:
        #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
        grads = []
        for g, _ in grad_and_vars:
            # Add 0 dimension to the gradients to represent the tower.
            expanded_g = tf.expand_dims(g, 0)

            # Append on a 'tower' dimension which we will average over below.
            grads.append(expanded_g)

        # Average over the 'tower' dimension.
        grad = tf.concat(axis=0, values=grads)
        grad = tf.reduce_mean(grad, 0)

        # Keep in mind that the Variables are redundant because they are shared
        # across towers. So .. we will just return the first tower's pointer to
        # the Variable.
        v = grad_and_vars[0][1]
        grad_and_var = (grad, v)
        average_grads.append(grad_and_var)
    return average_grads

# Construct model for each GPU, where variables are shared/updated by CPU
with tf.device('/cpu:0'):
    optimizer = tf.train.AdamOptimizer(learning_rate = rate)

# Calculate the gradients for each model tower.
tower_grads = []
logits_list = []
feed_x = []
feed_one_hot_y = []
keep_prob = tf.placeholder(tf.float32) # probability to keep units
with tf.variable_scope(tf.get_variable_scope()):
    for i in range(2):
        with tf.device('/gpu:%d' % i):
            x = tf.placeholder(tf.float32, (None, img_height, img_width, cNum))
            feed_x.append(x)
            one_hot_y = tf.placeholder(tf.int32, (None, n_classes))
            feed_one_hot_y.append(one_hot_y)

            logits = mnistCNN.create(x,keep_prob)
            logits_list.append(logits)

            cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_y,logits=logits)
            loss_op = tf.reduce_mean(cross_entropy)

            tf.get_variable_scope().reuse_variables()

            # Calculate the gradients for each batch of data on this model tower.
            grads = optimizer.compute_gradients(loss_op)
            # Keep track of the gradients across all towers.
            tower_grads.append(grads)

with tf.device('/cpu:0'):
    # We must calculate the mean of each gradient. Note that this is the
    # synchronization point across all towers.
    grads = average_gradients(tower_grads)
    
    # Apply the gradients to adjust the shared variables.
    apply_gradient_op = optimizer.apply_gradients(grads)
    
    training_op = apply_gradient_op
    
"""
Prediction/Inference Part
"""
# Define accuracy evaluation, calculate the average accuracy by calling evaluate(X_data, y_data)
# Using model that in the First GPU to calculate
correct_prediction = tf.equal(tf.argmax(logits_list[0], axis=1), tf.argmax(feed_one_hot_y[0], axis=1))
accuracy_operation = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

def evaluate(X_data, y_data):
    num_examples = len(X_data)
    total_accuracy = 0
    sess = tf.get_default_session()
    for offset in range(0, num_examples, BATCH_SIZE):
        batch_x, batch_y = X_data[offset:offset+BATCH_SIZE], y_data[offset:offset+BATCH_SIZE]
        accuracy = sess.run(accuracy_operation, feed_dict={feed_x[0]: batch_x, feed_one_hot_y[0]: batch_y, keep_prob: 1.0})
        total_accuracy += (accuracy * len(batch_x))
    return total_accuracy / num_examples
    
"""
Run Session
"""
# {feed_x[0]:batch_x_1, feed_x[1]:batch_x_2,\
# feed_one_hot_y[0]:batch_y_1, feed_one_hot_y[1]:batch_y_1, keep_prob:drop_out_keep_prob}
def gen_feed_dict(batch_x, batch_y, drop_out_keep_prob):
    assert(len(batch_x)==len(batch_y))
    assert(len(batch_x)%2==0)
    data_num = int(len(batch_x)/2)
    rtn_dict = {}
    rtn_dict[feed_x[0]] = batch_x[:data_num]
    rtn_dict[feed_x[1]] = batch_x[data_num:]
    rtn_dict[feed_one_hot_y[0]] = batch_y[:data_num]
    rtn_dict[feed_one_hot_y[1]] = batch_y[data_num:]
    rtn_dict[keep_prob] = drop_out_keep_prob
    return rtn_dict
    
### Train your model here.
import time
if not os.path.isdir('./models'):
    os.makedirs('./models')
#saver = tf.train.Saver()
accumulate_time = 0.0
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    num_examples = img_num
    
    print("Training...")
    print()
    train_accuracy = np.zeros(EPOCHS)
    validation_accuracy = np.zeros(EPOCHS)
    stime = time.time()
    for i in range(EPOCHS):
        stime = time.time()
        acc_train_accuracy = 0
        X_train, y_train = shuffle(images, labels1Hot)
        for offset in range(0, num_examples, BATCH_SIZE):
            end = offset + BATCH_SIZE
            batch_x, batch_y = X_train[offset:end], y_train[offset:end]
            feed_dict = gen_feed_dict(batch_x, batch_y, drop_out_keep_prob)
            sess.run(training_op, feed_dict=feed_dict)

        etime = time.time()
        accumulate_time += etime - stime
        validation_accuracy[i] = evaluate(images_valid, labels1Hot_valid)
        print("EPOCH {} ...".format(i+1))
        print("Validation Accuracy = {:.3f}".format(validation_accuracy[i]))
    
    print('Cost time: ' + str(accumulate_time) + ' sec.')

幾個注意處:

1. 記得建立 variables (tf.get_variable) 時要使用 with tf.device('/cpu:0'): 確保變量是存在 cpu 內

2. 可以跑一小段 code:

   with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
       sess.run(tf.global_variables_initializer())

   來觀察變量是否正確放在 cpu 上。

3. 延續 2. 若使用 jupyter notebook 可以這樣做 jupyter notebook > outputlog，執行完 2. 的 code 接著 cat outputlog | grep 'cpu' 觀察變量是否存在。
4. 使用 collections (如下) 來確認變數有正確分享 (養成好習慣)

   trainable_collection = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
   global_collection = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
   print('Without Scope: len(trainable_collection)={}; len(global_collection)={}'.format(len(trainable_collection),len(global_collection)))

---

Benchmark Results

batch size = 128 時，使用兩張 GPU 花的時間為一張的 0.78 倍。而 batch size = 512 時的效果更明顯，為 0.69 倍。

---

一點小結論

這種同步的架構適合在 batch size 大的時候，效果會更明顯。實驗起來兩張卡在 512 batch size 花的時間在一張卡的 0.7 倍。
不過相比使用兩張卡，一張卡其實有一點優勢是在變量全部放在 GPU 上，因此省去了 CPU <–> GPU 的傳輸代價。這也是主要只到 0.7 倍，而沒有接近 0.5 倍的關鍵原因。

---

Reference

1. Tensorflow 官網 同步式 data Parallelism 方法
2. Tensorflow github cifar10_multi_gpu_train.py