def get_representations(dataset, postprocess_dir, dataset_name):
batch_size = 32
module_path = os.path.join(postprocess_dir, "tfhub")
reps = []
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"])
for index in range(0, len(dataset.images), batch_size):
batch = dataset.images[
index:min(index + batch_size, dataset.images.shape[0]), :]
if dataset_name == "smallnorb":
batch = np.expand_dims(batch, axis=3)
rep = _representation_function(batch)
reps.append(rep)
reps = np.vstack(reps)
# factors
factors = cartesian(
[np.array(list(range(i))) for i in dataset.factors_num_values])
return factors, reps
def random_generated(amount):
lat = np.random.normal(size=(amount,10))
#print(lat)
# Path to TFHub module of previously trained model.
module_path = os.path.join(model_dir+"/model", "tfhub")
with hub.eval_function_for_module(module_path) as f:
output2 = f(dict(latent_vectors=lat), signature="decoder", as_dict=True)["images"]
imgs = 1/(1+np.exp(-output2))
return imgs
def _evaluate_with_tensorflow(module_path, evaluation_fn, dataset, random_seed):
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.
results_dict = evaluation_fn(
dataset,
_representation_function,
random_state=np.random.RandomState(random_seed))
return results_dict
def test_convolute(self):
# Create variables to use.
random_state = np.random.RandomState(0)
data = random_state.normal(size=(5, 10))
variable1 = random_state.normal(size=(10, 6))
variable2 = random_state.normal(size=(6, 2))
# Save variables to checkpoint.
checkpoint_path = os.path.join(self.get_temp_dir(), "checkpoint.ckpt")
convolute_hub.save_numpy_arrays_to_checkpoint(
checkpoint_path, variable1=variable1, variable2=variable2)
# Save a TFHub module that we will convolute.
module_path = os.path.join(self.get_temp_dir(), "module_path")
def module_fn():
tensor = tf.placeholder(tf.float64, shape=(None, 10))
variable1 = tf.get_variable("variable1", shape=(10, 6), dtype=tf.float64)
output = tf.matmul(tensor, variable1)
hub.add_signature(
name="multiplication1",
inputs={"tensor": tensor},
outputs={"tensor": output})
spec = hub.create_module_spec(module_fn)
spec.export(module_path, checkpoint_path=checkpoint_path)
# Function used for the convolution.
def _operation2(tensor):
variable2 = tf.get_variable("variable2", shape=(6, 2), dtype=tf.float64)
return dict(tensor=tf.matmul(tensor, variable2))
# Save the convolution as a new TFHub module
module_path_new = os.path.join(self.get_temp_dir(), "module_path_new")
convolute_hub.convolute_and_save(module_path, "multiplication1",
module_path_new, _operation2,
checkpoint_path, "convoluted")
# Check the first signature.
with hub.eval_function_for_module(module_path_new) as f:
module_result = f(dict(tensor=data), signature="convoluted", as_dict=True)
real_result = data.dot(variable1).dot(variable2)
self.assertAllClose(real_result, module_result["tensor"])
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 evaluate(model_dir,
output_dir,
overwrite=False,
postprocess_fn=gin.REQUIRED,
evaluation_fn=gin.REQUIRED,
random_seed=gin.REQUIRED,
pca_components=gin.REQUIRED,
name=""):
"""Loads a trained Gaussian encoder and decoder.
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 np.array(output["mean"]), np.array(output["logvar"])
def _decoder(z):
"""Encodes images using trained model."""
# Push images through the TFHub module.
output = f(dict(latent_vectors=z),
signature="decoder",
as_dict=True)
# Convert to numpy arrays and return.
return np.array(output['images'])
# 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)
print('\n\n\n Calculating recall')
# 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_list = evaluation_fn(
dataset,
_gaussian_encoder,
transform_fn,
_decoder,
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_list = evaluation_fn(
dataset,
_gaussian_encoder,
transform_fn,
_decoder,
random_state=np.random.RandomState(random_seed))
# Save the results (and all previous results in the pipeline) on disk.
for results_dict, pca_comp in list(zip(results_dict_list, pca_components)):
results_dir = os.path.join(output_dir, "results")
results_dict["elapsed_time"] = time.time() - experiment_timer
filename = "evaluation_pca_{}comp".format(pca_comp)
results.update_result_directory(results_dir, filename, results_dict)
def train(model_dir,
trained_model_dir,
overwrite=False,
model=gin.REQUIRED,
training_steps=gin.REQUIRED,
random_seed=gin.REQUIRED,
batch_size=gin.REQUIRED,
holdout_dataset_name=gin.REQUIRED,
eval_steps=gin.REQUIRED,
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.")
# Obtain the datasets.
dataset_train, dataset_valid = named_data.get_named_ground_truth_data()
dataset_holdout = named_data.get_named_ground_truth_data(
holdout_dataset_name)
# Create a numpy random state. We will sample the random seeds for training
# and evaluation from this.
random_state = np.random.RandomState(random_seed)
# Path to TFHub module of previously trained representation.
trained_module_path = os.path.join(trained_model_dir, "tfhub")
with hub.eval_function_for_module(trained_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"])
# 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_downstream_input_fn(
dataset_train, _representation_function,
random_state.randint(2**32),
util.tf_labeled_data_set_from_ground_truth_data),
steps=training_steps)
# Save model as a TFHub module.
representation_shape, observation_shape = \
named_data.get_named_ground_truth_data()[0].representation_observation_shapes
module_export_path = os.path.join(model_dir, "tfhub")
downstream_model.export_as_tf_hub(model, representation_shape,
observation_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_downstream_input_fn(
dataset_valid,
_representation_function,
random_state.randint(2**32),
util.tf_labeled_data_set_from_ground_truth_data,
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, "evaluate", results_dict)
holdout_results_dict = tpu_estimator.evaluate(
input_fn=_make_downstream_input_fn(
dataset_holdout,
_representation_function,
random_state.randint(2**32),
util.tf_labeled_data_set_from_ground_truth_data,
num_batches=eval_steps))
holdout_results_dict["elapsed_time"] = time.time() - experiment_timer
results.update_result_directory(results_dir, "evaluate_holdout",
holdout_results_dict)
random_normal_results_dict = tpu_estimator.evaluate(
input_fn=_make_downstream_input_fn(
dataset_holdout,
_representation_function,
random_state.randint(2**32),
util.tf_random_labeled_data_set_from_ground_truth_data,
random_fn="random_normal",
num_batches=eval_steps))
random_normal_results_dict["elapsed_time"] = time.time(
) - experiment_timer
random_uniform_results_dict = tpu_estimator.evaluate(
input_fn=_make_downstream_input_fn(
dataset_holdout,
_representation_function,
random_state.randint(2**32),
util.tf_random_labeled_data_set_from_ground_truth_data,
random_fn="random_uniform",
num_batches=eval_steps))
random_uniform_results_dict["elapsed_time"] = time.time(
) - experiment_timer
results.update_result_directory(
results_dir, "evaluate_random", {
"random_normal": random_normal_results_dict,
"random_uniform": random_uniform_results_dict
})
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)
def visualize(model_dir,
output_dir,
overwrite=False,
num_animations=5,
num_frames=20,
fps=10,
num_points_irs=10000):
"""Takes trained model from model_dir and visualizes it in output_dir.
Args:
model_dir: Path to directory where the trained model is saved.
output_dir: Path to output directory.
overwrite: Boolean indicating whether to overwrite output directory.
num_animations: Integer with number of distinct animations to create.
num_frames: Integer with number of frames in each animation.
fps: Integer with frame rate for the animation.
num_points_irs: Number of points to be used for the IRS plots.
"""
# Fix the random seed for reproducibility.
random_state = np.random.RandomState(0)
# Create the output directory if necessary.
if tf.gfile.IsDirectory(output_dir):
if overwrite:
tf.gfile.DeleteRecursively(output_dir)
else:
raise ValueError(
"Directory already exists and overwrite is False.")
# 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.
# 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)
gin.bind_parameter("dataset.name",
gin_dict["dataset.name"].replace("'", ""))
# Automatically infer the activation function from gin config.
activation_str = gin_dict["reconstruction_loss.activation"]
if activation_str == "'logits'":
activation = sigmoid
elif activation_str == "'tanh'":
activation = tanh
else:
raise ValueError(
"Activation function could not be infered from gin config.")
dataset = named_data.get_named_ground_truth_data()
num_pics = 64
module_path = os.path.join(model_dir, "tfhub")
with hub.eval_function_for_module(module_path) as f:
# Save reconstructions.
real_pics = dataset.sample_observations(num_pics, random_state)
raw_pics = f(dict(images=real_pics),
signature="reconstructions",
as_dict=True)["images"]
pics = activation(raw_pics)
paired_pics = np.concatenate((real_pics, pics), axis=2)
paired_pics = [
paired_pics[i, :, :, :] for i in range(paired_pics.shape[0])
]
results_dir = os.path.join(output_dir, "reconstructions")
if not gfile.IsDirectory(results_dir):
gfile.MakeDirs(results_dir)
visualize_util.grid_save_images(
paired_pics, os.path.join(results_dir, "reconstructions.jpg"))
# Save samples.
def _decoder(latent_vectors):
return f(dict(latent_vectors=latent_vectors),
signature="decoder",
as_dict=True)["images"]
num_latent = int(gin_dict["encoder.num_latent"])
num_pics = 64
random_codes = random_state.normal(0, 1, [num_pics, num_latent])
pics = activation(_decoder(random_codes))
results_dir = os.path.join(output_dir, "sampled")
if not gfile.IsDirectory(results_dir):
gfile.MakeDirs(results_dir)
visualize_util.grid_save_images(
pics, os.path.join(results_dir, "samples.jpg"))
# Save latent traversals.
result = f(
dict(images=dataset.sample_observations(num_pics, random_state)),
signature="gaussian_encoder",
as_dict=True)
means = result["mean"]
logvars = result["logvar"]
results_dir = os.path.join(output_dir, "traversals")
if not gfile.IsDirectory(results_dir):
gfile.MakeDirs(results_dir)
for i in range(means.shape[1]):
pics = activation(
latent_traversal_1d_multi_dim(_decoder, means[i, :], None))
file_name = os.path.join(results_dir, "traversals{}.jpg".format(i))
visualize_util.grid_save_images([pics], file_name)
# Save the latent traversal animations.
results_dir = os.path.join(output_dir, "animated_traversals")
if not gfile.IsDirectory(results_dir):
gfile.MakeDirs(results_dir)
# Cycle through quantiles of a standard Gaussian.
for i, base_code in enumerate(means[:num_animations]):
images = []
for j in range(base_code.shape[0]):
code = np.repeat(np.expand_dims(base_code, 0),
num_frames,
axis=0)
code[:, j] = visualize_util.cycle_gaussian(
base_code[j], num_frames)
images.append(np.array(activation(_decoder(code))))
filename = os.path.join(results_dir,
"std_gaussian_cycle%d.gif" % i)
visualize_util.save_animation(np.array(images), filename, fps)
# Cycle through quantiles of a fitted Gaussian.
for i, base_code in enumerate(means[:num_animations]):
images = []
for j in range(base_code.shape[0]):
code = np.repeat(np.expand_dims(base_code, 0),
num_frames,
axis=0)
loc = np.mean(means[:, j])
total_variance = np.mean(np.exp(logvars[:, j])) + np.var(
means[:, j])
code[:, j] = visualize_util.cycle_gaussian(
base_code[j],
num_frames,
loc=loc,
scale=np.sqrt(total_variance))
images.append(np.array(activation(_decoder(code))))
filename = os.path.join(results_dir,
"fitted_gaussian_cycle%d.gif" % i)
visualize_util.save_animation(np.array(images), filename, fps)
# Cycle through [-2, 2] interval.
for i, base_code in enumerate(means[:num_animations]):
images = []
for j in range(base_code.shape[0]):
code = np.repeat(np.expand_dims(base_code, 0),
num_frames,
axis=0)
code[:, j] = visualize_util.cycle_interval(
base_code[j], num_frames, -2., 2.)
images.append(np.array(activation(_decoder(code))))
filename = os.path.join(results_dir,
"fixed_interval_cycle%d.gif" % i)
visualize_util.save_animation(np.array(images), filename, fps)
# Cycle linearly through +-2 std dev of a fitted Gaussian.
for i, base_code in enumerate(means[:num_animations]):
images = []
for j in range(base_code.shape[0]):
code = np.repeat(np.expand_dims(base_code, 0),
num_frames,
axis=0)
loc = np.mean(means[:, j])
total_variance = np.mean(np.exp(logvars[:, j])) + np.var(
means[:, j])
scale = np.sqrt(total_variance)
code[:, j] = visualize_util.cycle_interval(
base_code[j], num_frames, loc - 2. * scale,
loc + 2. * scale)
images.append(np.array(activation(_decoder(code))))
filename = os.path.join(results_dir,
"conf_interval_cycle%d.gif" % i)
visualize_util.save_animation(np.array(images), filename, fps)
# Cycle linearly through minmax of a fitted Gaussian.
for i, base_code in enumerate(means[:num_animations]):
images = []
for j in range(base_code.shape[0]):
code = np.repeat(np.expand_dims(base_code, 0),
num_frames,
axis=0)
code[:, j] = visualize_util.cycle_interval(
base_code[j], num_frames, np.min(means[:, j]),
np.max(means[:, j]))
images.append(np.array(activation(_decoder(code))))
filename = os.path.join(results_dir,
"minmax_interval_cycle%d.gif" % i)
visualize_util.save_animation(np.array(images), filename, fps)
# Interventional effects visualization.
factors = dataset.sample_factors(num_points_irs, random_state)
obs = dataset.sample_observations_from_factors(factors, random_state)
batch_size = 64
num_outputs = 0
latents = []
while num_outputs < obs.shape[0]:
input_batch = obs[num_outputs:min(num_outputs +
batch_size, obs.shape[0])]
output_batch = f(dict(images=input_batch),
signature="gaussian_encoder",
as_dict=True)["mean"]
latents.append(output_batch)
num_outputs += batch_size
latents = np.concatenate(latents)
results_dir = os.path.join(output_dir, "interventional_effects")
vis_all_interventional_effects(factors, latents, results_dir)
# Finally, we clear the gin config that we have set.
gin.clear_config()
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 validate(model_dir,
output_dir,
overwrite=False,
validation_fn=gin.REQUIRED,
random_seed=gin.REQUIRED,
num_labelled_samples=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.
validation_fn: Function used to validate the representation (see metrics/
for examples).
random_seed: Integer with random seed used for training.
num_labelled_samples: How many labelled samples are available.
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()
observations, labels, _ = semi_supervised_utils.sample_supervised_data(
random_seed, dataset, num_labelled_samples)
# 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.
results_dict = validation_fn(observations, np.transpose(labels),
_representation_function)
# 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, "validation", results_dict,
original_results_dir)
def get_dataset(dataname, dataroot, numsamples):
'''
inputs: x
targets: c or y (y if using dlib)
'''
if numsamples == -1:
inds = slice(0, None, 1)
else:
inds = slice(0, numsamples, 1)
if dataname == 'dSprites':
# Dataroot is the directory containing the dsprites .npz file
dataset_zip = np.load(os.path.join(
dataroot, 'dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz'),
encoding='bytes')
inputs = np.expand_dims(dataset_zip['imgs'], axis=4)[inds]
targets = dataset_zip['latents_values'][inds]
elif dataname == 'faces':
# Dataroot is the directory containing the images
filenames = np.array(
json.load(open(os.path.join(dataroot, 'img_store')))[inds])
inputs = np.array([
os.path.join(dataroot, '/'.join(filenames[i].split('/')[1:]))
for i in range(len(filenames))
])
targets = np.load(os.path.join(dataroot, 'z_store'))[inds]
elif dataname == 'faces2':
# Dataroot is the directory containing the images
filenames = np.array(
json.load(open(os.path.join(dataroot, 'img_store')))[inds])
inputs = np.array([
os.path.join(dataroot, '/'.join(filenames[i].split('/')[1:]))
for i in range(len(filenames))
])
targets = np.load(os.path.join(dataroot, 'z_store_full'))[inds]
elif dataname == 'planes':
# Dataroot is the directory containing the images
filenames = np.array(
json.load(open(os.path.join(dataroot, 'img_store_planes')))[inds])
inputs = np.array([
os.path.join(dataroot, 'planes',
'/'.join(filenames[i].split('/')[-2:]))
for i in range(len(filenames))
])
targets = np.load(os.path.join(dataroot, 'z_store_trans'))[inds]
elif dataname == 'cars':
# Dataroot is the directory containing the images
filenames = np.array(
json.load(open(os.path.join(dataroot, 'img_store_cars')))[inds])
inputs = np.array([
os.path.join(dataroot, 'cars_for_transfer',
'/'.join(filenames[i].split('/')[-2:]))
for i in range(len(filenames))
])
targets = np.load(os.path.join(dataroot, 'z_store_new_cars'))[inds]
elif dataname == 'chairs':
# Dataroot is the directory containing the images
filenames = np.array(
json.load(open(os.path.join(dataroot, 'img_store_chairs')))[inds])
inputs = np.array([
os.path.join(dataroot, 'chairs_for_transfer',
'/'.join(filenames[i].split('/')[-2:]))
for i in range(len(filenames))
])
targets = np.load(os.path.join(dataroot, 'z_store_new_chairs'))[inds]
elif dataname == 'dlib_cars3d':
## dataroot: path to the tensorflow_hub directory
dta = cars3d.Cars3D()
with hub.eval_function_for_module(dataroot) as f:
# Save reconstructions.
inputs = dta.images
targets = f(dict(images=inputs),
signature="reconstructions",
as_dict=True)["images"]
elif dataname == 'dlib_faces3d':
## dataroot: path to the tensorflow_hub directory
dta = faces3d.Faces3D()
with hub.eval_function_for_module(dataroot) as f:
# Save reconstructions.
inputs = dta.images
targets = f(dict(images=inputs),
signature="reconstructions",
as_dict=True)["images"]
else:
raise NotImplementedError
return (inputs, targets)
def visualize_supervised(supervised_model_dir,
trained_vae_model_dir,
output_dir,
overwrite=False):
"""Takes trained model from model_dir and visualizes it in output_dir.
Args:
model_dir: Path to directory where the trained model is saved.
output_dir: Path to output directory.
overwrite: Boolean indicating whether to overwrite output directory.
num_animations: Integer with number of distinct animations to create.
num_frames: Integer with number of frames in each animation.
fps: Integer with frame rate for the animation.
num_points_irs: Number of points to be used for the IRS plots.
"""
# Fix the random seed for reproducibility.
random_state = np.random.RandomState(0)
# Create the output directory if necessary.
if tf.gfile.IsDirectory(output_dir):
if overwrite:
tf.gfile.DeleteRecursively(output_dir)
else:
raise ValueError(
"Directory already exists and overwrite is False.")
# 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.
# Obtain the dataset name from the gin config of the previous step.
gin_config_file = os.path.join(supervised_model_dir, "results", "gin",
"evaluate.gin")
gin_dict = results.gin_dict(gin_config_file)
gin.bind_parameter("dataset.name",
gin_dict["dataset.name"].replace("'", ""))
# Automatically infer the activation function from gin config.
activation_str = gin_dict["reconstruction_loss.activation"]
if activation_str == "'logits'":
activation = sigmoid
elif activation_str == "'tanh'":
activation = tanh
else:
raise ValueError(
"Activation function could not be infered from gin config.")
_, dataset = named_data.get_named_ground_truth_data()
num_pics = 64
supervised_module_path = os.path.join(supervised_model_dir, "tfhub")
with hub.eval_function_for_module(supervised_module_path) as f:
trained_vae_path = os.path.join(trained_vae_model_dir, "tfhub")
with hub.eval_function_for_module(trained_vae_path) as g:
def _representation_function(x):
"""Computes representation vector for input images."""
output = g(dict(images=x),
signature="representation",
as_dict=True)
return np.array(output["default"])
# Save reconstructions.
real_pics = dataset.sample_observations(num_pics, random_state)
# real_pics, _ = dataset.sample_observations_and_labels(num_pics, random_state)
representations = _representation_function(real_pics)
print(real_pics.shape, representations.shape)
decoded_pics = f(dict(representations=representations),
signature="reconstructions",
as_dict=True)['images']
pics = activation(decoded_pics)
paired_pics = np.concatenate((real_pics, pics), axis=2)
paired_pics = [
paired_pics[i, :, :, :] for i in range(paired_pics.shape[0])
]
results_dir = os.path.join(output_dir, "reconstructions")
if not gfile.IsDirectory(results_dir):
gfile.MakeDirs(results_dir)
visualize_util.grid_save_images(
paired_pics, os.path.join(results_dir, "reconstructions.jpg"))
# Finally, we clear the gin config that we have set.
gin.clear_config()
def get_eval_res(uid):
data = np.load('../../results/eval_output/{}.npz'.format(uid))
c = data['c']
x_rand = data['x_rand']
x_enc = data['x_enc']
x_randY = data['x_randY']
x_baseline_logits = data['x_baseline']
x_baseline = sigmoid(x_baseline_logits)
model_dir = dlib_model_path[:-6]
output_directory = '../../results/eval_output/{}/'.format(uid)
config = '/hdd_c/data/disentanglement_lib/disentanglement_lib/config/unsupervised_study_v1/postprocess_configs/mean.gin'
study = reproduce.STUDIES['unsupervised_study_v1']
random_state = np.random.RandomState(0)
postprocess_config_files = sorted(study.get_postprocess_config_files())
for config in postprocess_config_files:
post_name = os.path.basename(config).replace(".gin", "")
#logging.info("Extracting representation %s...", post_name)
post_dir = os.path.join(output_directory, "postprocessed", post_name)
postprocess_bindings = [
"postprocess.random_seed = {}".format(random_state.randint(2**32)),
"postprocess.name = '{}'".format(post_name)
]
postprocess.postprocess_with_gin(model_dir, post_dir, True, [config],
postprocess_bindings)
post_processed_dir = post_dir + '/tfhub'
with hub.eval_function_for_module(post_processed_dir) as f:
# Save reconstructions.
inputs = dta.images
if inputs.ndim < 4:
inputs = np.expand_dims(inputs, 3)
inputs = inputs[:c.shape[0]]
assert inputs.shape == x_baseline.shape
inputs_c = f(dict(images=inputs),
signature="representation",
as_dict=True)["default"]
baseline_c = f(dict(images=x_baseline),
signature="representation",
as_dict=True)["default"]
x_rand_c = f(dict(images=x_rand),
signature="representation",
as_dict=True)["default"]
x_enc_c = f(dict(images=x_enc),
signature="representation",
as_dict=True)["default"]
x_randY_c = f(dict(images=x_randY),
signature="representation",
as_dict=True)["default"]
eval_bindings = [
"evaluation.random_seed = {}".format(random_state.randint(2**32)),
"evaluation.name = 'MI'"
]
gin_config_files = [
'/hdd_c/data/disentanglement_lib/disentanglement_lib/config/unsupervised_study_v1/metric_configs/mig.gin'
]
gin.parse_config_files_and_bindings(gin_config_files, eval_bindings)
def compute_mi_matrix(mus_train,
ys_train,
need_discretized_1=False,
need_discretized_2=False):
score_dict = {}
if need_discretized_1:
mus_train = utils.make_discretizer(mus_train)
if need_discretized_2:
ys_train = utils.make_discretizer(ys_train)
m = utils.discrete_mutual_info(mus_train, ys_train)
assert m.shape[0] == mus_train.shape[0]
assert m.shape[1] == ys_train.shape[0]
# m is [num_latents, num_factors]
entropy = utils.discrete_entropy(ys_train)
return m, entropy
# compute MI matrix
x_rand_mi_matrix, x_rand_entropy = compute_mi_matrix(
np.transpose(x_rand_c), np.transpose(inputs_c), True, True)
x_enc_mi_matrix, x_enc_entropy = compute_mi_matrix(np.transpose(x_enc_c),
np.transpose(inputs_c),
True, True)
baseline_mi_matrix, baseline_entropy = compute_mi_matrix(
np.transpose(baseline_c), np.transpose(inputs_c), True, True)
x_randY_mi_matrix, x_randY_entropy = compute_mi_matrix(
np.transpose(x_randY_c), np.transpose(inputs_c), True, True)
x_enc_mi_matrix_gd, x_enc_entropy_gd = compute_mi_matrix(
np.transpose(x_enc_c), np.transpose(c), True, False)
baseline_mi_matrix_gd, baseline_entropy_gd = compute_mi_matrix(
np.transpose(baseline_c), np.transpose(c), True, False)
# compute MI and MIG
x_enc_mi_average = (
np.trace(np.divide(x_enc_mi_matrix, baseline_entropy)) /
float(baseline_mi_matrix.shape[0]))
x_enc_m = np.divide(x_enc_mi_matrix_gd, baseline_entropy_gd)
sorted_x_enc_m = np.sort(x_enc_m, axis=0)[::-1]
x_enc_MIG = (np.mean(sorted_x_enc_m[0, :] - sorted_x_enc_m[1, :]))
x_rand_mi_average = (
np.trace(np.divide(x_rand_mi_matrix, baseline_entropy)) /
float(baseline_mi_matrix.shape[0]))
randY_mi_average = (
np.trace(np.divide(x_randY_mi_matrix, x_randY_entropy)) /
float(x_randY_mi_matrix.shape[0]))
baseline_mi_average = (
np.trace(np.divide(baseline_mi_matrix, baseline_entropy)) /
float(baseline_mi_matrix.shape[0]))
baseline_m = baseline_mi_matrix_gd
sorted_baseline_m = np.sort(baseline_m, axis=0)[::-1]
baseline_MIG = (np.mean(
np.divide(sorted_baseline_m[0, :] - sorted_baseline_m[1, :],
baseline_entropy_gd)))
def get_fid_with_uid(uid):
convert2image_path = '../../results/eval_output/converted_images/'
originaldata_path = '../../dataset/{}'.format(args.dataset)
data = np.load('../../results/eval_output/{}.npz'.format(uid))
path2inceptionnet = '../../inception'
c = data['c']
x_rand = data['x_rand']
x_enc = data['x_enc']
x_baseline_logits = data['x_baseline']
x_baseline = sigmoid(x_baseline_logits)
fid_list = get_fid_from_array([x_baseline, x_rand, x_enc],
convert2image_path, originaldata_path,
path2inceptionnet)
return fid_list
fid_list = get_fid_with_uid(args.id)
print('Evaluation results:')
print('Beta-TCVAE FID: {}'.format(fid_list[0]))
print('DS-VAE FID (Random Y): {}'.format(fid_list[1]))
print('DS-VAE FID: {}'.format(fid_list[2]))
print('DS-VAE MIG: {}'.format(x_enc_MIG))
print('Beta-TCVAE MIG: {}'.format(baseline_MIG))
print('DS-VAE MI: {}'.format(x_enc_mi_average))
print('DS-VAE MI (Random Z): {}'.format(x_rand_mi_average))
print('DS-VAE MI (Random Y): {}'.format(randY_mi_average))
print('Beta-TCVAE MI: {}'.format(baseline_mi_average))
if not settings.showplot:
matplotlib.use('Agg')
import matplotlib.pyplot as plt
# get ground truth factors
if args.dataset == 'cars3d':
dta = cars3d.Cars3D()
factors = get_cars_factors()
elif args.dataset == 'norb':
dta = norb.SmallNORB()
factors = get_norb_factors()
else:
raise Error
with hub.eval_function_for_module(dlib_model_path) as f:
# Save reconstructions.
inputs = dta.images
if inputs.ndim < 4:
inputs = np.expand_dims(inputs, 3)
targets = f(dict(images=inputs), signature="reconstructions",
as_dict=True)["images"]
if settings.dataname == 'faces' or settings.dataname == 'faces2' or settings.dataname == 'planes' or settings.dataname == 'cars' or settings.dataname == 'chairs' or settings.dataname == 'dlib_cars3d' or settings.dataname == 'dlib_faces3d':
input_shape = [64, 64, 3]
elif settings.dataname == 'dlib_smallnorb':
input_shape = [64, 64, 1]
else:
input_shape = list(tr_data_loader.inputs[0].shape)
print('--- Created Dataset ---')
