Double Lop

double_lop.py

import time
import os
import numpy as np
import theano
import theano.tensor as T
from IPython.display import Image

def Rop(f, x, v):
    if isinstance(f, (list, tuple)):
        u = [T.zeros_like(fi) for fi in f]
    else:
        u = T.zeros_like(f)
    return T.Lop( T.Lop(f, x, u), u, v)

def print_graphs(f, w, v, shapes=None, folder=None):
    if isinstance(w, (list, tuple)):
        inputs = w + v
    else:
        inputs = [w, v]
    jvp = T.Rop(f, w, v)
    a_jvp = Rop(f, w, v)
    compiled_jvp = theano.function(inputs, jvp)
    compiled_a_jvp = theano.function(inputs, a_jvp)
    if folder is not None:
        theano.printing.pydotprint(compiled_jvp, os.path.join(folder, "rop.png"),
                                             var_with_name_simple=True)
        theano.printing.pydotprint(compiled_jvp, os.path.join(folder, "double_lop.png"),
                                             var_with_name_simple=True)
    display(Image(theano.printing.pydotprint(compiled_jvp,
                                             return_image=True, 
                                             var_with_name_simple=True)))
    display(Image(theano.printing.pydotprint(compiled_a_jvp,
                                             return_image=True, 
                                             var_with_name_simple=True)))
    if shapes is not None:
        np_inputs = [np.random.randn(*s) for s in shapes] + [np.random.randn(*s) for s in shapes]
        np_inputs = [array.astype(theano.config.floatX) for array in np_inputs]
        for i in range(20):
            f1 = compiled_a_jvp(*np_inputs)
        t = time.time()
        for i in range(100):
            f1 = compiled_a_jvp(*np_inputs)
        t1 = time.time() - t
        for i in range(20):
            f2 = compiled_jvp(*np_inputs)
        t = time.time()
        for i in range(100):
            f2 = compiled_jvp(*np_inputs)
        t2 = time.time() - t
        if isinstance(f1, (list, tuple)):
            abs_diff = sum(np.sum(np.abs(f1_i - f2_i)) for f1_i, f2_i in zip(f1, f2))
        else:
            abs_diff = np.sum(np.abs(f1 - f2))
        print("Absolute difference:", abs_diff)
        print("Time regular Rop: {:.3f}s, time double Lop: {:.3f}s".format(t2, t1))
        
x = T.fmatrix('x')
y = T.fmatrix('y')
w = [x, y]
u = [T.fmatrix('vx'), T.fmatrix('vy')]
f = T.sum(T.sin(T.sin(T.sin(x))) * T.cos(T.cos(T.cos(y))))
print_graphs(f, w, u, shapes=[(15, 15), (15, 15)])

from lasagne.layers import Conv2DLayer as conv
from lasagne.layers import DenseLayer as dense
from lasagne.layers import MaxPool2DLayer as maxpool
from lasagne.layers import InputLayer, get_output


def lenet5(inputs_shape, folder=None):
    w = T.ftensor4("input")
    l = InputLayer(inputs_shape, input_var=w)
    l = conv(l, 32, 5)
    l = maxpool(l, 2)

    l = conv(l, 64, 5)
    l = maxpool(l, 2)
    l = dense(l, 512)
    # output layers
    logits = dense(l, 10, nonlinearity=None)
    
    f = get_output(logits)
    v = T.ftensor4("vi")
    print_graphs(f, w, v, [inputs_shape], folder=folder)

lenet5([100, 1, 28, 28], "/home/alex/work/python")