class TestBohamiann(unittest.TestCase):
def setUp(self):
self.X = np.random.rand(10, 3)
self.y = np.sinc(self.X * 10 - 5).sum(axis=1)
self.model = Bohamiann(normalize_input=True, normalize_output=True, use_double_precision=True)
self.model.train(self.X, self.y, num_burn_in_steps=20, num_steps=100, keep_every=10)
def test_predict(self):
X_test = np.random.rand(10, self.X.shape[1])
m, v = self.model.predict(X_test)
assert len(m.shape) == 1
assert m.shape[0] == X_test.shape[0]
assert len(v.shape) == 1
assert v.shape[0] == X_test.shape[0]
def test_gradient_mean(self):
X_test = np.random.rand(10, self.X.shape[1])
def wrapper(x):
return self.model.predict([x])[0]
def wrapper_grad(x):
return self.model.predictive_mean_gradient(x)
grad = self.model.predictive_mean_gradient(X_test[0])
assert grad.shape[0] == X_test.shape[1]
for xi in X_test:
err = check_grad(wrapper, wrapper_grad, xi, epsilon=1e-6)
assert err < 1e-5
def test_gradient_variance(self):
X_test = np.random.rand(10, self.X.shape[1])
def wrapper(x):
v = self.model.predict([x])[1]
return v
def wrapper_grad(x):
return self.model.predictive_variance_gradient(x)
grad = self.model.predictive_variance_gradient(X_test[0])
assert grad.shape[0] == X_test.shape[1]
for xi in X_test:
err = check_grad(wrapper, wrapper_grad, xi, epsilon=1e-6)
assert err < 1e-5
class WrapperBohamiann(BaseModel):
def __init__(self, get_net=get_default_network, lr=1e-5, use_double_precision=False, verbose=False):
"""
Wrapper around pybnn Bohamiann implementation. It automatically adjusts the length by the MCMC chain,
by performing 100 times more burnin steps than we have data points and sampling ~100 networks weights.
Parameters
----------
get_net: func
Architecture specification
lr: float
The MCMC step length
use_double_precision: Boolean
Use float32 or float64 precision. Note: Using float64 makes the training slower.
verbose: Boolean
Determines whether to print pybnn output.
"""
self.lr = lr
self.verbose = verbose
self.bnn = Bohamiann(get_network=get_net, use_double_precision=use_double_precision)
def train(self, X, y, **kwargs):
self.X = X
self.y = y
self.bnn.train(X, y, lr=self.lr,
num_burn_in_steps=X.shape[0] * 100,
num_steps=X.shape[0] * 100 + 10000, verbose=self.verbose)
def predict(self, X_test):
return self.bnn.predict(X_test)
class BayesianNN(object):
def __init__(self):
self.model = Bohamiann(print_every_n_steps=1000,
sampling_method="adaptive_sghmc")
self.trained = False
def fit(self, x, y):
self.model.train(x,
y.flatten(),
num_steps=10000 + 100 * len(x),
num_burn_in_steps=100 * len(x),
keep_every=200,
lr=1e-2,
verbose=True,
continue_training=self.trained)
def predict(self, x):
mean, var = self.model.predict(x)
return mean, 1.96 * np.sqrt(var)
class OrionBohamiannWrapper(BaseModel):
"""
Wrapper for PyBNN's BOHAMIANN model
Parameters
----------
normalize_input: bool
Normalize the input based on the provided bounds (zero mean and unit standard deviation).
Defaults to ``True``.
normalize_output: bool
Normalize the output based on data (zero mean and unit standard deviation).
Defaults to ``False``.
burnin_steps: int or None.
The number of burnin steps before the sampling procedure starts.
If ``None``, ``burnin_steps = n_dims * 100`` where ``n_dims`` is the dimensionality
of the search space. Defaults to ``None``.
sampling_method: str
Can be one of ``['adaptive_sghmc', 'sgld', 'preconditioned_sgld', 'sghmc']``.
Defaults to ``"adaptive_sghmc"``. See PyBNN samplers'
`code <https://github.com/automl/pybnn/tree/master/pybnn/sampler>`_ for more information.
use_double_precision: bool
Use double precision if using ``bohamiann``. Note that it can run faster on GPU
if using single precision. Defaults to ``True``.
num_steps: int or None
Number of sampling steps to perform after burn-in is finished.
In total, ``num_steps // keep_every`` network weights will be sampled.
If ``None``, ``num_steps = n_dims * 100 + 10000`` where ``n_dims`` is the
dimensionality of the search space.
keep_every: int
Number of sampling steps (after burn-in) to perform before keeping a sample.
In total, ``num_steps // keep_every`` network weights will be sampled.
learning_rate: float
Learning rate. Defaults to 1e-2.
batch_size: int
Batch size for training the neural network. Defaults to 20.
epsilon: float
epsilon for numerical stability. Defaults to 1e-10.
mdecay: float
momemtum decay. Defaults to 0.05.
verbose: bool
Write progress logs in stdout. Defaults to ``False``.
"""
def __init__(self,
lower,
upper,
sampling_method="adaptive_sghmc",
use_double_precision=True,
num_steps=None,
keep_every=100,
burnin_steps=None,
learning_rate=1e-2,
batch_size=20,
epsilon=1e-10,
mdecay=0.05,
verbose=False,
**kwargs):
self.num_steps = num_steps
self.keep_every = keep_every
self.burnin_steps = burnin_steps
self.learning_rate = learning_rate
self.batch_size = batch_size
self.epsilon = epsilon
self.mdecay = mdecay
self.verbose = verbose
self.bnn = Bohamiann(get_network=get_default_network,
sampling_method=sampling_method,
use_double_precision=use_double_precision,
**kwargs)
self.burnin_steps = burnin_steps
self.lower = lower
self.upper = upper
# pylint:disable=no-self-use
def set_state(self, state_dict):
"""Restore the state of the optimizer"""
torch.random.set_rng_state(state_dict["torch"])
# pylint:disable=no-self-use
def state_dict(self):
"""Return the current state of the optimizer so that it can be restored"""
return {"torch": torch.random.get_rng_state()}
def seed(self, seed):
"""Seed all internal RNGs"""
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = False
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.manual_seed(seed)
def train(self, X, y, **kwargs):
"""
Sets num_steps and burnin_steps before training with parent's train()
"""
self.X = X
self.y = y
if self.num_steps:
num_steps = self.num_steps
else:
num_steps = X.shape[0] * 100 + 10000
if self.burnin_steps is None:
burnin_steps = X.shape[0] * 100
else:
burnin_steps = self.burnin_steps
self.bnn.train(X,
y,
num_steps=num_steps,
keep_every=self.keep_every,
num_burn_in_steps=burnin_steps,
lr=self.learning_rate,
batch_size=self.batch_size,
epsilon=self.epsilon,
mdecay=self.mdecay,
continue_training=False,
verbose=self.verbose,
**kwargs)
def predict(self, X_test):
"""Predict using bnn.predict()"""
return self.bnn.predict(X_test)
class BOHAMIANNWarp(BaseModel):
"""
A Wrapper for MC Dropout for a fully connected
feed forward neural network..
"""
def __init__(self,
num_samples=6000,
keep_every=50,
lr=1e-2,
normalize_input: bool = True,
normalize_output: bool = True,
verbose=True,
seed=42):
self.verbose = verbose
self.num_samples = num_samples
self.keep_every = keep_every
self.lr = lr
self.model = Bohamiann(normalize_input=normalize_input,
normalize_output=normalize_output,
seed=seed)
def _create_model(self, X, Y):
Y = Y.flatten()
num_burn_in_steps = X.shape[0] * 100
num_steps = X.shape[0] * 100 + self.num_samples
self.model.train(X,
Y,
num_steps=num_steps,
num_burn_in_steps=num_burn_in_steps,
keep_every=self.keep_every,
lr=self.lr,
verbose=self.verbose)
def _update_model(self, X_all, Y_all):
"""
Updates the model with new observations.
"""
Y_all = Y_all.flatten()
num_burn_in_steps = X_all.shape[0] * 100
num_steps = X_all.shape[0] * 100 + self.num_samples
if self.model is None:
self._create_model(X_all, Y_all)
else:
self.model.train(X_all,
Y_all,
num_steps=num_steps,
num_burn_in_steps=num_burn_in_steps,
keep_every=self.keep_every,
lr=self.lr,
verbose=self.verbose)
def predict(self, X):
"""
Predictions with the model. Returns predictive means and standard deviations at X.
"""
X = np.atleast_2d(X)
m, v = self.model.predict(X)
# m and v have shape (N,)
s = np.sqrt(v)
return m[:, None], s[:, None]
def predict_withGradients(self, X):
"""
Returns the mean, standard deviation, mean gradient and standard deviation gradient at X.
"""
return print('Not Implemented')
plt.xlim(0, 1)
plt.show()
# -- Train Model ---
model = Bohamiann(print_every_n_steps=1000)
model.train(x[:, None],
y,
num_steps=20000,
num_burn_in_steps=2000,
keep_every=50,
lr=1e-2,
verbose=True)
# -- Predict with Model ---
m, v = model.predict(grid[:, None])
plt.plot(x, y, "ro")
plt.grid()
plt.plot(grid, fvals, "k--")
plt.plot(grid, m, "blue")
plt.fill_between(grid,
m + np.sqrt(v),
m - np.sqrt(v),
color="orange",
alpha=0.8)
plt.fill_between(grid,
m + 2 * np.sqrt(v),
m - 2 * np.sqrt(v),
color="orange",
alpha=0.6)
plt.fill_between(grid,