def postprocess(model_dir,
output_dir,
overwrite=False,
postprocess_fn=gin.REQUIRED,
random_seed=gin.REQUIRED,
name=""):
"""Loads a trained Gaussian encoder and extracts representation.
Args:
model_dir: String with path to directory where the model is saved.
output_dir: String with the path where the representation should be saved.
overwrite: Boolean indicating whether to overwrite output directory.
postprocess_fn: Function used to extract the representation (see methods.py
for examples).
random_seed: Integer with random seed used for postprocessing (may be
unused).
name: Optional string with name of the representation (can be used to name
representations).
"""
# We do not use the variable 'name'. Instead, it can be used to name
# representations as it will be part of the saved gin config.
del name
# Delete the output directory if it already exists.
if tf.gfile.IsDirectory(output_dir):
if overwrite:
tf.gfile.DeleteRecursively(output_dir)
else:
raise ValueError("Directory already exists and overwrite is False.")
# Set up timer to keep track of elapsed time in results.
experiment_timer = time.time()
# Automatically set the proper data set if necessary. We replace the active
# gin config as this will lead to a valid gin config file where the data set
# is present.
if gin.query_parameter("dataset.name") == "auto":
# Obtain the dataset name from the gin config of the previous step.
gin_config_file = os.path.join(model_dir, "results", "gin", "train.gin")
gin_dict = results.gin_dict(gin_config_file)
with gin.unlock_config():
gin.bind_parameter("dataset.name", gin_dict["dataset.name"].replace(
"'", ""))
dataset = named_data.get_named_ground_truth_data()
# Path to TFHub module of previously trained model.
module_path = os.path.join(model_dir, "tfhub")
with hub.eval_function_for_module(module_path) as f:
def _gaussian_encoder(x):
"""Encodes images using trained model."""
# Push images through the TFHub module.
output = f(dict(images=x), signature="gaussian_encoder", as_dict=True)
# Convert to numpy arrays and return.
return {key: np.array(values) for key, values in output.items()}
# Run the postprocessing function which returns a transformation function
# that can be used to create the representation from the mean and log
# variance of the Gaussian distribution given by the encoder. Also returns
# path to a checkpoint if the transformation requires variables.
transform_fn, transform_checkpoint_path = postprocess_fn(
dataset, _gaussian_encoder, np.random.RandomState(random_seed),
output_dir)
# Takes the "gaussian_encoder" signature, extracts the representation and
# then saves under the signature "representation".
tfhub_module_dir = os.path.join(output_dir, "tfhub")
convolute_hub.convolute_and_save(
module_path, "gaussian_encoder", tfhub_module_dir, transform_fn,
transform_checkpoint_path, "representation")
# We first copy over all the prior results and configs.
original_results_dir = os.path.join(model_dir, "results")
results_dir = os.path.join(output_dir, "results")
results_dict = dict(elapsed_time=time.time() - experiment_timer)
results.update_result_directory(results_dir, "postprocess", results_dict,
original_results_dir)
def train(model_dir,
overwrite=False,
model=gin.REQUIRED,
training_steps=gin.REQUIRED,
random_seed=gin.REQUIRED,
batch_size=gin.REQUIRED,
eval_steps=1000,
name="",
model_num=None):
"""Trains the estimator and exports the snapshot and the gin config.
The use of this function requires the gin binding 'dataset.name' to be
specified as that determines the data set used for training.
Args:
model_dir: String with path to directory where model output should be saved.
overwrite: Boolean indicating whether to overwrite output directory.
model: GaussianEncoderModel that should be trained and exported.
training_steps: Integer with number of training steps.
random_seed: Integer with random seed used for training.
batch_size: Integer with the batch size.
eval_steps: Optional integer with number of steps used for evaluation.
name: Optional string with name of the model (can be used to name models).
model_num: Optional integer with model number (can be used to identify
models).
"""
# We do not use the variables 'name' and 'model_num'. Instead, they can be
# used to name results as they will be part of the saved gin config.
del name, model_num
# Delete the output directory if it already exists.
if tf.gfile.IsDirectory(model_dir):
if overwrite:
tf.gfile.DeleteRecursively(model_dir)
else:
raise ValueError("Directory already exists and overwrite is False.")
# Create a numpy random state. We will sample the random seeds for training
# and evaluation from this.
random_state = np.random.RandomState(random_seed)
# Obtain the dataset.
dataset = named_data.get_named_ground_truth_data()
# We create a TPUEstimator based on the provided model. This is primarily so
# that we could switch to TPU training in the future. For now, we train
# locally on GPUs.
run_config = tf.contrib.tpu.RunConfig(
tf_random_seed=random_seed,
keep_checkpoint_max=1,
tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=500))
tpu_estimator = tf.contrib.tpu.TPUEstimator(
use_tpu=False,
model_fn=model.model_fn,
model_dir=os.path.join(model_dir, "tf_checkpoint"),
train_batch_size=batch_size,
eval_batch_size=batch_size,
config=run_config)
# Set up time to keep track of elapsed time in results.
experiment_timer = time.time()
# Do the actual training.
tpu_estimator.train(
input_fn=_make_input_fn(dataset, random_state.randint(2**32)),
steps=training_steps)
# Save model as a TFHub module.
output_shape = named_data.get_named_ground_truth_data().observation_shape
module_export_path = os.path.join(model_dir, "tfhub")
gaussian_encoder_model.export_as_tf_hub(model, output_shape,
tpu_estimator.latest_checkpoint(),
module_export_path)
# Save the results. The result dir will contain all the results and config
# files that we copied along, as we progress in the pipeline. The idea is that
# these files will be available for analysis at the end.
results_dict = tpu_estimator.evaluate(
input_fn=_make_input_fn(
dataset, random_state.randint(2**32), num_batches=eval_steps))
results_dir = os.path.join(model_dir, "results")
results_dict["elapsed_time"] = time.time() - experiment_timer
results.update_result_directory(results_dir, "train", results_dict)
def evaluate(model_dir,
output_dir,
overwrite=False,
evaluation_fn=gin.REQUIRED,
random_seed=gin.REQUIRED,
name=""):
"""Loads a representation TFHub module and computes disentanglement metrics.
Args:
model_dir: String with path to directory where the representation function
is saved.
output_dir: String with the path where the results should be saved.
overwrite: Boolean indicating whether to overwrite output directory.
evaluation_fn: Function used to evaluate the representation (see metrics/
for examples).
random_seed: Integer with random seed used for training.
name: Optional string with name of the metric (can be used to name metrics).
"""
# We do not use the variable 'name'. Instead, it can be used to name scores
# as it will be part of the saved gin config.
del name
# Delete the output directory if it already exists.
if tf.gfile.IsDirectory(output_dir):
if overwrite:
tf.gfile.DeleteRecursively(output_dir)
else:
raise ValueError("Directory already exists and overwrite is False.")
# Set up time to keep track of elapsed time in results.
experiment_timer = time.time()
# Automatically set the proper data set if necessary. We replace the active
# gin config as this will lead to a valid gin config file where the data set
# is present.
if gin.query_parameter("dataset.name") == "auto":
# Obtain the dataset name from the gin config of the previous step.
gin_config_file = os.path.join(model_dir, "results", "gin",
"postprocess.gin")
gin_dict = results.gin_dict(gin_config_file)
with gin.unlock_config():
gin.bind_parameter("dataset.name", gin_dict["dataset.name"].replace(
"'", ""))
dataset = named_data.get_named_ground_truth_data()
# Path to TFHub module of previously trained representation.
module_path = os.path.join(model_dir, "tfhub")
with hub.eval_function_for_module(module_path) as f:
def _representation_function(x):
"""Computes representation vector for input images."""
output = f(dict(images=x), signature="representation", as_dict=True)
return np.array(output["default"])
# Computes scores of the representation based on the evaluation_fn.
if _has_kwarg_or_kwargs(evaluation_fn, "artifact_dir"):
artifact_dir = os.path.join(model_dir, "artifacts")
results_dict = evaluation_fn(
dataset,
_representation_function,
random_state=np.random.RandomState(random_seed),
artifact_dir=artifact_dir)
else:
# Legacy code path to allow for old evaluation metrics.
warnings.warn(
"Evaluation function does not appear to accept an"
" `artifact_dir` argument. This may not be compatible with "
"future versions.", DeprecationWarning)
results_dict = evaluation_fn(
dataset,
_representation_function,
random_state=np.random.RandomState(random_seed))
# Save the results (and all previous results in the pipeline) on disk.
original_results_dir = os.path.join(model_dir, "results")
results_dir = os.path.join(output_dir, "results")
results_dict["elapsed_time"] = time.time() - experiment_timer
results.update_result_directory(results_dir, "evaluation", results_dict,
original_results_dir)
def reason(
input_dir,
output_dir,
overwrite=False,
model=gin.REQUIRED,
num_iterations=gin.REQUIRED,
training_steps_per_iteration=gin.REQUIRED,
eval_steps_per_iteration=gin.REQUIRED,
random_seed=gin.REQUIRED,
batch_size=gin.REQUIRED,
name="",
):
"""Trains the estimator and exports the snapshot and the gin config.
The use of this function requires the gin binding 'dataset.name' to be
specified if a model is trained from scratch as that determines the data set
used for training.
Args:
input_dir: String with path to directory where the representation function
is saved.
output_dir: String with the path where the results should be saved.
overwrite: Boolean indicating whether to overwrite output directory.
model: GaussianEncoderModel that should be trained and exported.
num_iterations: Integer with number of training steps.
training_steps_per_iteration: Integer with number of training steps per
iteration.
eval_steps_per_iteration: Integer with number of validationand test steps
per iteration.
random_seed: Integer with random seed used for training.
batch_size: Integer with the batch size.
name: Optional string with name of the model (can be used to name models).
"""
# We do not use the variable 'name'. Instead, it can be used to name results
# as it will be part of the saved gin config.
del name
# Delete the output directory if it already exists.
if tf.gfile.IsDirectory(output_dir):
if overwrite:
tf.gfile.DeleteRecursively(output_dir)
else:
raise ValueError(
"Directory already exists and overwrite is False.")
# Create a numpy random state. We will sample the random seeds for training
# and evaluation from this.
random_state = np.random.RandomState(random_seed)
# Automatically set the proper data set if necessary. We replace the active
# gin config as this will lead to a valid gin config file where the data set
# is present.
if gin.query_parameter("dataset.name") == "auto":
if input_dir is None:
raise ValueError(
"Cannot automatically infer data set for methods with"
" no prior model directory.")
# Obtain the dataset name from the gin config of the previous step.
gin_config_file = os.path.join(input_dir, "results", "gin",
"postprocess.gin")
gin_dict = results.gin_dict(gin_config_file)
with gin.unlock_config():
gin.bind_parameter("dataset.name",
gin_dict["dataset.name"].replace("'", ""))
dataset = pgm_data.get_pgm_dataset()
# Set the path to the TFHub embedding if we are training based on a
# pre-trained embedding..
if input_dir is not None:
tfhub_dir = os.path.join(input_dir, "tfhub")
with gin.unlock_config():
gin.bind_parameter("HubEmbedding.hub_path", tfhub_dir)
# We create a TPUEstimator based on the provided model. This is primarily so
# that we could switch to TPU training in the future. For now, we train
# locally on GPUs.
run_config = tf.contrib.tpu.RunConfig(
tf_random_seed=random_seed,
keep_checkpoint_max=1,
tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=500))
tpu_estimator = tf.contrib.tpu.TPUEstimator(use_tpu=False,
model_fn=model.model_fn,
model_dir=os.path.join(
output_dir,
"tf_checkpoint"),
train_batch_size=batch_size,
eval_batch_size=batch_size,
config=run_config)
# Set up time to keep track of elapsed time in results.
experiment_timer = time.time()
# Create a dictionary to keep track of all relevant information.
results_dict_of_dicts = {}
validation_scores = []
all_dicts = []
for i in range(num_iterations):
steps_so_far = i * training_steps_per_iteration
tf.logging.info("Training to %d steps.", steps_so_far)
# Train the model for the specified steps.
tpu_estimator.train(input_fn=dataset.make_input_fn(
random_state.randint(2**32)),
steps=training_steps_per_iteration)
# Compute validation scores used for model selection.
validation_results = tpu_estimator.evaluate(
input_fn=dataset.make_input_fn(
random_state.randint(2**32),
num_batches=eval_steps_per_iteration))
validation_scores.append(validation_results["accuracy"])
tf.logging.info("Validation results %s", validation_results)
# Compute test scores for final results.
test_results = tpu_estimator.evaluate(input_fn=dataset.make_input_fn(
random_state.randint(2**32), num_batches=eval_steps_per_iteration),
name="test")
dict_at_iteration = results.namespaced_dict(val=validation_results,
test=test_results)
results_dict_of_dicts["step{}".format(
steps_so_far)] = dict_at_iteration
all_dicts.append(dict_at_iteration)
# Select the best number of steps based on the validation scores and add it as
# as a special key to the dictionary.
best_index = np.argmax(validation_scores)
results_dict_of_dicts["best"] = all_dicts[best_index]
# Save the results. The result dir will contain all the results and config
# files that we copied along, as we progress in the pipeline. The idea is that
# these files will be available for analysis at the end.
if input_dir is not None:
original_results_dir = os.path.join(input_dir, "results")
else:
original_results_dir = None
results_dict = results.namespaced_dict(**results_dict_of_dicts)
results_dir = os.path.join(output_dir, "results")
results_dict["elapsed_time"] = time.time() - experiment_timer
results.update_result_directory(results_dir, "abstract_reasoning",
results_dict, original_results_dir)
def evaluate(model_dirs,
output_dir,
evaluation_fn=gin.REQUIRED,
random_seed=gin.REQUIRED,
name=""):
"""Loads a trained estimator and evaluates it according to beta-VAE metric."""
# The name will be part of the gin config and can be used to tag results.
del name
# Set up time to keep track of elapsed time in results.
experiment_timer = time.time()
# Automatically set the proper dataset if necessary. We replace the active
# gin config as this will lead to a valid gin config file where the dataset
# is present.
if gin.query_parameter("dataset.name") == "auto":
# Obtain the dataset name from the gin config of the previous step.
gin_config_file = os.path.join(model_dirs[0], "results", "gin",
"train.gin")
gin_dict = results.gin_dict(gin_config_file)
with gin.unlock_config():
print(gin_dict["dataset.name"])
gin.bind_parameter("dataset.name",
gin_dict["dataset.name"].replace("'", ""))
output_dir = os.path.join(output_dir)
if tf.io.gfile.isdir(output_dir):
tf.io.gfile.rmtree(output_dir)
dataset = named_data.get_named_ground_truth_data()
with contextlib.ExitStack() as stack:
representation_functions = []
eval_functions = [
stack.enter_context(
hub.eval_function_for_module(os.path.join(model_dir, "tfhub")))
for model_dir in model_dirs
]
for f in eval_functions:
def _representation_function(x, f=f):
def compute_gaussian_kl(z_mean, z_logvar):
return np.mean(
0.5 *
(np.square(z_mean) + np.exp(z_logvar) - z_logvar - 1),
axis=0)
encoding = f(dict(images=x),
signature="gaussian_encoder",
as_dict=True)
return np.array(encoding["mean"]), compute_gaussian_kl(
np.array(encoding["mean"]), np.array(encoding["logvar"]))
representation_functions.append(_representation_function)
results_dict = evaluation_fn(
dataset,
representation_functions,
random_state=np.random.RandomState(random_seed))
original_results_dir = os.path.join(model_dirs[0], "results")
results_dir = os.path.join(output_dir, "results")
results_dict["elapsed_time"] = time.time() - experiment_timer
results.update_result_directory(results_dir, "evaluation", results_dict,
original_results_dir)