In [1]:
import numpy as np
import algopy
def func(x0, x1, x2):
"Toy function (calibrated to take roughly the same time as real one)"
return 8.3145 * x0 * (x1 * algopy.log(x1) + x2 * algopy.log(x2)) + \
2000*x1*x2 + 100*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + 2000*x1*x2 + 100*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + \
2400*x1*x2 + 120*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + 2000*x1*x2 + 100*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + \
2000*x1*x2 + 160*x1*x2*(x1-x2) + 1000*x1*x2*(x1-x2)**2 + 2000*x1*x2 + 100*x1*x2*(x1-x2) + 500*x1*x2*(x1-x2)**2
In [2]:
# We repeat the same two values over and over to simulate lots of different input
C = 8000
inp_arr = np.tile([[300,0.5,0.5], [600, 0.4, 0.6]], (int(C/2),1))
print(inp_arr.shape)
In [3]:
def extract_hessian(N, y):
"Vectorized version of extract_hessian"
H = np.zeros((y.data.shape[1], N,N), dtype=y.data.dtype)
for n in range(N):
for m in range(n):
a = sum(range(n+1))
b = sum(range(m+1))
k = sum(range(n+2)) - m - 1
#print 'k,a,b=', k,a,b
if n!=m:
tmp = (y.data[2, :, k] - y.data[2, :, a] - y.data[2, :, b])
H[:, m,n]= H[:, n,m]= tmp
a = sum(range(n+1))
H[:, n,n] = 2*y.data[2, :, a]
return H
In [4]:
# generate directions
N = func.__code__.co_argcount
M = (N*(N+1))/2
L = (N*(N-1))/2
S = np.zeros((N,M))
s = 0
i = 0
for n in range(1,N+1):
S[-n:, s:s+n] = np.eye(n)
S[-n, s:s+n] = np.ones(n)
s += n
i += 1
#print(S)
S = S[::-1].T
#print(S)
x = algopy.UTPM(np.zeros((3, inp_arr.shape[0]) + S.shape))
x.data[0, :, :, :] = inp_arr[..., None, :]
x.data[1, :, :] = S
In [5]:
#%%timeit
y = func(*[x[..., i] for i in range(N)])
obj_algopy = y.data[0, :, 0]
grad_algopy = y.data[1, :, np.cumsum(np.arange(N), dtype=np.int)].T
hess_algopy = extract_hessian(N, y)
#print('OBJ SHAPE', obj_algopy.shape)
#print('GRAD SHAPE', grad_algopy.shape)
#print('HESS SHAPE', hess_algopy.shape)
In [6]:
import autograd.numpy as anp
from autograd import elementwise_grad, jacobian
from itertools import chain
def autograd_func(x0, x1, x2):
"Toy function (calibrated to take roughly the same time as real one)"
return 8.3145 * x0 * (x1 * anp.log(x1) + x2 * anp.log(x2)) + \
2000*x1*x2 + 100*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + 2000*x1*x2 + 100*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + \
2400*x1*x2 + 120*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + 2000*x1*x2 + 100*x1*x2*(x1-x2) + 200*x1*x2*(x1-x2)**2 + \
2000*x1*x2 + 160*x1*x2*(x1-x2) + 1000*x1*x2*(x1-x2)**2 + 2000*x1*x2 + 100*x1*x2*(x1-x2) + 500*x1*x2*(x1-x2)**2
def elementwise_hess(fun, argnum=0):
"""
From https://github.com/HIPS/autograd/issues/60
"""
def sum_latter_dims(x):
return anp.sum(x.reshape(x.shape[0], -1), 1)
def sum_grad_output(*args, **kwargs):
return sum_latter_dims(elementwise_grad(fun)(*args, **kwargs))
return jacobian(sum_grad_output, argnum)
def build_functions():
obj = autograd_func
def argwrapper(args):
return obj(*args)
def grad_func(*args):
inp = anp.array(anp.broadcast_arrays(*args))
result = anp.atleast_2d(elementwise_grad(argwrapper)(inp))
# Put 'gradient' axis at end
axes = list(range(len(result.shape)))
result = result.transpose(*chain(axes[1:], [axes[0]]))
return result
def hess_func(*args):
result = anp.atleast_3d(elementwise_hess(argwrapper)(anp.array(anp.broadcast_arrays(*args))))
# Put 'hessian' axes at end
axes = list(range(len(result.shape)))
result = result.transpose(*chain(axes[2:], axes[0:2]))
return result
return obj, grad_func, hess_func
autograd_obj, autograd_grad, autograd_hess = build_functions()
In [7]:
#%%timeit
ag_o = autograd_obj(*[inp_arr[..., i] for i in range(inp_arr.shape[-1])])
ag_g = autograd_grad(*[inp_arr[..., i] for i in range(inp_arr.shape[-1])])
ag_h = autograd_hess(*[inp_arr[..., i] for i in range(inp_arr.shape[-1])])
In [8]:
# Have to comment out '%%timeit' lines to actually run this
import numpy.testing
numpy.testing.assert_allclose(obj_algopy, ag_o)
numpy.testing.assert_allclose(grad_algopy, ag_g)
numpy.testing.assert_allclose(hess_algopy, ag_h)
print('equivalent')
In [ ]: