AdaGrad + stochastic gradient descent using nested theano.scan()

nested_theano_scan_adagrad.py

import numpy as np
import theano
import theano.tensor as T
from theano.tensor.shared_randomstreams import RandomStreams

# AdaGrad + stochastic gradient descent using nested theano.scan()

train_x = np.random.rand(100)
train_y = train_x + np.random.rand(100) * 0.01

init_params = T.dvector()
init_r = T.dvector()
data_x = T.dvector()
data_y = T.dvector()
rnd = RandomStreams()

def fn1(_, init_params, init_r):
    indices = rnd.permutation(n=data_y.shape[0])

    def fn2(i, params, r, data_x, data_y, learning_rate):
        y = params[0] * data_x[i] + params[1]
        cost = (data_y[i] - y) ** 2
        g = T.grad(cost, params)
        r += g ** 2
        return params - learning_rate / T.sqrt(r) * g, r

    result2, updates2 = theano.scan(fn=fn2,
                                    sequences=indices,
                                    outputs_info=(init_params, init_r),
                                    non_sequences=(data_x, data_y, 3))
    return result2[0][-1], result2[1][-1]

result, updates = theano.scan(fn=fn1,
                              sequences=T.arange(100),
                              outputs_info=(init_params, init_r))

f = theano.function([init_params],
                    result[0][-1],
                    givens={
                        init_r: np.array([1e-8, 1e-8]),
                        data_x: train_x,
                        data_y: train_y
                    },
                    updates=updates)

print(f(np.array([0.5, 0.5])))  # [about 1, about 0]