def compute_train_elbos(results_root: str, name: str) -> List[float]:
with tf.Graph().as_default() as graph:
with tf.Session(graph=graph).as_default():
args = experiment_common.parse_args_from_csv(name, results_root)
checkpoint_path, _ = experiment_common.create_paths(
name, results_root)
data = experiment_common.get_data(args)
# Disable minibatching so we just compute the ELBO over the entire dataset at once.
args.minibatch_size = None
model = build_model(args, data.X_train, data.Y_train)
sess = model.enquire_session()
model.init_op(sess)
step = 300000
saver = tf.train.Saver(var_list=[
v for v in tf.all_variables() if "dataholder" not in v.name
])
checkpoint_file = _get_checkpoint_file(checkpoint_path, step)
saver.restore(sess, checkpoint_file)
assert sess.run(model.global_step) == step
return [
sess.run(model.likelihood_tensor).item() for _ in range(10)
]
def compute_tlls(results_root: str, name: str) -> List[float]:
with tf.Graph().as_default() as graph:
with tf.Session(graph=graph).as_default():
args = experiment_common.parse_args_from_csv(name, results_root)
checkpoint_path, _ = experiment_common.create_paths(
name, results_root)
data = experiment_common.get_data(args)
model = build_model(args, data.X_train, data.Y_train)
sess = model.enquire_session()
model.init_op(sess)
step = 300000
saver = tf.train.Saver()
checkpoint_file = _get_checkpoint_file(checkpoint_path, step)
saver.restore(sess, checkpoint_file)
assert sess.run(model.global_step) == step
return [
metrics.compute_log_likelihood(model,
data.X_test,
data.Y_test,
n_samples=10000)
for _ in range(10)
]
def _sample_ys_and_plot_histogram(ax, results_root: str, name: str,
point_i: int):
with tf.Graph().as_default() as graph:
with tf.Session(graph=graph).as_default():
args = experiment_common.parse_args_from_csv(
name, results_root)
checkpoint_path, _ = experiment_common.create_paths(
name, results_root)
data = experiment_common.get_data(args)
# Disable minibatching so we just compute the ELBO over the entire dataset at once.
args.minibatch_size = None
model = build_model(args, data.X_train, data.Y_train)
sess = model.enquire_session()
model.init_op(sess)
step = 300000
saver = tf.train.Saver(var_list=[
v for v in tf.all_variables() if "dataholder" not in v.name
])
checkpoint_file = _get_checkpoint_file(checkpoint_path, step)
saver.restore(sess, checkpoint_file)
assert sess.run(model.global_step) == step
n_samples = 10000
ys = model.predict_y_samples(
data.X_test[point_i].reshape(1, -1), n_samples)
if "fixq" in name:
color = highlight_color
else:
color = non_dreg_color
ax.hist(ys.reshape(-1), bins=200, density=True, color=color)
def _sample_and_save_gradients(
name: str,
batch_size: int,
Ks: List[int],
num_samples: int,
results_path: Optional[str],
step_to_load: Optional[int],
dreg: bool,
gpu: bool,
):
print(f"Loading {name}")
args = parse_args_from_csv(name, results_path)
checkpoint_path, _ = experiment_common.create_paths(name, results_path)
file_prefix = _get_file_prefix("gradients", batch_size, dreg)
print("Computing gradients...")
for K in _get_Ks_to_compute(Ks, checkpoint_path, file_prefix):
gradients = _sample_gradients_from_checkpoint(
args,
checkpoint_path,
step_to_load,
batch_size,
batch_range=(0, 1),
K=K,
dreg=dreg,
gpu=gpu,
num_samples=num_samples,
)
file_path = os.path.join(checkpoint_path, f"{file_prefix}k{K}.npy")
np.save(file_path, gradients)
def table_1(results_root: str, metric: str) -> None:
datasets = [
"wilson_forest", "power", "wilson_pol", "wilson_solar", "winewhite",
"winered"
]
values = defaultdict(lambda: {False: {}, True: {}})
for dataset in datasets:
n_splits = 0
for split in itertools.count(0, 1):
if split > 40:
raise ValueError(f"Not enough splits for {dataset}")
reg_value = None
dreg_value = None
for dreg in [False, True]:
suffix = "_noise0.2" if dataset == "demo_normalized" else ""
dreg_text = "_dreg" if dreg else ""
name = f"{dataset}_L1_G5_IWAE_SF{split}_K50_BS64{dreg_text}{suffix}"
checkpoint_path, _ = experiment_common.create_paths(
name, results_root)
final_iter_file = os.path.join(checkpoint_path,
f"{name}_300000.csv")
if not os.path.exists(final_iter_file):
# print(f"split {split} {dataset} {dreg} missing")
continue
if metric == "elbo":
value = compute_train_elbos(results_root, name)
elif metric == "tll":
value = compute_tlls(results_root, name)
else:
raise ValueError
if dreg:
dreg_value = value
else:
reg_value = value
if reg_value is None or dreg_value is None:
continue
values[dataset][False][split] = reg_value
values[dataset][True][split] = dreg_value
n_splits += 1
if n_splits == 20:
break
with open(f"figure_data/table1_{metric}.txt", mode="w") as results_file:
for dataset in datasets:
for dreg in [False, True]:
dreg_str = "T" if dreg else "F"
for split, results in sorted(values[dataset][dreg].items()):
results_str = " ".join([str(r) for r in results])
results_file.write(
f"{dataset} {dreg_str} {split} {results_str}")
def _compute_and_save_aggregates(
name: str,
batch_size: int,
num_datapoints: int,
Ks: List[int],
num_samples: int,
results_path: Optional[str],
step_to_load: Optional[int],
dreg: bool,
gpu: bool,
):
print(f"Loading {name}")
args = parse_args_from_csv(name, results_path)
checkpoint_path, _ = experiment_common.create_paths(name, results_path)
file_prefix = _get_file_prefix("gradags", batch_size, dreg)
if num_datapoints != 1 and batch_size == 1:
datapoints = range(num_datapoints)
batch_ranges = [(i, i + 1) for i in datapoints]
elif num_datapoints == 1:
datapoints = [None]
batch_ranges = [(0, batch_size)]
else:
raise ValueError(
"Cannot aggregate snrs over more than one datapoint if batch size is not 1"
)
print("Computing snrs...")
for K in _get_Ks_to_compute(Ks, checkpoint_path, file_prefix):
for datapoint, batch_range in zip(datapoints, batch_ranges):
print(
f"Computing for datapoint {datapoint}, batch range {batch_range}"
)
gradients = _sample_gradients_from_checkpoint(
args,
checkpoint_path,
step_to_load,
batch_size,
batch_range,
K,
dreg,
gpu,
num_samples,
)
aggregates: Dict[str, Dict[str, ndarray]] = {"mean": {}, "std": {}}
for param, grads in gradients.items():
aggregates["mean"][param] = np.mean(grads, axis=0)
aggregates["std"][param] = np.std(grads, axis=0)
if np.any(aggregates["mean"][param] == 0.0):
print(f"Gradient mean was zero for {param} at K={K}")
datapoint_suffix = "" if num_datapoints == 1 else f"_dp{datapoint}"
file_path = os.path.join(
checkpoint_path, f"{file_prefix}k{K}{datapoint_suffix}.npy")
np.save(file_path, aggregates)
def figure_4(results_root: str) -> None:
datasets = [
"wilson_forest", "power", "wilson_pol", "wilson_solar", "winewhite",
"winered"
]
want_n_splits = 10
values = defaultdict(lambda: {False: {}, True: {}})
for dataset in datasets:
n_splits = 0
for split in itertools.count(0, 1):
if split > 40:
raise ValueError(f"Not enough splits for {dataset}")
reg_value = None
fixq_value = None
for fixq in [False, True]:
fixq_text = "_fixq" if fixq else ""
name = f"{dataset}_L1_G5_IWAE_SF{split}_K50_BS64{fixq_text}"
checkpoint_path, _ = experiment_common.create_paths(
name, results_root)
final_iter_file = os.path.join(checkpoint_path,
f"{name}_300000.csv")
if not os.path.exists(final_iter_file):
continue
if fixq:
fixq_value = compute_tlls(results_root, name)
else:
reg_value = compute_tlls(results_root, name)
if reg_value is None or fixq_value is None:
continue
values[dataset][False][split] = reg_value
values[dataset][True][split] = fixq_value
n_splits += 1
if n_splits == want_n_splits:
break
with open(f"figure_data/figure_4.txt", mode="w") as results_file:
for dataset in datasets:
for fixq in [False, True]:
fixq_str = "T" if fixq else "F"
for split, results in sorted(values[dataset][fixq].items()):
results_str = " ".join([str(r) for r in results])
results_file.write(
f"{dataset} {fixq_str} {split} {results_str}")
def main():
#################################### args
ARGS = experiment_common.parse_arguments()
print("\n", "ARGS:", "\n", ARGS, "\n")
if ARGS.plot_freq is not None and not ARGS.dataset.startswith("demo"):
raise ValueError("Plotting only supported for demo dataset.")
random.seed(ARGS.seed)
np.random.seed(ARGS.seed)
tf.set_random_seed(ARGS.seed)
#################################### paths
file_name = experiment_common.get_file_name(ARGS)
checkpoint_path, tensorboard_path = experiment_common.create_paths(
file_name, ARGS.results_path
)
#################################### data
data = experiment_common.get_data(ARGS)
#################################### model
model = build_model(ARGS, data.X_train, data.Y_train)
#################################### init
sess = model.enquire_session()
model.init_op(sess)
#################################### monitoring
def _write_dict_to_csv(data: Dict, step: int):
csvsavepath = os.path.join(checkpoint_path, f"{file_name}_{step}.csv")
with open(csvsavepath, "w") as file:
writer = csv.writer(file)
for key, val in data.items():
writer.writerow([key, val])
print("CSV WRITTEN " + csvsavepath)
#################################### training
tensorboard_writer = LogdirWriter(tensorboard_path)
checkpoint_task = _create_checkpoint_task(checkpoint_path)
snr_log_task = _create_snr_log_task(ARGS, checkpoint_path, tensorboard_writer, checkpoint_task)
tensorboard_task = _create_tensorboard_task(model, tensorboard_writer, ARGS.log_main_freq)
monitor_tasks = [
checkpoint_task, #
snr_log_task, #
PrintTimingsTask()
.with_name("print")
.with_condition(PeriodicIterationCondition(interval=100)),
tensorboard_task,
]
with Monitor(monitor_tasks, sess, model.global_step, print_summary=True) as monitor:
try:
restore_session(sess, checkpoint_path)
except ValueError:
pass
initial_global_step = sess.run(model.global_step)
iterations_to_go = max([ARGS.iterations - initial_global_step, 0])
if initial_global_step == 0:
# Log initial values. Bit dodgy.
tensorboard_task.run(monitor._context)
if ARGS.log_snr_freq is not None:
snr_log_task.run(monitor._context)
if ARGS.plot_freq is not None:
demo_dataset.plot_data_and_predictions(
data, model, tensorboard_writer, sess, step=0
)
print(
"Already run {} iterations. Running {} iterations".format(
initial_global_step, iterations_to_go
)
)
epoch_train_elbos = []
epoch_train_dreg_objectives = []
datapoints_since_last_epoch = 0
batching_enabled = ARGS.minibatch_size is not None
minibatch_size = ARGS.minibatch_size if batching_enabled else len(data.X_train)
for it in range(iterations_to_go):
monitor()
if isinstance(model, DregModel) and hasattr(model, "train_op"):
_, train_elbo, train_dreg_objective = sess.run(
[
model.train_op,
model.likelihood_tensor,
model.get_dreg_objective_for_encoder_params(),
]
)
epoch_train_elbos.append(train_elbo)
epoch_train_dreg_objectives.append(train_dreg_objective)
else:
model.train_func(sess)
global_step = sess.run(model.global_step)
datapoints_since_last_epoch += minibatch_size
# If batching is disabled then we use the entire dataset each iteration, so there is no point in recording
# separate epoch statistics.
if batching_enabled and datapoints_since_last_epoch >= len(data.X_train):
# We have passed over the entire dataset, so compute epoch stats.
epoch_train_elbo = np.mean(np.stack(epoch_train_elbos, axis=0))
tensorboard_writer.add_summary(
_create_scalar_summary("optimisation/epoch_train_elbo", epoch_train_elbo),
global_step,
)
epoch_train_dreg_objective = np.mean(np.stack(epoch_train_dreg_objectives, axis=0))
tensorboard_writer.add_summary(
_create_scalar_summary(
"optimisation/epoch_train_dreg_objective", epoch_train_dreg_objective
),
global_step,
)
datapoints_since_last_epoch = 0
epoch_train_elbos = []
epoch_train_dreg_objectives = []
if (
ARGS.plot_freq is not None
and (global_step - 1) % ARGS.plot_freq == 0
or it == iterations_to_go
):
demo_dataset.plot_data_and_predictions(
data, model, tensorboard_writer, sess, global_step
)
if (global_step - 1) % ARGS.log_test_freq == 0 or it == iterations_to_go:
print("Iteration: {}".format(it))
#################################### evaluation
test_elbo = model.compute_log_likelihood(data.X_test)
loglik, rmse, median_shapiro_W = metrics.compute_metrics(
model, data.X_test, data.Y_test, ARGS.num_predict_samples
)
res = {}
res["test_loglik"] = loglik
res["train_elbo"] = model.compute_log_likelihood(data.X_train)
res["test_elbo"] = test_elbo
res["test_shapiro_W_median"] = median_shapiro_W
res["test_rmse"] = rmse
res.update(ARGS.__dict__)
print(res)
_write_dict_to_csv(res, step=sess.run(model.global_step) - 1)
tensorboard_writer.add_summary(
_create_scalar_summary("optimisation/test_elbo", test_elbo), global_step
)
tensorboard_writer.add_summary(
_create_scalar_summary("optimisation/test_loglik", loglik), global_step
)
model.anchor(sess)
print(model.as_pandas_table())
####################################
loglik, rmse, median_shapiro_W = metrics.compute_metrics(
model, data.X_test, data.Y_test, ARGS.num_predict_samples
)
res = {}
res["test_loglik"] = loglik
res["train_elbo"] = model.compute_log_likelihood(data.X_train)
res["test_elbo"] = model.compute_log_likelihood(data.X_test)
res["test_shapiro_W_median"] = median_shapiro_W
res["test_rmse"] = rmse
res.update(ARGS.__dict__)
print(res)
################################### save results as csv files for tighter bounds
_write_dict_to_csv(res, step=sess.run(model.global_step))