def bbb_copy_slang(data_set='australian_presplit',
model_params=defaults.MLP_CLASS_PARAMETERS,
optimizer_params=defaults.BBB_COPYSLANG_PARAMETERS,
objective='avneg_elbo_bernoulli',
metrics=metric_factory.BAYES_BINCLASS,
normalize={
'x': False,
'y': False
},
save_params=defaults.SAVE_PARAMETERS,
seed=123,
use_cuda=torch.cuda.is_available(),
init_hook=None):
"""
Base class for BBB experiments that are to be initialized with the same model parameters as a specific SLANG experiment.
BBB_COPYSLANG_PARAMETERS are modified BBB parameters that specify the initialization of the SLANG experiment in question.
This experiment is useful for comparing the convergence rates of BBB and SLANG from the same initial model.
"""
set_seeds(seed)
data = Dataset(data_set=data_set, data_folder=DEFAULT_DATA_FOLDER)
model, predict_fn, kl_fn, closure_factory, optimizer = init_bbb_copy_slang_experiment(
data, model_params, optimizer_params, use_cuda=use_cuda)
if init_hook is not None:
init_hook(model, optimizer)
results_dict = run_experiment(data, model, model_params, predict_fn, kl_fn,
optimizer, optimizer_params, objective,
metrics, closure_factory, normalize,
save_params, seed, use_cuda)
return results_dict
def slang_cv(data_set='australian_presplit',
n_splits=10,
model_params=defaults.MLP_CLASS_PARAMETERS,
optimizer_params=defaults.SLANG_PARAMETERS,
objective='avneg_loglik_bernoulli',
metrics=[
'avneg_elbo_bernoulli', 'pred_avneg_loglik_bernoulli',
'sigmoid_predictive_accuracy'
],
normalize={
'x': False,
'y': False
},
save_params=defaults.SAVE_PARAMETERS,
seed=123,
use_cuda=torch.cuda.is_available()):
set_seeds(seed)
data = DatasetCV(data_set=data_set,
n_splits=n_splits,
seed=seed,
data_folder=DEFAULT_DATA_FOLDER)
results_dict = run_cv_experiment(data, n_splits, init_slang_experiment,
model_params, optimizer_params, objective,
metrics, normalize, save_params, seed,
use_cuda)
return results_dict
def bbb_base(data_set='australian_presplit',
model_params=defaults.MLP_CLASS_PARAMETERS,
optimizer_params=defaults.VI_PARAMETERS,
objective='avneg_elbo_bernoulli',
metrics=metric_factory.BAYES_BINCLASS,
normalize={
'x': False,
'y': False
},
save_params=defaults.SAVE_PARAMETERS,
seed=123,
use_cuda=torch.cuda.is_available(),
init_hook=None):
set_seeds(seed)
data = Dataset(data_set=data_set, data_folder=DEFAULT_DATA_FOLDER)
model, predict_fn, kl_fn, closure_factory, optimizer = init_bbb_experiment(
data, model_params, optimizer_params, use_cuda=use_cuda)
if init_hook is not None:
init_hook(model, optimizer)
results_dict = run_experiment(data, model, model_params, predict_fn, kl_fn,
optimizer, optimizer_params, objective,
metrics, closure_factory, normalize,
save_params, seed, use_cuda)
return results_dict
def bbb_cv(data_set='australian_presplit',
n_splits=10,
model_params=defaults.MLP_CLASS_PARAMETERS,
optimizer_params=defaults.VI_PARAMETERS,
objective='avneg_elbo_bernoulli',
metrics=metric_factory.BAYES_BINCLASS,
normalize={
'x': False,
'y': False
},
save_params=defaults.SAVE_PARAMETERS,
seed=123,
use_cuda=torch.cuda.is_available()):
set_seeds(seed)
data = DatasetCV(data_set=data_set,
n_splits=n_splits,
seed=seed,
data_folder=DEFAULT_DATA_FOLDER)
results_dict = run_cv_experiment(data, n_splits, init_bbb_experiment,
model_params, optimizer_params, objective,
metrics, normalize, save_params, seed,
use_cuda)
return results_dict
def uci_slang_bo(data_set='boston0',
n_splits=5,
model_params=UCI_BO_PARAMETERS,
optimizer_params=UCI_SLANG_PARAMETERS,
objective='avneg_loglik_gaussian',
metrics=['avneg_elbo_gaussian', 'pred_avneg_loglik_gaussian', 'predictive_rmse'],
normalize={'x': True, 'y': True},
save_params=MIN_SAVE_PARAMETERS,
param_bounds = {'log_noise_prec': (0, 5), 'log_prior_prec': (-4, 4)},
gp_params = {'kernel': Matern(nu=2.5, length_scale=[1, 2]), 'alpha': 1e-2},
bo_params = {'acq': 'ei', 'init_points': 5, 'n_iter': 25},
seed=123,
use_cuda=False):
set_seeds(seed)
data_cv = DatasetCV(data_set=data_set, n_splits=n_splits, seed=seed, data_folder=DEFAULT_DATA_FOLDER)
def get_cv_average(dict_of_dicts, key):
value = 0
for split, result_dict in dict_of_dicts.items():
value += result_dict["final_metrics"][key][0]
value = value/len(dict_of_dicts)
return value
def cv_exp(log_noise_prec, log_prior_prec):
model_params_cv = copy.deepcopy(model_params)
model_params_cv['prior_precision'] = math.exp(log_prior_prec)
model_params_cv['noise_precision'] = math.exp(log_noise_prec)
try:
# Run experiment
results_dicts = run_cv_experiment(data_cv, n_splits, init_slang_experiment, model_params_cv, optimizer_params, objective, metrics, normalize, save_params, seed, use_cuda)
# Return Avg. Test LL
logloss = get_cv_average(results_dicts, key='test_pred_logloss')
return -logloss
except:
logloss = 5.0
return -logloss
# Run BO
bo = BayesianOptimization(cv_exp, param_bounds, random_state=seed)
bo.maximize(init_points=bo_params['init_points'], n_iter=bo_params['n_iter'], acq=bo_params['acq'], **gp_params)
# Run final experiment
model_params_final = copy.deepcopy(model_params)
model_params_final['prior_precision'] = math.exp(bo.res['max']['max_params']['log_prior_prec'])
model_params_final['noise_precision'] = math.exp(bo.res['max']['max_params']['log_noise_prec'])
data = Dataset(data_set=data_set, data_folder=DEFAULT_DATA_FOLDER)
train_set_size=data.get_train_size()
model, predict_fn, kl_fn, closure_factory, optimizer = init_slang_experiment(data, model_params_final, optimizer_params, train_set_size=train_set_size, use_cuda=use_cuda)
save_params = defaults.SAVE_PARAMETERS
results_dict = run_experiment(data, model, model_params_final, predict_fn, kl_fn, optimizer, optimizer_params, objective, metrics, closure_factory, normalize, save_params, seed, use_cuda)
results = dict(final_run=results_dict, bo_results=bo.res)
return results
def slang_continue(data_set='mnist',
model_params=MNIST_MODEL_PARAM,
optimizer_params=MNIST_OPTIM_PARAM,
model_state_dict=None,
optimizer_state_dict=None,
new_train_set_size=60000,
objective='avneg_loglik_categorical',
metrics=['pred_avneg_loglik_categorical', 'softmax_predictive_accuracy', 'avneg_elbo_categorical'],
normalize={'x': False, 'y': False},
save_params=defaults.SAVE_PARAMETERS,
seed=123,
use_cuda=torch.cuda.is_available(),
init_hook=None,
iter_hook=None,
end_hook=None):
set_seeds(seed)
data = Dataset(data_set=data_set, data_folder=DEFAULT_DATA_FOLDER)
train_set_size=data.get_train_size()
model, predict_fn, kl_fn, closure_factory, optimizer = init_slang_experiment(data, model_params, optimizer_params, train_set_size=train_set_size, use_cuda=use_cuda)
if model_state_dict is None or optimizer_state_dict is None:
raise ValueError("Previous experiment values are not loaded. Records from the previous experiment are needed.")
model.load_state_dict(model_state_dict)
if use_cuda:
model = model.cuda()
optimizer.load_state_dict(optimizer_state_dict)
optimizer.defaults['train_set_size'] = new_train_set_size
if use_cuda:
optimizer.defaults['prior_prec'] = optimizer.defaults['prior_prec'].cuda()
optimizer.state['mean'] = optimizer.state['mean'].cuda()
optimizer.state['momentum_grad'] = optimizer.state['momentum_grad'].cuda()
optimizer.state['prec_diag'] = optimizer.state['prec_diag'].cuda()
optimizer.state['prec_factor'] = optimizer.state['prec_factor'].cuda()
if init_hook is not None:
init_hook(model, optimizer)
results_dict = run_experiment(data, model, model_params, predict_fn, kl_fn, optimizer, optimizer_params, objective, metrics, closure_factory, normalize, save_params, seed, use_cuda, iter_hook)
if end_hook is not None:
end_hook(results_dict, model, optimizer)
return results_dict
def slang_base(data_set='australian_presplit',
model_params=defaults.MLP_CLASS_PARAMETERS,
optimizer_params=defaults.SLANG_PARAMETERS,
objective='avneg_loglik_bernoulli',
metrics=[
'avneg_elbo_bernoulli', 'pred_avneg_loglik_bernoulli',
'pred_avneg_loglik_bernoulli', 'sigmoid_predictive_accuracy'
],
normalize={
'x': False,
'y': False
},
save_params=defaults.SAVE_PARAMETERS,
seed=123,
use_cuda=torch.cuda.is_available(),
init_hook=None,
iter_hook=None,
end_hook=None):
set_seeds(seed)
data = Dataset(data_set=data_set, data_folder=DEFAULT_DATA_FOLDER)
train_set_size = data.get_train_size()
model, predict_fn, kl_fn, closure_factory, optimizer = init_slang_experiment(
data,
model_params,
optimizer_params,
train_set_size=train_set_size,
use_cuda=use_cuda)
if init_hook is not None:
init_hook(model, optimizer)
results_dict = run_experiment(data, model, model_params, predict_fn, kl_fn,
optimizer, optimizer_params, objective,
metrics, closure_factory, normalize,
save_params, seed, use_cuda, iter_hook)
if end_hook is not None:
end_hook(results_dict, model, optimizer)
return results_dict
def get_accuracy(experiment_base,
experiment_name,
variant,
mc_10_multiplier,
data_set=None):
set_seeds(123)
# Retrieve results
exp = experiment_base.get_variant(experiment_name).get_variant(variant)
record = exp.get_latest_record()
if record.has_result():
result = record.get_result()
# Instantiate model and optimizer
params = exp.get_args()
if data_set is None:
data_set = params['data_set']
data = Dataset(data_set=data_set, data_folder=DEFAULT_DATA_FOLDER)
train_set_size = data.get_train_size()
model, predict_fn, kl_fn, closure_factory, optimizer = init_slang_experiment(
data=data,
model_params=params['model_params'],
optimizer_params=params['optimizer_params'],
train_set_size=train_set_size,
use_cuda=torch.cuda.is_available())
# Get state dicts
model_state_dict = result['model']
optim_state_dict = result['optimizer']
# Set model and optimizer state
model.load_state_dict(model_state_dict)
optimizer.load_state_dict(optim_state_dict)
# Make predictions
test_x, test_y = data.load_full_test_set()
for i in range(mc_10_multiplier):
print(i)
with torch.no_grad():
pred_y = predict_fn(test_x, 10)
if i == 0:
preds = pred_y
else:
preds = torch.cat([preds, pred_y], 0)
# Compute accuracy
pred_acc = softmax_predictive_accuracy(preds, test_y).item()
else:
pred_acc = np.nan
return pred_acc
def predict_fn(x, mc_samples):
noise = optimizer.distribution().rsample(mc_samples).t()
preds = model(x, noise, False)
if output_size == 1:
preds = preds.reshape(mc_samples, -1)
return preds
def kl_fn():
return optimizer.kl_divergence()
##################
### Experiment ###
##################
set_seeds(123)
# Set objective and metrics:
objective = 'avneg_loglik_bernoulli'
metrics = ['avneg_elbo_bernoulli', 'pred_avneg_loglik_bernoulli', 'pred_avneg_loglik_bernoulli', 'sigmoid_predictive_accuracy']
normalize = {'x': False, 'y': False}
objective = metric_factory.make_objective_closure("avneg_elbo_bernoulli", kl_fn, N)
metrics = metric_factory.BAYES_BINCLASS
train_metrics, test_metrics = metric_factory.make_metric_closures(metrics, kl_fn, N)
#####################################
########### Training Loop ###########
#####################################
metric_history = {}
def slang_complete(val_data_set='mnist_val',
continue_data_set='mnist',
model_params=MNIST_MODEL_PARAM,
optimizer_params=MNIST_OPTIM_PARAM,
continue_train_set_size=60000,
num_continues=3,
objective='avneg_loglik_categorical',
metrics=[
'pred_avneg_loglik_categorical',
'softmax_predictive_accuracy', 'avneg_elbo_categorical'
],
normalize={
'x': False,
'y': False
},
save_params=defaults.SAVE_PARAMETERS,
val_seed=123,
continue_seeds=[123, 123, 123],
use_cuda=torch.cuda.is_available(),
init_hook=None,
iter_hook=None,
end_hook=None):
######################################
### Run Validation-set Experiment: ###
######################################
set_seeds(val_seed)
data = Dataset(data_set=val_data_set, data_folder=DEFAULT_DATA_FOLDER)
train_set_size = data.get_train_size()
model, predict_fn, kl_fn, closure_factory, optimizer = init_slang_experiment(
data,
model_params,
optimizer_params,
train_set_size=train_set_size,
use_cuda=use_cuda)
if init_hook is not None:
init_hook(model, optimizer)
results_dict = run_experiment(data, model, model_params, predict_fn, kl_fn,
optimizer, optimizer_params, objective,
metrics, closure_factory, normalize,
save_params, val_seed, use_cuda, iter_hook)
if end_hook is not None:
end_hook(results_dict, model, optimizer)
#####################################
### Run Continuation Experiments: ###
#####################################
for continuation in range(num_continues):
set_seeds(continue_seeds[continuation])
data = Dataset(data_set=continue_data_set,
data_folder=DEFAULT_DATA_FOLDER)
train_set_size = data.get_train_size()
model, predict_fn, kl_fn, closure_factory, optimizer = init_slang_experiment(
data,
model_params,
optimizer_params,
train_set_size=train_set_size,
use_cuda=use_cuda)
results_dict['optimizer']['state']['mean'] = results_dict['optimizer'][
'state']['mean'].detach()
model, optimizer = load_dicts(model, optimizer, results_dict,
continue_train_set_size, use_cuda)
if init_hook is not None:
init_hook(model, optimizer)
results_dict = run_experiment(data, model, model_params, predict_fn,
kl_fn, optimizer, optimizer_params,
objective, metrics, closure_factory,
normalize, save_params,
continue_seeds[continuation], use_cuda,
iter_hook)
if end_hook is not None:
end_hook(results_dict, model, optimizer)
return results_dict