cudnn_tensorflow_speed_test.py

import tensorflow as tf
import numpy as np
from tensorflow.python import gen_batchnorm_training_op

with tf.device("/gpu:1"):

    def tf_bn(shape, data_format):
        """ output tensor for tensorflow implemented batch norm"""
        data = np.random.randn(*shape)*10 + 5
        inp = tf.Variable(data.astype("float32"), tf.float32)
        if data_format == "NHWC":
            scale = tf.Variable(np.ones([shape[-1],], dtype=np.float32), tf.float32)
            bias = tf.Variable(np.zeros([shape[-1],], dtype=np.float32), tf.float32)
            axies = [0, 1, 2]
        else:
            scale = tf.Variable(np.ones([1, shape[1], 1, 1], dtype=np.float32), tf.float32)
            bias = tf.Variable(np.zeros([1, shape[1], 1, 1], dtype=np.float32), tf.float32)
            axies = [0, 2, 3]

        eps = 0.001
        mean, varience = tf.nn.moments(inp, axies, keep_dims=True)
        out = tf.nn.batch_normalization(inp, mean, varience,
                                        bias,
                                        scale,
                                        eps)
        grads = tf.gradients(tf.reduce_sum(out), [inp, scale, bias])

        return grads

    def cudnn_bn(shape, data_format):
        """ make tensor for cudnn implementation of batch norm"""

        data = np.random.randn(*shape)*10 + 5
        inp = tf.Variable(data.astype("float32"), tf.float32)
        if data_format == "NHWC":
            raise NotImplemented()
        else:
            scale = tf.Variable(np.ones([1, shape[1], 1, 1], dtype=np.float32), tf.float32)
            bias = tf.Variable(np.zeros([1, shape[1], 1, 1], dtype=np.float32), tf.float32)
            axies = [0, 2, 3]

            running = np.zeros([1, shape[1], 1, 1])
            running_mean = tf.Variable(running.astype("float32"), tf.float32)
            running_inv_var = tf.Variable(running.astype("float32"), tf.float32)

        eps = 0.001

        out, save_mean, save_inv_var = gen_batchnorm_training_op.batch_norm_training(inp, scale, bias,
                running_mean, running_inv_var, 0.01)
        print out
        grads = tf.gradients(tf.reduce_sum(out), [inp, scale, bias])

        return grads

    #shape = [32, 8, 8, 32]
    #shape = [256, 16, 16, 32]
    #shape = [128, 64, 64, 32]
    shape = [128, 4, 4, 32]
    #shape = [128, 16, 16, 128]

    # tf_NHWC
    #speed = tf_bn(shape, "NHWC")

    shape = [shape[0], shape[3], shape[1], shape[2]]

    # tf_NCHWC
    #speed = tf_bn(shape, "NCHW")

    # cudnn_NCHWC
    speed = cudnn_bn(shape, "NCHW")


init = tf.initialize_all_variables()
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
sess.run([init])

# warmup
for i in range(100):
    grads = sess.run(speed)

import time
stime = time.time()
for i in range(100):
    grads = sess.run(speed)
print (time.time() - stime) * 1000