def test_load_frozen_ensemble(self,
subnetwork_fns,
features_to_freeze,
features_to_load,
want_logits,
bias=0,
features_placeholder=None):
freezer = _EnsembleFreezer()
filename = os.path.join(self.test_subdirectory, "frozen.pbtxt")
bias_value = bias
# First freeze ensemble.
with tf.Graph().as_default() as g, self.test_session(graph=g) as sess:
bias = tf.constant(bias, name="bias")
if features_to_freeze is not None:
features_to_freeze = tu.tensor_features(features_to_freeze)
elif features_placeholder is not None:
features_to_freeze = {
k: tf.placeholder(dtype=tf.float32, name=k, shape=shape)
for k, shape in features_placeholder.items()
}
features_to_freeze = freezer.wrapped_features(features_to_freeze)
weighted_subnetworks = []
for name, fn in subnetwork_fns:
with tf.variable_scope(
"adanet/iteration_0/ensemble_{}/".format(name)):
weighted_subnetworks.append(fn(features_to_freeze))
init = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init)
freezer.freeze_ensemble(filename=filename,
weighted_subnetworks=weighted_subnetworks,
bias=bias,
features=features_to_freeze)
# Load frozen ensemble into a new graph with potentially different
# features than those used when saving.
with tf.Graph().as_default() as g, self.test_session(graph=g) as sess:
features_to_load = tu.tensor_features(features_to_load)
features_to_load = freezer.wrapped_features(features_to_load)
frozen_ensemble, frozen_bias, _ = freezer.load_frozen_ensemble(
filename=filename, features=features_to_load)
want_ensemble = [
fn(features_to_load) for (_, fn) in subnetwork_fns
]
init = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init)
want_ensemble = sess.run(want_ensemble)
frozen_ensemble = sess.run(frozen_ensemble)
self.assertEqual([w.name for w in want_ensemble],
[w.name for w in frozen_ensemble])
self.assertAllClose([(w.logits, w.weight, w.subnetwork)
for w in want_ensemble],
[(w.logits, w.weight, w.subnetwork)
for w in frozen_ensemble])
self.assertAllEqual(bias_value, sess.run(frozen_bias))
def test_load_frozen_ensemble_colocation_bug(self, features_to_freeze,
features_to_load):
"""Test colocation bug b/74595432."""
freezer = _EnsembleFreezer()
filename = os.path.join(self.test_subdirectory, "frozen.pbtxt")
# First freeze ensemble.
with tf.Graph().as_default() as g, self.test_session(graph=g) as sess:
features_to_freeze = tu.tensor_features(features_to_freeze)
features_to_freeze = freezer.wrapped_features(features_to_freeze)
for k, feature in features_to_freeze.items():
with tf.colocate_with(feature):
if isinstance(feature, tf.SparseTensor):
feature = tf.SparseTensor(
tf.identity(feature.indices,
name="colocated_indices"),
tf.identity(feature.values,
name="colocated_values"),
tf.identity(feature.dense_shape,
name="colocated_dense_shape"))
else:
feature = tf.identity(feature, name="colocated")
features_to_freeze[k] = feature
weighted_subnetworks = [_dnn_subnetwork_fn()(features_to_freeze)]
init = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init)
freezer.freeze_ensemble(filename=filename,
weighted_subnetworks=weighted_subnetworks,
bias=tf.constant(0, name="bias"),
features=features_to_freeze)
# Verify that repeatedly freezing and reloading frozen ensembles works.
features_to_load_copy = features_to_load.copy()
for _ in range(5):
with tf.Graph().as_default() as g, self.test_session(
graph=g) as sess:
features_to_load = tu.tensor_features(features_to_load_copy)
features_to_load = freezer.wrapped_features(features_to_load)
frozen_ensemble, frozen_bias, _ = freezer.load_frozen_ensemble(
filename=filename, features=features_to_load)
init = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init)
freezer.freeze_ensemble(filename=filename,
weighted_subnetworks=frozen_ensemble,
bias=frozen_bias,
features=features_to_freeze)