cudnn_v3_algo_tests.py

import theano
from theano import tensor as T
import numpy as np
from theano.sandbox.cuda.dnn import dnn_conv
from theano.sandbox.cuda import host_from_gpu
from contexttimer import Timer
import theano.compile.mode
import theano.printing

def speed_test(_func, input_shape):
    x_val = np.random.randn(*input_shape).astype("float32")
    X.set_value(x_val)
    avg = []
    for i in range(40):
        with Timer() as t:
            _func()
        avg.append(t.elapsed)
    print np.mean(avg), np.std(avg)

def shared_floatx(*shape):
    return theano.shared(np.random.randn(*shape).astype("float32"))

X = shared_floatx(128, 64, 112, 112)
W = shared_floatx(32, 64, 7, 7)

y = dnn_conv(X, W, algo="small", border_mode=(0,0))[0,0,0,0]
_func = theano.function([], y)
speed_test(_func, X.get_value().shape)

y = dnn_conv(X, W, algo="large", border_mode=(0,0))[0,0,0,0]
_func = theano.function([], y)
speed_test(_func, X.get_value().shape)

y = dnn_conv(X, W, algo="fft", border_mode=(0,0))[0,0,0,0]
_func = theano.function([], y)
speed_test(_func, X.get_value().shape)