How to reuse Theano function with different shared variables without rebuilding graph?

Question

I have a Theano function that is called several times, each time with different shared variables. The way it is implemented now, the Theano function gets redefined every time it is run. I assume, that this make the whole program slow, because every time the Theano functions gets defined the graph is rebuild.

def sumprod_shared(T_shared_array1, T_shared_array2):
    f = theano.function([], (T_shared_array1 * T_shared_array2).sum(axis=0))
    return f()

for factor in range(10):
    m1 = theano.shared(factor * array([[1, 2, 4], [5, 6, 7]]))
    m2 = theano.shared(factor * array([[1, 2, 4], [5, 6, 7]]))
    print sumprod_shared(m1, m2)

For non shared (normal) variables I can define the function once and then call it with different variables without redefining.

def sumprod_init():
    T_matrix1 = T.lmatrix('T_matrix1')
    T_matrix2 = T.lmatrix('T_matrix2')
    return theano.function([T_matrix1, T_matrix2], (T_matrix1 * T_matrix2).sum(axis=0))    

sumprod = sumprod_init()
for factor in range(10):
    np_array1 = factor * array([[1, 2, 4], [5, 6, 7]])
    np_array2 = factor * array([[1, 2, 4], [5, 6, 7]])
    print sumprod(np_array1, np_array2)

Is this possible also for shared variables?

Answer 1

You can use the givens keyword in theano.function for that. Basically, you do the following.

m1 = theano.shared(name='m1', value = np.zeros((3,2)) )
m2 = theano.shared(name='m2', value = np.zeros((3,2)) )

x1 = theano.tensor.dmatrix('x1')
x2 = theano.tensor.dmatrix('x2')

y = (x1*x2).sum(axis=0)
f = theano.function([],y,givens=[(x1,m1),(x2,m2)],on_unused_input='ignore')

then to loop through values you just set the value of the shared variables to the value you'd like. You have to set the on_unused_input to 'ignore' to use functions with no arguments in theano, by the way. Like this:

array1 = array([[1,2,3],[4,5,6]])
array2 = array([[2,4,6],[8,10,12]])

for i in range(10):
    m1.set_value(i*array1)
    m2.set_value(i*array2)
    print f()

It should work, at least that's how I've been working around it.

Answer 2

Currently it is not easily possible to reuse a Theano function with different shared variable.

But you have alternative:

1. Is it really a bottleneck? In the example, it is, but I suppose it is a simplified case. The only way to know is to profile it.
2. You compile 1 Theano function with the first shared variable. Then you can call the get_value/set_value on those shared variables before calling the Theano function. This way, you won't need to recompile the Theano function.