#!/usr/bin/env python2
import theano
import theano.tensor as tt
import lasagne
import numpy as np
from lasagne.layers import InputLayer, DenseLayer
from kusanagi import utils
from kusanagi.ghost.regression import layers
from kusanagi.ghost.optimizers import SGDOptimizer
from kusanagi.ghost.regression import nonlinearities
from kusanagi.ghost.regression import BaseRegressor
from kusanagi.ghost.regression import objectives
floatX = theano.config.floatX
def mlp(input_dims, output_dims, hidden_dims=[200]*4, batchsize=None,
nonlinearities=nonlinearities.rectify,
output_nonlinearity=nonlinearities.linear,
W_init=lasagne.init.GlorotUniform('relu'),
b_init=lasagne.init.Uniform(0.01),
name='mlp', **kwargs):
if not isinstance(nonlinearities, list):
nonlinearities = [nonlinearities]*len(hidden_dims)
if not isinstance(W_init, list):
W_init = [W_init]*(len(hidden_dims)+1)
if not isinstance(b_init, list):
b_init = [b_init]*(len(hidden_dims)+1)
network_spec = []
# input layer
input_shape = (batchsize, input_dims)
network_spec.append((InputLayer,
dict(shape=input_shape,
name=name+'_input')))
# hidden layers
for i in range(len(hidden_dims)):
layer_type = DenseLayer
network_spec.append((layer_type,
dict(num_units=hidden_dims[i],
nonlinearity=nonlinearities[i],
W=W_init[i],
b=b_init[i],
name=name+'_fc%d' % (i))))
# output layer
layer_type = DenseLayer
network_spec.append((layer_type,
dict(num_units=output_dims,
nonlinearity=output_nonlinearity,
W=W_init[-1],
b=b_init[-1],
name=name+'_output')))
return network_spec
def dropout_mlp(input_dims, output_dims, hidden_dims=[200]*4, batchsize=None,
nonlinearities=nonlinearities.rectify,
output_nonlinearity=nonlinearities.linear,
W_init=lasagne.init.GlorotUniform('relu'),
b_init=lasagne.init.Uniform(0.01),
p=0.5, p_input=0.2,
dropout_class=layers.DenseDropoutLayer,
name='dropout_mlp'):
if not isinstance(p, list):
p = [p]*(len(hidden_dims))
p = [p_input] + p
network_spec = mlp(input_dims, output_dims, hidden_dims, batchsize,
nonlinearities, output_nonlinearity, W_init, b_init,
name)
# first layer is input layer, we skip it
for i in range(len(p)):
layer_class, layer_args = network_spec[i+1]
if layer_class == DenseLayer and p[i] != 0:
layer_args['p'] = p[i]
network_spec[i+1] = (dropout_class, layer_args)
return network_spec
class BNN(BaseRegressor):
''' Bayesian neural net regressor '''
def __init__(self, idims, odims, n_samples=100,
heteroscedastic=True, name='BNN',
filename=None, network=None, network_spec=None,
likelihood=objectives.gaussian_log_likelihood,
**kwargs):
self.D = idims
self.E = odims
self.name = name
self.should_recompile = False
self.trained = False
self.heteroscedastic = heteroscedastic
self.likelihood = likelihood
sn = (np.ones((self.E,))*1e-3).astype(floatX)
sn = np.log(np.exp(sn)-1)
self.unconstrained_sn = theano.shared(sn, name='%s_sn' % (self.name))
eps = np.finfo(np.__dict__[floatX]).eps
self.sn = tt.nnet.softplus(self.unconstrained_sn) + eps
self.network = network
if type(network_spec) is list:
self.network_spec = network_spec
elif type(network_spec) is dict:
network_spec['input_dims'] = idims
odims_ = odims*2 if self.heteroscedastic else odims
network_spec['output_dims'] = odims_
network_spec['name'] = network_spec.get('name', self.name)
build_fn = network_spec.pop('build_fn', dropout_mlp)
self.network_spec = build_fn(**network_spec)
else:
self.network_spec = None
self.network_params = None
samples = np.array(n_samples).astype('int32')
samples_name = "%s>n_samples" % (self.name)
self.n_samples = theano.shared(samples, name=samples_name)
self.X = None
self.Y = None
self.Xm = None
self.iXs = None
self.Ym = None
self.Ys = None
# filename for saving
fname = '%s_%d_%d_%s_%s' % (self.name, self.D, self.E,
theano.config.device,
theano.config.floatX)
self.filename = fname if filename is None else filename
BaseRegressor.__init__(self, name=name, filename=self.filename)
if filename is not None:
self.load()
# optimizer options
max_evals = kwargs['max_evals'] if 'max_evals' in kwargs else 2000
conv_thr = kwargs['conv_thr'] if 'conv_thr' in kwargs else 1e-12
min_method = kwargs['min_method'] if 'min_method' in kwargs else 'ADAM'
self.optimizer = SGDOptimizer(min_method, max_evals,
conv_thr, name=self.name+'_opt')
# register theano shared variables for saving
self.register_types([tt.sharedvar.SharedVariable])
self.register(['sn', 'network_params', 'network_spec'])
def load(self, output_folder=None, output_filename=None):
n_samples = self.n_samples.get_value()
super(BNN, self).load(output_folder, output_filename)
self.n_samples.set_value(n_samples)
if self.network_params is None:
self.network_params = {}
if self.network_spec is not None:
self.build_network(self.network_spec,
params=self.network_params,
name=self.name)
if hasattr(self, 'unconstrained_sn'):
eps = np.finfo(np.__dict__[floatX]).eps
self.sn = tt.nnet.softplus(self.unconstrained_sn) + eps
def save(self, output_folder=None, output_filename=None):
# store references to the network shared variables, so we can save
# and load them correctly
self.network_params = []
for layer in lasagne.layers.get_all_layers(self.network):
layer_params = dict([(p.name.split('>')[-1], p)
for p in layer.get_params()])
self.network_params.append((layer.name, layer_params))
self.network_params = dict(self.network_params)
super(BNN, self).save(output_folder, output_filename)
def get_intermediate_outputs(self):
ret = super(BNN, self).get_intermediate_outputs()
ret += lasagne.layers.get_all_params(self.network, unwrap_shared=True)
for l in lasagne.layers.get_all_layers(self.network):
if hasattr(l, 'get_intermediate_outputs'):
ret += l.get_intermediate_outputs()
return list(set(ret))
def set_dataset(self, X_dataset, Y_dataset, **kwargs):
# set dataset
super(BNN, self).set_dataset(X_dataset.astype(floatX),
Y_dataset.astype(floatX))
# extra operations when setting the dataset (specific to this class)
self.update_dataset_statistics(X_dataset, Y_dataset)
if not self.trained:
# default log of measurement noise variance is set to 2.5% of
# dataset variation
s = (0.05*Y_dataset.std(0)).astype(floatX)
s = np.log(np.exp(s, dtype=floatX)-1.0)
self.unconstrained_sn.set_value(s)
def append_dataset(self, X_dataset, Y_dataset):
# set dataset
super(BNN, self).append_dataset(X_dataset.astype(floatX),
Y_dataset.astype(floatX))
# extra operations when setting the dataset (specific to this class)
self.update_dataset_statistics(self.X.get_value(), self.Y.get_value())
def update_dataset_statistics(self, X_dataset, Y_dataset):
# add small amount of noise for smoothing
X_dataset += 1e-6*np.random.randn(*X_dataset.shape)
Xm = np.atleast_1d(X_dataset.mean(0).astype(floatX))
Xc = np.atleast_2d(
(np.cov(X_dataset-Xm, rowvar=False, ddof=1).astype(floatX)))
iXs = np.linalg.cholesky(
np.linalg.inv(Xc)).astype(floatX)
if self.Xm is None:
self.Xm = theano.shared(Xm, name='%s>Xm' % (self.name))
self.iXs = theano.shared(iXs, name='%s>Xs' % (self.name))
else:
self.Xm.set_value(Xm)
self.iXs.set_value(iXs)
Ym = np.atleast_1d(Y_dataset.mean(0).astype(floatX))
Yc = np.atleast_2d(
np.cov(Y_dataset-Ym, rowvar=False, ddof=1).astype(floatX))
Ys = np.linalg.cholesky(Yc).T.astype(floatX)
if self.Ym is None:
self.Ym = theano.shared(Ym, name='%s>Ym' % (self.name))
self.Ys = theano.shared(Ys, name='%s>Ys' % (self.name))
else:
self.Ym.set_value(Ym)
self.Ys.set_value(Ys)
def build_network(self, network_spec=None, input_shape=None,
params={}, return_net=False, name=None):
''' Builds a network according to the specification in the
network_spec argument. network_spec should be a list containing
tuples where the first element is a class in lasagne.layers and the
second element is a dictionary with jeyword arguments for the class
constructor; i.e. [(layer_class_1, constructor_argss_1), ... ,
(layer_class_N, constructor_argss_N)].
Optionally, you can also pass a dictionary of parameters where the
keys are 'layer_name' and the values are dictionaries with the
trainable parameters indexed by names (e.g. W or b). The trainable
parameter values should be numpy arrays, theano shared variables, or
theano expressions to set the trainable parameters of the lasagne
layers. e.g params = {'layer_name_1': {'W': theano_shared_1,
'b': some_np_array}}
would set the weights of layer 1 to be the shared variable
theano_shared_1 and the biases to be initialized with the values of
some_np_array.'''
# set default values
if name is None:
name = self.name
if network_spec is None:
if self.network_spec is None:
idims = self.D
odims = self.E*2 if self.heteroscedastic else self.E
network_spec = dropout_mlp(
idims, odims, hidden_dims=[200]*4,
p=0.5, p_input=0.0,
nonlinearities=nonlinearities.rectify,
dropout_class=layers.DenseLogNormalDropoutLayer)
else:
network_spec = self.network_spec
utils.print_with_stamp('Building network', self.name)
self.network_spec = network_spec
# create input layer
assert network_spec[0][0] is lasagne.layers.InputLayer\
or input_shape is not None
if network_spec[0][0] is lasagne.layers.InputLayer:
if input_shape is not None:
# change the input shape
network_spec[0][1]['shape'] = input_shape
layer_class, layer_args = network_spec[0]
print(layer_class.__name__, layer_args)
network = layer_class(**layer_args)
network_spec = network_spec[1:]
else:
network = lasagne.layers.InputLayer(input_shape,
name=name+'_input')
# create hidden layers
for layer_class, layer_args in network_spec:
layer_name = layer_args['name']
if layer_name in params:
layer_args.update(params[layer_name])
print(layer_class.__name__, layer_args)
network = layer_class(network, **layer_args)
# change the periods in variable names
for p in network.get_params():
p.name = p.name.replace('.', '>')
# force rebuilding the prediction functions, as they will be
# out of date
self.predict_fn = None
self.predict_ic_fn = None
if return_net:
return network
else:
if self.network:
# remove old network params from this class params
params = lasagne.layers.get_all_params(
self.network, trainable=True)
self.remove_params(
[p.name for p in params if p.name is not None])
# add network parameters to this class params
params = lasagne.layers.get_all_params(network, trainable=True)
self.set_params(dict([(p.name, p) for p in params]))
self.network = network
self.update()
def get_loss(self):
''' initializes the loss function for training '''
# build the network
if self.network is None:
params = self.network_params\
if self.network_params is not None\
else {}
self.build_network(self.network_spec,
params=params,
name=self.name)
utils.print_with_stamp('Initialising loss function', self.name)
# Input variables
input_lengthscale = tt.scalar('%s>input_lengthscale' % (self.name))
hidden_lengthscale = tt.scalar('%s>hidden_lengthscale' % (self.name))
train_inputs = tt.matrix('%s>train_inputs' % (self.name))
train_targets = tt.matrix('%s>train_targets' % (self.name))
# evaluate nework output for batch
train_predictions, sn = self.predict(
train_inputs, None, deterministic=False,
iid_per_eval=True, return_samples=True)
# build the dropout loss function ( See Gal and Ghahramani 2015)
M = train_targets.shape[0].astype(theano.config.floatX)
N = self.X.shape[0].astype(theano.config.floatX)
# compute negative log likelihood
# note that if we have sn_std be a 1xD vector, broadcasting
# rules apply
lml = self.likelihood(
train_targets, train_predictions, sn)
# compute regularization term
# this is only for binary dropout layers
input_ls = input_lengthscale
hidden_ls = hidden_lengthscale
reg = objectives.dropout_gp_kl(
self.network, input_ls, hidden_ls)
# this is only for gaussian dropout layers
reg += objectives.gaussian_dropout_kl(
self.network, input_ls, hidden_ls)
# this is only for log normal dropout layers
reg += objectives.log_normal_kl(
self.network, input_ls, hidden_ls)
loss = -lml/M + reg/N
inputs = [train_inputs, train_targets,
input_lengthscale, hidden_lengthscale]
updates = theano.updates.OrderedUpdates()
# get trainable network parameters
params = lasagne.layers.get_all_params(self.network, trainable=True)
# if we are learning the noise
if not self.heteroscedastic:
params.append(self.unconstrained_sn)
self.set_params(dict([(p.name, p) for p in params]))
return loss, inputs, updates
def get_updates(self, network=None):
''' returns an updates dictionary, collected from layers in the
networks that provide the get_updates method'''
if network is None:
network = self.network
layer_updates = theano.updates.OrderedUpdates()
for l in lasagne.layers.get_all_layers(network):
if hasattr(l, 'get_updates'):
layer_updates += l.get_updates()
return layer_updates
def predict(self, mx, Sx=None, deterministic=False,
iid_per_eval=False, return_samples=False,
whiten_inputs=True, whiten_outputs=True, **kwargs):
''' returns symbolic expressions for the evaluations of this objects
neural network. If Sx is specified, the output will correspond to the
mean, covariance and input-output covariance of the network
predictions'''
# build the network if nedded
if self.network is None:
params = self.network_params\
if self.network_params is not None\
else {}
self.build_network(self.network_spec,
params=params,
name=self.name)
if Sx is not None:
# generate random samples from input (assuming gaussian
# distributed inputs)
# standard uniform samples (one sample per network sample)
z_std = utils.get_mrng().normal((self.n_samples, self.D))
# scale and center particles
Lx = tt.slinalg.cholesky(Sx)
x = mx + z_std.dot(Lx.T)
else:
x = mx[None, :] if mx.ndim == 1 else mx
if (whiten_inputs and hasattr(self, 'Xm') and self.Xm is not None):
# standardize inputs
x = (x - self.Xm).dot(self.iXs)
# unless we set the shared_axes parameter on the dropout layers,
# the noise samples should be different per input sample
ret = lasagne.layers.get_output(self.network, x,
deterministic=deterministic,
fixed_noise_samples=not iid_per_eval)
y = ret[:, :self.E]
sn = (0.1*tt.nnet.sigmoid(ret[:, self.E:])
if self.heteroscedastic
else tt.tile(self.sn, (y.shape[0], 1)))
# fudge factor
sn += 1e-6
if whiten_outputs and hasattr(self, 'Ym') and self.Ym is not None:
# scale and center outputs
y = y.dot(self.Ys) + self.Ym
# rescale variances
sn = sn*tt.diag(self.Ys)
y.name = '%s>output_samples' % (self.name)
if return_samples:
# nothing else to do!
return y, sn
else:
n = tt.cast(y.shape[0], dtype=theano.config.floatX)
# empirical mean
M = y.mean(axis=0)
# empirical covariance
deltay = y - M
S = deltay.T.dot(deltay)/(n-1)
# noise
S += tt.diag((sn**2).mean(axis=0))
# empirical input output covariance
if Sx is not None:
deltax = x - x.mean(0)
C = deltax.T.dot(deltay)/(n-1)
else:
C = tt.zeros((self.D, self.E))
return [M, S, C]
def update(self, n_samples=None):
''' Updates the dropout masks'''
if n_samples is not None:
if isinstance(n_samples, tt.sharedvar.SharedVariable):
self.n_samples = n_samples
self.update_fn = None
else:
self.n_samples.set_value(n_samples)
# increase the size of the masks
updts = self.get_updates()
for v, updt in updts.items():
v_shape = v.shape.eval()
if v_shape[0] != n_samples:
v_shape[0] = n_samples
v.set_value(np.zeros(v_shape, dtype=v.dtype))
if not hasattr(self, 'update_fn') or self.update_fn is None:
# get prediction with non deterministic samples
mx = tt.zeros((self.n_samples, self.D))
self.predict(mx, iid_per_eval=False)
# create a function to update the masks manually. Here the dropout
# masks should be shared variables
updts = self.get_updates()
# compile optimmized
mode = theano.compile.mode.get_mode('FAST_RUN')
self.update_fn = theano.function([], [], updates=updts,
allow_input_downcast=True,
mode=mode)
# draw samples from the networks
self.update_fn()
def train(self, batch_size=100,
input_ls=None, hidden_ls=None, lr=1e-4,
optimizer=None, callback=None):
if optimizer is None:
optimizer = self.optimizer
if optimizer.loss_fn is None or self.should_recompile:
loss, inps, updts = self.get_loss()
# we pass the learning rate as an input, and as a parameter to the
# updates method
learning_rate = theano.tensor.scalar('lr')
inps.append(learning_rate)
optimizer.set_objective(loss, self.get_params(symbolic=True),
inps, updts, learning_rate=learning_rate)
if input_ls is None:
# by default, be less strict with the input layer
input_ls = 1.0
if hidden_ls is None:
hidden_ls = 1.0
optimizer.minibatch_minimize(self.X.get_value(), self.Y.get_value(),
input_ls, hidden_ls, lr,
batch_size=batch_size,
callback=callback)
self.trained = True
self.update()
