def test_periodicity_group(capsys):
"""Test that groups are called at different periods."""
task_a = ExecuteCallback(lambda: print("a", end=" "))
task_b = ExecuteCallback(lambda: print("b", end=" "))
task_X = ExecuteCallback(lambda: print("X", end=" "))
group_often = MonitorTaskGroup([task_a, task_b], period=1)
group_seldom = MonitorTaskGroup([task_X], period=3)
monitor = Monitor(group_often, group_seldom)
for i in range(7):
monitor(i)
out, _ = capsys.readouterr()
expected = "a b X a b a b a b X a b a b a b X "
assert out == expected
# AutoGraph mode
compiled_monitor = tf.function(monitor)
for i in tf.range(7):
compiled_monitor(i)
# When using TF's range and compiling the monitoring we only expected the python prints once.
out, _ = capsys.readouterr()
assert "a b X"
def parse_fast_tasks(fast_tasks_period, training_loss, model, log_dir):
if fast_tasks_period > 0:
elbo_task = ScalarToTensorBoard(log_dir, training_loss, "elbo")
model_task = ModelToTensorBoard(log_dir, model)
return MonitorTaskGroup([model_task, elbo_task],
period=fast_tasks_period)
else:
return None
def test_ExecuteCallback_arguments(capsys):
def cb1(x=None, **_):
assert x is not None
print(x)
def cb2(**_):
print(2)
def cb3(y=None, **_):
assert y is not None
print(y)
group1 = MonitorTaskGroup([ExecuteCallback(cb1), ExecuteCallback(cb2)])
group2 = MonitorTaskGroup(ExecuteCallback(cb3))
monitor = Monitor(group1, group2)
monitor(0, x=1, y=3)
out, _ = capsys.readouterr()
assert out == "1\n2\n3\n"
def monitor(model, tmp_path):
tmp_path = str(tmp_path)
def lml_callback():
return model.log_marginal_likelihood()
def print_callback():
print("foo")
return Monitor(
MonitorTaskGroup(
[
ModelToTensorBoard(tmp_path, model),
ScalarToTensorBoard(tmp_path, lml_callback, "lml"),
],
period=2,
),
MonitorTaskGroup(ExecuteCallback(print_callback), period=1),
)
def tf_monitor_task_group(self, log_dir, slow_period=500):
ncols = 1
nrows_experts = self.num_experts
nrows_y = self.output_dim
image_task_experts_f = ImageToTensorBoard(
log_dir,
self.plot_experts_f,
name="experts_latent_function_posterior",
fig_kw={'figsize': (10, 4)},
subplots_kw={
'nrows': nrows_experts,
'ncols': self.output_dim
})
image_task_experts_y = ImageToTensorBoard(
log_dir,
self.plot_experts_y,
name="experts_output_posterior",
fig_kw={'figsize': (10, 4)},
subplots_kw={
'nrows': nrows_experts,
'ncols': self.output_dim
})
image_task_gating_gps = ImageToTensorBoard(
log_dir,
self.plot_gating_gps,
name="gating_network_gps_posteriors",
subplots_kw={
'nrows': nrows_experts,
'ncols': 1
})
image_task_gating = ImageToTensorBoard(
log_dir,
self.plot_gating_network,
name="gating_network_mixing_probabilities",
subplots_kw={
'nrows': 1,
'ncols': 1
})
image_task_y = ImageToTensorBoard(log_dir,
self.plot_y,
name="predictive_posterior_samples")
# image_tasks = [
# image_task_experts_y, image_task_experts_f, image_task_gating
# ]
image_tasks = [
image_task_experts_y, image_task_experts_f, image_task_gating_gps,
image_task_gating, image_task_y
]
return MonitorTaskGroup(image_tasks, period=slow_period)
def test_MonitorTaskGroup_and_Monitor(task_or_tasks):
group = MonitorTaskGroup(task_or_tasks, period=2)
# check that the tasks is actually a list (custom setter)
isinstance(group.tasks, list)
# Smoke test the __call__
group(0)
compiled_group = tf.function(group)
compiled_group(0)
# Smoke test the Monitor wrapper
monitor = Monitor(group)
monitor(0)
compiled_monitor = tf.function(monitor)
compiled_monitor(0)
def create_monitor(self, model):
model_task = ModelToTensorBoard(self.monitor_path, model)
self.monitor = Monitor(MonitorTaskGroup([model_task]), period=5)
# data_minibatch = (
# tf.data.Dataset.from_tensor_slices(data)
# .prefetch(autotune)
# .repeat()
# .shuffle(N)
# .batch(batch_size)
# )
#nat grad loop
# gamma_start = 1e-2 # deliberately chosen to be too large for this example
# gamma_max = 1e-1 # same max value as before
# gamma_step = 1e-2 # this is much more aggressive increase
# gamma = tf.Variable(gamma_start, dtype=tf.float64)
# gamma_incremented = tf.where(tf.less(gamma, gamma_max), gamma + gamma_step, gamma_max)
# op_ng = NatGradOptimizer(gamma).make_optimize_tensor(model, var_list=[[model.q_mu, model.q_sqrt]])
# op_adam = AdamOptimizer(0.001).make_optimize_tensor(model)
# op_increment_gamma = tf.assign(gamma, gamma_incremented)
# gamma_fallback = 1e-1 # we'll reduce by this factor if there's a cholesky failure
# op_fallback_gamma = tf.assign(gamma, gamma * gamma_fallback)
# sess.run(tf.variables_initializer([gamma]))
# for it in range(1000):
# try:
# sess.run(op_ng)
# sess.run(op_increment_gamma)
# except tf.errors.InvalidArgumentError:
# g = sess.run(gamma)
# print('gamma = {} on iteration {} is too big! Falling back to {}'.format(it, g, g * gamma_fallback))
# sess.run(op_fallback_gamma)
# sess.run(op_adam)
# if it % 100 == 0:
# print('{} gamma={:.4f} ELBO={:.4f}'.format(it, *sess.run([gamma, model.likelihood_tensor])))
def tf_monitor_task_group(self, log_dir, slow_period=500):
ncols = 2
image_task_f = ImageWithCbarToTensorBoard(
log_dir,
self.plot_f,
name="predict_f",
fig_kw={"figsize": (10, 4)},
subplots_kw={
"nrows": self.output_dim,
"ncols": ncols
},
)
image_task_y = ImageWithCbarToTensorBoard(
log_dir,
self.plot_y,
name="predict_y",
fig_kw={"figsize": (10, 4)},
subplots_kw={
"nrows": self.output_dim,
"ncols": ncols
},
)
image_tasks = [image_task_f, image_task_y]
return MonitorTaskGroup(image_tasks, period=slow_period)
def test_scipy_monitor_called(model):
task = DummyTask()
monitor = Monitor(MonitorTaskGroup(task, period=1))
opt = gpflow.optimizers.Scipy()
opt.minimize(model.training_loss, model.trainable_variables, step_callback=monitor)
assert task.current_step > 1
def print_cb(epoch_id=None, data=None):
tf.print(f"Epoch {epoch_id}: ELBO (train)", model.elbo(data))
def elbo_cb(data=None, **_):
return model.elbo(data)
output_logdir = enumerated_logdir()
model_task = ModelToTensorBoard(output_logdir, model)
elbo_task = ScalarToTensorBoard(output_logdir, elbo_cb, "elbo")
print_task = ExecuteCallback(callback=print_cb)
# We group these tasks and specify a period of `100` steps for them
fast_tasks = MonitorTaskGroup([model_task, elbo_task, print_task], period=100)
# We also want to see the model's fit during the optimisation
image_task = ImageToTensorBoard(output_logdir, plot_model, "samples_image")
# We typically don't want to plot too frequently during optimisation,
# which is why we specify a larger period for this task.
slow_taks = MonitorTaskGroup(image_task, period=500)
monitor = Monitor(fast_tasks, slow_taks)
def monitored_training_loop(epochs: int):
tf_optimization_step = tf.function(optimization_step)
batches = iter(train_dataset)
def tf_monitor_task_group(self, log_dir, slow_period=500):
ncols = 2
nrows_y = self.output_dim
nrows_experts = self.num_experts * self.output_dim
if self.num_experts > 2:
# num_gating_gps = self.num_experts * self.output_dim
num_gating_gps = self.num_experts
else:
# num_gating_gps = 1 * self.output_dim
num_gating_gps = 1
image_task_experts_f = ImageWithCbarToTensorBoard(
log_dir,
self.plot_experts_f,
name="experts_latent_function_posterior",
fig_kw={'figsize': (10, 4)},
subplots_kw={
'nrows': nrows_experts,
'ncols': ncols
})
image_task_experts_y = ImageWithCbarToTensorBoard(
log_dir,
self.plot_experts_y,
name="experts_output_posterior",
fig_kw={'figsize': (10, 4)},
subplots_kw={
'nrows': nrows_experts,
'ncols': ncols
})
image_task_gating_gps = ImageWithCbarToTensorBoard(
log_dir,
self.plot_gating_gps,
name="gating_network_gps_posteriors",
fig_kw={'figsize': (10, 2 * num_gating_gps)},
subplots_kw={
'nrows': num_gating_gps,
'ncols': 2
})
image_task_gating = ImageWithCbarToTensorBoard(
log_dir,
self.plot_gating_network,
name="gating_network_mixing_probabilities",
subplots_kw={
'nrows': self.num_experts,
'ncols': 1
# 'ncols': self.output_dim
})
image_task_y = ImageWithCbarToTensorBoard(
log_dir,
self.plot_y,
name="predictive_posterior_moment_matched",
fig_kw={'figsize': (10, 2)},
subplots_kw={
'nrows': nrows_y,
'ncols': ncols
})
# image_tasks = [
# image_task_experts_y, image_task_experts_f, image_task_gating
# ]
image_tasks = [
image_task_experts_y, image_task_experts_f, image_task_gating,
image_task_gating_gps, image_task_y
# image_task_y
]
# image_tasks = [
# image_task_experts_y, image_task_experts_f, image_task_y
# ]
return MonitorTaskGroup(image_tasks, period=slow_period)
# %%
log_dir = "logs" # Directory where TensorBoard files will be written.
model_task = ModelToTensorBoard(log_dir, model)
image_task = ImageToTensorBoard(log_dir, plot_prediction, "image_samples")
lml_task = ScalarToTensorBoard(log_dir, lambda: model.training_loss(),
"training_objective")
# %% [markdown]
# We now group the tasks in a set of fast and slow tasks and pass them to the monitor.
# This allows us to execute the groups at a different frequency.
# %%
# Plotting tasks can be quite slow. We want to run them less frequently.
# We group them in a `MonitorTaskGroup` and set the period to 5.
slow_tasks = MonitorTaskGroup(image_task, period=5)
# The other tasks are fast. We run them at each iteration of the optimisation.
fast_tasks = MonitorTaskGroup([model_task, lml_task], period=1)
# Both groups are passed to the monitor.
# `slow_tasks` will be run five times less frequently than `fast_tasks`.
monitor = Monitor(fast_tasks, slow_tasks)
# %%
training_loss = model.training_loss_closure(
compile=True) # compile=True (default): compiles using tf.function
opt = tf.optimizers.Adam()
for step in range(optimisation_steps):
opt.minimize(training_loss, model.trainable_variables)
def configure_tensorboard_monitor(self,
scalar_period,
imgs_period,
nb_images=1,
do_phase_space=None):
if do_phase_space is None:
do_phase_space = self.model.phase_space_dim == 2
if self.experiment.tensorboard_dir is None or scalar_period < 1:
return None
def create_bloss_tasks(directory):
bloss_names = [
'-ly', '-lx', 'penalty_term', 'alpha_term', '-H', '+KL'
]
bloss_tasks = []
def create_lambda(i):
return lambda train_bloss=None, **kwargs: train_bloss[i]
for i, name in enumerate(bloss_names):
bloss_tasks.append(
ScalarToTensorBoard(directory, create_lambda(i),
'bloss/' + name))
return bloss_tasks
train_dir = os.path.join(self.experiment.tensorboard_dir, 'train')
test_dir = os.path.join(self.experiment.tensorboard_dir, 'test')
# diff_task = ModelToTensorBoard(train_dir, self.model.sde_model.diffusion)
# drift_task = ModelToTensorBoard(train_dir, self.model.sde_model.drift_svgp)
diff_task = []
drift_task = []
train_loss = ScalarToTensorBoard(
train_dir, lambda train_loss=None, **kwargs: train_loss, 'loss')
test_loss = ScalarToTensorBoard(
test_dir,
lambda epoch=None, kl_scheduler=None, **kwargs: self.test_loss(
epoch, kl_scheduler),
'loss')
train_bloss_list = create_bloss_tasks(train_dir)
# train_bloss_list = [] # TODO: remove or add
generator = self.experiment.test_dataset if self.experiment.has_test else self.experiment.train_dataset
y_inputs = []
y_targets = []
for y in generator.take(1):
for y_input, y_target in self.tbptt_chunks_generator(y):
break
# y_inputs.append(y_input)
# y_targets.append(y_target)
# y_input = tf.concat(y_inputs, axis=1)
# y_target = tf.concat(y_targets, axis=1)
def calc_drift_error(**kwargs):
samples, entropies, encoded_dist, q0_stats, states = draw_fast_samples(
self.model, None, y_input)
fx, var_fx = self.model.sde_model.drift_svgp.predict_f(
tf.reshape(samples, (-1, samples.shape[-1])))
fx = tf.reshape(fx, samples.shape)
return tf.reduce_mean(
tf.square(samples[..., 1:, :] - samples[..., :-1, :] -
fx[..., :-1, :]))
drift_error = ScalarToTensorBoard(train_dir, calc_drift_error,
'drift_error')
beta_alpha = ScalarToTensorBoard(
train_dir, lambda **kwargs: tf.reduce_mean(
self.model.sde_model.diffusion.expected_diffusion()),
'beta_div_alpha')
if imgs_period > 0:
print('Creating image callbacks')
images_dir = os.path.join(self.experiment.tensorboard_dir,
'images')
nrows = 2 if self.model.phase_space_dim > 3 else 1
encoded_samples = ImageToTensorBoard(
images_dir,
lambda f, a: plot_encoded_samples(f, a, self.model, y_input),
'encoded_samples',
fig_kw={'figsize': (12, 12)},
subplots_kw={
'nrows': nrows,
'ncols': np.ceil(5 / 2).astype(int)
})
def plot_synth(fig, axes):
plot_synthetic_samples(fig,
axes,
self.model,
y_input,
y_target,
simulation_steps=y.shape[-2])
nrows = 2 if do_phase_space else 1
synthetic_samples = ImageToTensorBoard(
images_dir,
plot_synth,
'synthetic_samples',
fig_kw={'figsize': (12, 12)},
subplots_kw={
'nrows': nrows,
'ncols': nb_images
})
def plot_dec(fig, axes):
plot_decoder(fig, axes, self.model, y_input, y_target)
nrows = 2 if self.experiment.batch_size > 1 else 1
dec_images = ImageToTensorBoard(
images_dir,
plot_dec,
'decoder',
fig_kw={'figsize': (12, 12)},
subplots_kw={
'nrows': nrows,
'ncols': min(self.experiment.batch_size // nrows, 2)
})
drift_images = ImageToTensorBoard(
images_dir,
lambda fig, axes: plot_drift_predictions(
fig, axes, self.model, y_input),
'drift',
fig_kw={'figsize': (12, 12)},
subplots_kw={
'nrows': nrows,
'ncols': self.model.sde_model.dimension
})
monitor = Monitor(
MonitorTaskGroup([train_loss, test_loss] + train_bloss_list,
period=scalar_period),
# MonitorTaskGroup([drift_error, beta_alpha], period=scalar_period),
MonitorTaskGroup([
synthetic_samples, dec_images, encoded_samples,
drift_images
],
period=imgs_period))
print('done')
else:
monitor = Monitor(
MonitorTaskGroup([train_loss, test_loss] + train_bloss_list,
period=scalar_period),
MonitorTaskGroup([drift_error, beta_alpha],
period=scalar_period),
)
return monitor
