def _save_restore_keras_model(model):
_, keras_file = tempfile.mkstemp('.h5')
keras.models.save_model(model, keras_file)
with prune.prune_scope():
loaded_model = keras.models.load_model(keras_file)
return loaded_model
def _save_restore_saved_model(model):
tmpdir = tempfile.mkdtemp()
saved_model_experimental.export_saved_model(model, tmpdir)
with prune.prune_scope():
loaded_model = saved_model_experimental.load_from_saved_model(tmpdir)
loaded_model.compile(loss='categorical_crossentropy',
optimizer='sgd',
metrics=['accuracy'])
return loaded_model
def testPruneScope_NeededForKerasModel(self):
model = keras_test_utils.build_simple_dense_model()
pruned_model = prune.prune_low_magnitude(model)
_, keras_model = tempfile.mkstemp('.h5')
pruned_model.save(keras_model)
with self.assertRaises(ValueError):
tf.keras.models.load_model(keras_model)
# works with `prune_scope`
with prune.prune_scope():
tf.keras.models.load_model(keras_model)
def testPruneFunctionalModelPreservesBuiltState(self):
i1 = keras.Input(shape=(10, ))
i2 = keras.Input(shape=(10, ))
x1 = layers.Dense(10)(i1)
x2 = layers.Dense(10)(i2)
outputs = layers.Add()([x1, x2])
model = keras.Model(inputs=[i1, i2], outputs=outputs)
self.assertEqual(model.built, True)
pruned_model = prune.prune_low_magnitude(model, **self.params)
self.assertEqual(model.built, True)
# Test built state preserves across serialization
with prune.prune_scope():
loaded_model = keras.models.model_from_config(
json.loads(pruned_model.to_json()))
self.assertEqual(loaded_model.built, True)
def testPruneScope_NeededForTF1SavedModel(self):
# TODO(b/185726968): replace with shared v1 test_util.
is_v1_apis = hasattr(tf, 'assign')
if not is_v1_apis:
return
model = keras_test_utils.build_simple_dense_model()
pruned_model = prune.prune_low_magnitude(model)
saved_model_dir = tempfile.mkdtemp()
tf.keras.experimental.export_saved_model(pruned_model, saved_model_dir)
with self.assertRaises(ValueError):
tf.keras.experimental.load_from_saved_model(saved_model_dir)
# works with `prune_scope`
with prune.prune_scope():
tf.keras.experimental.load_from_saved_model(saved_model_dir)
def train_and_save(model, x_train, y_train, x_test, y_test):
model.compile(
loss="categorical_crossentropy",
optimizer="adam",
metrics=["accuracy"])
# Print the model summary.
model.summary()
# Add a pruning step callback to peg the pruning step to the optimizer's
# step. Also add a callback to add pruning summaries to tensorboard
callbacks = [
pruning_callbacks.UpdatePruningStep(),
#pruning_callbacks.PruningSummaries(log_dir=tempfile.mkdtemp())
pruning_callbacks.PruningSummaries(log_dir="/tmp/mnist_prune")
]
model.fit(
x_train,
y_train,
batch_size=batch_size,
epochs=epochs,
verbose=1,
callbacks=callbacks,
validation_data=(x_test, y_test))
score = model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", score[0])
print("Test accuracy:", score[1])
print_model_sparsity(model)
# Export and import the model. Check that accuracy persists.
_, keras_file = tempfile.mkstemp(".h5")
print("Saving model to: ", keras_file)
save_model(model, keras_file)
print("Reloading model")
with prune.prune_scope():
loaded_model = load_qmodel(keras_file)
score = loaded_model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", score[0])
print("Test accuracy:", score[1])
def train_and_save(models, x_train, y_train, x_test, y_test):
for model in models:
model.compile(
loss=tf.keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy'])
# Print the model summary.
model.summary()
# Add a pruning step callback to peg the pruning step to the optimizer's
# step. Also add a callback to add pruning summaries to tensorboard
callbacks = [
pruning_callbacks.UpdatePruningStep(),
pruning_callbacks.PruningSummaries(log_dir=FLAGS.output_dir)
]
model.fit(
x_train,
y_train,
batch_size=batch_size,
epochs=epochs,
verbose=1,
callbacks=callbacks,
validation_data=(x_test, y_test))
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])
# Export and import the model. Check that accuracy persists.
_, keras_file = tempfile.mkstemp('.h5')
print('Saving model to: ', keras_file)
keras.models.save_model(model, keras_file)
with prune.prune_scope():
loaded_model = keras.models.load_model(keras_file)
score = loaded_model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])
def _save_restore_tf_model(model):
tmpdir = tempfile.mkdtemp()
tf.keras.models.save_model(model, tmpdir, save_format='tf')
with prune.prune_scope():
loaded_model = tf.keras.models.load_model(tmpdir)
return loaded_model
