def test_fixed_preconditioner(self):
pmodel = sgmcmc_testlib.Normal2D(noise_sigma=1.0)
model = tf.keras.Sequential([pmodel])
model.build(input_shape=(1, 1))
var0 = model.trainable_variables[0]
optimizer = sgmcmc.NaiveSymplecticEulerMCMC(total_sample_size=1,
learning_rate=0.01,
momentum_decay=0.7,
preconditioner='fixed')
# Initial preconditioner: identity
precond_dict0 = {var0.name: 1.0}
optimizer.set_preconditioner_dict(precond_dict0,
model.trainable_variables)
sgmcmc_testlib.sample_model(model, optimizer, 2000)
self._check_kinetic_temperature_regions(model, optimizer)
# Adjust preconditioner
mom_old = tf.identity(optimizer.get_slot(var0, 'moments'))
precond_dict1 = {var0.name: 100.0}
optimizer.set_preconditioner_dict(precond_dict1,
model.trainable_variables)
mom_new = tf.identity(optimizer.get_slot(var0, 'moments'))
# Ensure moments are properly scaled and kinetic temperatures are ok
mom_new_target = tf.sqrt(100.0 / 1.0) * mom_old
self.assertAllClose(
mom_new,
mom_new_target,
msg='Moments not adjusted on preconditioner update.')
self._check_kinetic_temperature_regions(model, optimizer)
# Check kinetic temperature is ok after adjustment
sgmcmc_testlib.sample_model(model, optimizer, 2000)
self._check_kinetic_temperature_regions(model, optimizer)
def test_timestep_factor(self):
pmodel = sgmcmc_testlib.Normal2D(correlation=0.99,
noise_sigma=0.25,
uniform_noise=True)
model = tf.keras.Sequential([pmodel])
optimizer = sgmcmc.NaiveSymplecticEulerMCMC(total_sample_size=1,
learning_rate=0.01,
momentum_decay=0.9,
timestep_factor=1.0)
nburnin = 4096
nsamples = 262144
# Check that accuracy improves when we half the timestep_factor and
# efficiency as measured by ESS goes down by half.
kl_prev = None
efficiency_prev = None
for timestep_factor in [1.0, 0.5, 0.25, 0.125]:
optimizer.timestep_factor.assign(timestep_factor)
samples = sgmcmc_testlib.sample_model(model, optimizer, nburnin)
samples = sgmcmc_testlib.sample_model(model, optimizer, nsamples)
mean = np.mean(samples, axis=1)
cov = np.cov(samples)
_, efficiency_all = sgmcmc_testlib.compute_ess_multidimensional(
samples)
efficiency = np.min(efficiency_all)
kl = self._kl2d(pmodel, mean, cov)
name = optimizer.get_config()['name']
lr = optimizer.get_config()['learning_rate']
momentum_decay = optimizer.get_config().get('momentum_decay', -1.0)
dlr, dmomentum_decay = optimizer.dynamics_parameters(tf.float32)
dlr = float(dlr)
dmomentum_decay = float(dmomentum_decay)
logging.info(
'%s(lr=%.4f, momentum_decay=%.4f, timestep_factor=%.4f) => '
'(dlr=%.4f, dmomentum_decay=%.4f) => '
'mean (%.4f,%.4f) kl %.4f eff %.5f', name, lr,
momentum_decay, timestep_factor, dlr, dmomentum_decay, mean[0],
mean[1], kl, efficiency)
# Check values
if kl_prev is not None:
self.assertLess(
kl,
kl_prev + 0.006,
msg='Decreasing timestep_factor to %.4f increased KL '
'from %.5f to %.5f' % (timestep_factor, kl_prev, kl))
self.assertAlmostEqual(
efficiency / efficiency_prev,
0.5,
delta=0.05,
msg='Decreasing timestep_factor to %.4f '
'produced efficiency %.5f, compared to previous '
'efficiency of %.5f, but ratio %.5f not close '
'to 1/2.' % (timestep_factor, efficiency, efficiency_prev,
efficiency / efficiency_prev))
kl_prev = kl
efficiency_prev = efficiency
def main(argv):
del argv # unused arg
tf.io.gfile.makedirs(FLAGS.output_dir)
# Load data
tf.random.set_seed(DATASET_SEED)
if FLAGS.dataset == 'cifar10':
dataset_train, ds_info = datasets.load_cifar10(
tfds.Split.TRAIN, with_info=True,
data_augmentation=FLAGS.cifar_data_augmentation,
subsample_n=FLAGS.subsample_train_size)
dataset_test = datasets.load_cifar10(tfds.Split.TEST)
logging.info('CIFAR10 dataset loaded.')
elif FLAGS.dataset == 'imdb':
dataset, ds_info = datasets.load_imdb(
with_info=True, subsample_n=FLAGS.subsample_train_size)
dataset_train, dataset_test = datasets.get_generators_from_ds(dataset)
logging.info('IMDB dataset loaded.')
else:
raise ValueError('Unknown dataset {}.'.format(FLAGS.dataset))
# Prepare data for SG-MCMC methods
dataset_size = ds_info['train_num_examples']
dataset_size_orig = ds_info.get('train_num_examples_orig', dataset_size)
dataset_train = dataset_train.repeat().shuffle(10 * FLAGS.batch_size).batch(
FLAGS.batch_size)
test_batch_size = 100
validation_steps = ds_info['test_num_examples'] // test_batch_size
dataset_test_single = dataset_test.batch(FLAGS.batch_size)
dataset_test = dataset_test.repeat().batch(test_batch_size)
# If --pretend_batch_size flag is provided any cycle/epoch-length computation
# is done using this pretend_batch_size. Real batches are all still
# FLAGS.batch_size of length. This feature is used in the batch size ablation
# study.
#
# Also, always determine number of iterations from original data set size
if FLAGS.pretend_batch_size >= 1:
steps_per_epoch = dataset_size_orig // FLAGS.pretend_batch_size
else:
steps_per_epoch = dataset_size_orig // FLAGS.batch_size
# Set seed for the experiment
tf.random.set_seed(FLAGS.seed)
# Build model using pfac for proper priors
reg_weight = 1.0 / float(dataset_size)
if FLAGS.pfac == 'default':
pfac = priorfactory.DefaultPriorFactory(weight=reg_weight)
elif FLAGS.pfac == 'gaussian':
pfac = priorfactory.GaussianPriorFactory(prior_stddev=1.0,
weight=reg_weight)
else:
raise ValueError('Choose pfac from: default, gaussian.')
input_shape = ds_info['input_shape']
if FLAGS.model == 'cnnlstm':
assert FLAGS.dataset == 'imdb'
model = models.build_cnnlstm(ds_info['num_words'],
ds_info['sequence_length'],
pfac)
elif FLAGS.model == 'resnet':
assert FLAGS.dataset == 'cifar10'
model = models.build_resnet_v1(
input_shape=input_shape,
depth=20,
num_classes=ds_info['num_classes'],
pfac=pfac,
use_frn=FLAGS.resnet_use_frn,
use_internal_bias=FLAGS.resnet_bias)
else:
raise ValueError('Choose model from: cnnlstm, resnet.')
model.summary()
# Setup callbacks executed in keras.compile loop
callbacks = []
# Setup preconditioner
precond_dict = dict()
if FLAGS.use_preconditioner:
precond_dict['preconditioner'] = 'fixed'
logging.info('Use fixed preconditioner.')
else:
logging.info('No preconditioner is used.')
# Always append preconditioner callback to compute ctemp statistics
precond_estimator_cb = keras_utils.EstimatePreconditionerCallback(
gradest_train_fn,
iter(dataset_train),
every_nth_epoch=1,
batch_count=32,
raw_second_moment=True,
update_precond=FLAGS.use_preconditioner)
callbacks.append(precond_estimator_cb)
# Setup MCMC method
if FLAGS.method == 'sgmcmc':
# SG-MCMC optimizer, first-order symplectic Euler integrator
optimizer = sgmcmc.NaiveSymplecticEulerMCMC(
total_sample_size=dataset_size,
learning_rate=FLAGS.init_learning_rate,
momentum_decay=FLAGS.momentum_decay,
temp=FLAGS.temperature,
**precond_dict)
logging.info('Use symplectic Euler integrator.')
elif FLAGS.method == 'baoab':
# SG-MCMC optimizer, second-order accurate BAOAB integrator
optimizer = sgmcmc.BAOABMCMC(
total_sample_size=dataset_size,
learning_rate=FLAGS.init_learning_rate,
momentum_decay=FLAGS.momentum_decay,
temp=FLAGS.temperature,
**precond_dict)
logging.info('Use BAOAB integrator.')
else:
raise ValueError('Choose method from: sgmcmc, baoab.')
# Statistics for evaluation of ensemble performance
perf_stats = {
'ens_gce': stats.MeanStatistic(stats.ClassificationGibbsCrossEntropy()),
'ens_ce': stats.MeanStatistic(stats.ClassificationCrossEntropy()),
'ens_ce_sem': stats.StandardError(stats.ClassificationCrossEntropy()),
'ens_brier': stats.MeanStatistic(stats.BrierScore()),
'ens_brier_unc': stats.BrierUncertainty(),
'ens_brier_res': stats.BrierResolution(),
'ens_brier_reliab': stats.BrierReliability(),
'ens_ece': stats.ECE(10),
'ens_gacc': stats.MeanStatistic(stats.GibbsAccuracy()),
'ens_acc': stats.MeanStatistic(stats.Accuracy()),
'ens_acc_sem': stats.StandardError(stats.Accuracy()),
}
perf_stats_l, perf_stats_s = zip(*(perf_stats.items()))
# Setup ensemble
ens = ensemble.EmpiricalEnsemble(model, input_shape)
last_ens_eval = {'size': 0} # ensemble size from last evaluation
def cycle_ens_eval_maybe():
"""Ensemble evaluation callback, only evaluate at end of cycle."""
if len(ens) > last_ens_eval['size']:
last_ens_eval['size'] = len(ens)
logging.info('... evaluate ensemble on %d members', len(ens))
return ens.evaluate_ensemble(
dataset=dataset_test_single, statistics=perf_stats_s)
else:
return None
ensemble_eval_cb = keras_utils.EvaluateEnsemblePartial(
cycle_ens_eval_maybe, perf_stats_l)
callbacks.append(ensemble_eval_cb)
# Setup cyclical learning rate and temperature schedule for sgmcmc
if FLAGS.method == 'sgmcmc' or FLAGS.method == 'baoab':
# setup cyclical learning rate schedule
cyclic_sampler_cb = keras_utils.CyclicSamplerCallback(
ens,
FLAGS.cycle_length * steps_per_epoch, # number of iterations per cycle
FLAGS.cycle_start_sampling, # sampling phase start epoch
schedule=FLAGS.cycle_schedule,
min_value=0.0) # timestep_factor min value
callbacks.append(cyclic_sampler_cb)
# Setup temperature ramp-up schedule
begin_ramp_epoch = FLAGS.cycle_start_sampling - FLAGS.cycle_length
if begin_ramp_epoch < 0:
raise ValueError(
'cycle_start_sampling must be greater equal than cycle_length.')
ramp_iterations = FLAGS.cycle_length
tempramp_cb = keras_utils.TemperatureRampScheduler(
0.0, FLAGS.temperature, begin_ramp_epoch * steps_per_epoch,
ramp_iterations * steps_per_epoch)
# T0, Tf, begin_iter, ramp_epochs
callbacks.append(tempramp_cb)
# Additional callbacks
# Plot additional logs
def plot_logs(epoch, logs):
del epoch # unused
logs['lr'] = optimizer.get_config()['learning_rate']
if FLAGS.method == 'sgmcmc':
logs['timestep_factor'] = optimizer.get_config()['timestep_factor']
logs['ens_size'] = len(ens)
plot_logs_cb = tf.keras.callbacks.LambdaCallback(on_epoch_end=plot_logs)
# Write logs to tensorboard
tensorboard_cb = tf.keras.callbacks.TensorBoard(
log_dir=FLAGS.output_dir, write_graph=False)
# Output ktemp
diag_cb = keras_utils.PrintDiagnosticsCallback(10)
callbacks.extend([
diag_cb,
plot_logs_cb,
keras_utils.TemperatureMetric(),
keras_utils.SamplerTemperatureMetric(),
tensorboard_cb, # Should be after all callbacks that write logs
tf.keras.callbacks.CSVLogger(os.path.join(FLAGS.output_dir, 'logs.csv'))
])
# Keras train model
metrics = [
tf.keras.metrics.SparseCategoricalCrossentropy(
name='negative_log_likelihood',
from_logits=True),
tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy')]
model.compile(
optimizer,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=metrics)
logging.info('Model input shape: %s', model.input_shape)
logging.info('Model output shape: %s', model.output_shape)
logging.info('Model number of weights: %s', model.count_params())
model.fit(
dataset_train,
steps_per_epoch=steps_per_epoch,
epochs=FLAGS.train_epochs,
validation_data=dataset_test,
validation_steps=validation_steps,
callbacks=callbacks)
# Evaluate final ensemble performance
logging.info('Ensemble has %d members, computing final performance metrics.',
len(ens))
if ens.weights_list:
ens_perf_stats = ens.evaluate_ensemble(dataset_test_single, perf_stats_s)
print('Test set metrics:')
for label, stat_value in zip(perf_stats_l, ens_perf_stats):
stat_value = float(stat_value)
logging.info('%s: %.5f', label, stat_value)
print('%s: %.5f' % (label, stat_value))
# Add experiment info to experiment metadata csv file in *parent folder*
if FLAGS.write_experiment_metadata_to_csv:
csv_path = pathlib.Path.joinpath(
pathlib.Path(FLAGS.output_dir).parent, 'run_sweeps.csv')
data = {
'id': [FLAGS.experiment_id],
'seed': [FLAGS.seed],
'temperature': [FLAGS.temperature],
'dir': ['run_{}'.format(FLAGS.experiment_id)]
}
if tf.io.gfile.exists(csv_path):
sweeps_df = pd.read_csv(csv_path)
sweeps_df = sweeps_df.append(
pd.DataFrame.from_dict(data), ignore_index=True).set_index('id')
else:
sweeps_df = pd.DataFrame.from_dict(data).set_index('id')
# save experiment metadata csv file
sweeps_df.to_csv(csv_path)