def testPruneStopAndRestartOnModel(self, save_restore_fn): params = { 'pruning_schedule': pruning_schedule.PolynomialDecay(0.2, 0.6, 0, 4, 3, 1), 'block_size': (1, 1), 'block_pooling_type': 'AVG' } model = prune.prune_low_magnitude( keras_test_utils.build_simple_dense_model(), **params) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) # Model hasn't been trained yet. Sparsity 0.0 test_utils.assert_model_sparsity(self, 0.0, model) model.fit(np.random.rand(20, 10), np_utils.to_categorical(np.random.randint(5, size=(20, 1)), 5), batch_size=20, callbacks=[pruning_callbacks.UpdatePruningStep()]) # Training has run only 1 step. Sparsity 0.2 (initial_sparsity) test_utils.assert_model_sparsity(self, 0.2, model) model = save_restore_fn(model) model.fit(np.random.rand(20, 10), np_utils.to_categorical(np.random.randint(5, size=(20, 1)), 5), batch_size=20, epochs=3, callbacks=[pruning_callbacks.UpdatePruningStep()]) # Training has run all 4 steps. Sparsity 0.6 (final_sparsity) test_utils.assert_model_sparsity(self, 0.6, model) self._check_strip_pruning_matches_original(model, 0.6)
def testPrunesModel_CustomTrainingLoop_ReachesTargetSparsity(self): pruned_model = prune.prune_low_magnitude( keras_test_utils.build_simple_dense_model()) batch_size = 20 x_train = np.random.rand(20, 10) y_train = keras.utils.to_categorical( np.random.randint(5, size=(batch_size, 1)), 5) loss = keras.losses.categorical_crossentropy optimizer = keras.optimizers.SGD() unused_arg = -1 step_callback = pruning_callbacks.UpdatePruningStep() step_callback.set_model(pruned_model) pruned_model.optimizer = optimizer step_callback.on_train_begin() # 2 epochs for _ in range(2): step_callback.on_train_batch_begin(batch=unused_arg) inp = np.reshape(x_train, [batch_size, 10]) # original shape: from [10]. with tf.GradientTape() as tape: logits = pruned_model(inp, training=True) loss_value = loss(y_train, logits) grads = tape.gradient(loss_value, pruned_model.trainable_variables) optimizer.apply_gradients( zip(grads, pruned_model.trainable_variables)) step_callback.on_epoch_end(batch=unused_arg) test_utils.assert_model_sparsity(self, 0.5, pruned_model)
def train(model, x_train, y_train, x_test, y_test): 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='/tmp/logs') ] 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]) model = prune.strip_pruning(model) return model
def testPrunesMnist_ReachesTargetSparsity(self, model_type): model = test_utils.build_mnist_model(model_type, self.params) if model_type == 'layer_list': model = keras.Sequential( prune.prune_low_magnitude(model, **self.params)) elif model_type in ['sequential', 'functional']: model = prune.prune_low_magnitude(model, **self.params) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model, rtol=1e-4, atol=1e-4) model.fit(np.random.rand(32, 28, 28, 1), keras.utils.to_categorical( np.random.randint(10, size=(32, 1)), 10), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity(self, 0.5, model, rtol=1e-4, atol=1e-4) self._check_strip_pruning_matches_original(model, 0.5)
def testPruneCheckpoints_CheckpointsNotSparse(self): is_model_sparsity_not_list = [] # Run multiple times since problem doesn't always happen. for _ in range(3): model = keras_test_utils.build_simple_dense_model() pruned_model = prune.prune_low_magnitude(model, **self.params) checkpoint_dir = tempfile.mkdtemp() checkpoint_path = checkpoint_dir + '/weights.{epoch:02d}.tf' callbacks = [ pruning_callbacks.UpdatePruningStep(), tf.keras.callbacks.ModelCheckpoint( filepath=checkpoint_path, save_weights_only=True, save_freq=1) ] # Train one step. Sparsity reaches final sparsity. self._train_model(pruned_model, epochs=1, callbacks=callbacks) test_utils.assert_model_sparsity(self, 0.5, pruned_model) latest_checkpoint = tf.train.latest_checkpoint(checkpoint_dir) same_architecture_model = keras_test_utils.build_simple_dense_model() pruned_model = prune.prune_low_magnitude(same_architecture_model, **self.params) # Sanity check. test_utils.assert_model_sparsity(self, 0, pruned_model) pruned_model.load_weights(latest_checkpoint) is_model_sparsity_not_list.append( test_utils.is_model_sparsity_not(0.5, pruned_model)) self.assertTrue(any(is_model_sparsity_not_list))
def testPrunesSimpleDenseModel(self, distribution): with distribution.scope(): model = prune.prune_low_magnitude( keras_test_utils.build_simple_dense_model(), **self.params) model.compile( loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) # Model hasn't been trained yet. Sparsity 0.0 test_utils.assert_model_sparsity(self, 0.0, model) # Simple unpruned model. No sparsity. model.fit( np.random.rand(20, 10), keras.utils.np_utils.to_categorical( np.random.randint(5, size=(20, 1)), 5), epochs=2, callbacks=[pruning_callbacks.UpdatePruningStep()], batch_size=20) model.predict(np.random.rand(20, 10)) test_utils.assert_model_sparsity(self, 0.5, 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) test_utils.assert_model_sparsity(self, 0.5, loaded_model)
def testRNNLayersWithRNNCellParams(self): model = keras.Sequential() model.add( prune.prune_low_magnitude(keras.layers.RNN([ layers.LSTMCell(10), layers.GRUCell(10), layers.PeepholeLSTMCell(10), layers.SimpleRNNCell(10) ]), input_shape=(3, 4), **self.params)) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model) model.fit(np.random.randn( *self._batch(model.input.get_shape().as_list(), 32)), np.random.randn( *self._batch(model.output.get_shape().as_list(), 32)), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity(self, 0.5, model) self._check_strip_pruning_matches_original(model, 0.5)
def testPrunesPreviouslyUnprunedModel(self): model = keras_test_utils.build_simple_dense_model() model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) # Simple unpruned model. No sparsity. model.fit(np.random.rand(20, 10), np_utils.to_categorical(np.random.randint(5, size=(20, 1)), 5), epochs=2, batch_size=20) test_utils.assert_model_sparsity(self, 0.0, model) # Apply pruning to model. model = prune.prune_low_magnitude(model, **self.params) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) # Since newly compiled, iterations starts from 0. model.fit(np.random.rand(20, 10), np_utils.to_categorical(np.random.randint(5, size=(20, 1)), 5), batch_size=20, callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity(self, 0.5, model) self._check_strip_pruning_matches_original(model, 0.5)
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 testMbyNSparsityPruning_SupportedLayers(self, layer_type, layer_arg, input_shape, m_by_n=(2, 4), sparsity_ratio=0.50): """Check that we prune supported layers with m by n sparsity.""" self.params.update({'sparsity_m_by_n': m_by_n}) model = keras.Sequential() model.add( prune.prune_low_magnitude(layer_type(*layer_arg), input_shape=input_shape, **self.params)) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model) model.fit(np.random.randn( *self._batch(model.input.get_shape().as_list(), 32)), np.random.randn( *self._batch(model.output.get_shape().as_list(), 32)), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity_m_by_n(self, model, m_by_n) self._check_strip_pruning_matches_original(model, sparsity_ratio)
def _train_model(model, epochs=1, x_train=None, y_train=None, callbacks=None): if x_train is None: x_train = np.random.rand(20, 10), if y_train is None: y_train = keras.utils.to_categorical( np.random.randint(5, size=(20, 1)), 5) if model.optimizer is None: model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) if callbacks is None: callbacks = [] if PruneIntegrationTest._is_pruned(model): callbacks = [pruning_callbacks.UpdatePruningStep()] model.fit(x_train, y_train, epochs=epochs, batch_size=20, callbacks=callbacks)
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 if not os.path.exists(FLAGS.output_dir): os.makedirs(FLAGS.output_dir) # print(FLAGS.output_dir) 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. saved_model_dir = '/tmp/saved_model' if not os.path.exists(saved_model_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]) os.makedirs(saved_model_dir) print('Saving model to: ', saved_model_dir) tf.keras.models.save_model(model, saved_model_dir, save_format='tf') print('Loading model from: ', saved_model_dir) loaded_model = tf.keras.models.load_model(saved_model_dir) score = loaded_model.evaluate(x_test, y_test, verbose=0) print('Test loss:', score[0]) print('Test accuracy:', score[1])
def testDeepLayerUpdatePruningSteps(self): pruned_model, x_train, y_train = self._pruned_model_setup(has_deep_layer=True) pruned_model.fit( x_train, y_train, batch_size=self._BATCH_SIZE, epochs=3, callbacks=[ pruning_callbacks.UpdatePruningStep() ]) self.assertEqual(2, pruned_model.layers[-1]._layers[0].pruning_step)
def fit(self, X_train, y_train): with self.graph.as_default(), self.session.as_default(): callbacks = None if self.prune_params: callbacks = [ pruning_callbacks.UpdatePruningStep(), ] self.model.fit(X_train, y_train, epochs=self.train_epochs, batch_size=self.batch_size, callbacks=callbacks)
def fitModel(model, train_data, val_data, test_data, stepsPerEpoch, evalStepsPerEpoch): """Runs Keras fit and saves model. Arguments: STRATEGY: Mirrored strategy models: list of models to train train_data: training data val_data: validation data Returns: None """ if not os.path.exists(FLAGS.outdir + '/%s/' % model.name): os.system('mkdir ' + FLAGS.outdir + '/%s/' % model.name) callbacks = getCallbacks(FLAGS.outdir + '/%s/' % model.name) if FLAGS.prune == True: callbacks.append(pruning_callbacks.UpdatePruningStep()) start = time.time() LOSS = tf.keras.losses.CategoricalCrossentropy() OPTIMIZER = Adam(learning_rate=FLAGS.lr, beta_1=FLAGS.beta_1, beta_2=FLAGS.beta_2, epsilon=FLAGS.epsilon, amsgrad=True) model.compile(loss=LOSS, optimizer=OPTIMIZER, metrics=["accuracy"]) model.summary() history = model.fit(train_data, epochs=FLAGS.epochs, validation_data=val_data, callbacks=callbacks, verbose=1) model.load_weights(FLAGS.outdir + '/%s/weights_best.h5' % model.name) history_dict = history.history pd.DataFrame.from_dict(history.history).to_csv( FLAGS.outdir + '/%s/history_dict.csv' % model.name, index=False) test_score = model.evaluate(test_data) print("Done training model {}".format(model.name)) print('\n Test loss:', test_score[0]) print('\n Test accuracy:', test_score[1]) np.savez(FLAGS.outdir + '/%s/scores' % model.name, test_score) if FLAGS.prune == True: model_stripped = strip_pruning(model) model_stripped.save(FLAGS.outdir + '/%s/%s.h5' % (model.name, model.name)) else: model.save(FLAGS.outdir + '/%s/%s.h5' % (model.name, model.name)) end = time.time() print('\n It took {} minutes to train!\n'.format((end - start) / 60.))
def _prune_model(original_model): """Apply the pruning wrapper, compile and train the model.""" prune_epoch = 1 pruning_params = { 'pruning_schedule': pruning_schedule.ConstantSparsity(0.50, begin_step=0, frequency=10) } pruning_model = prune.prune_low_magnitude(original_model, **pruning_params) callbacks = [pruning_callbacks.UpdatePruningStep()] pruning_model = _train_model(pruning_model, callbacks, prune_epoch) pruning_model_stripped = prune.strip_pruning(pruning_model) return pruning_model, pruning_model_stripped
def testUpdatesPruningStep(self): model = prune.prune_low_magnitude( keras_test_utils.build_simple_dense_model()) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) model.fit(np.random.rand(20, 10), keras.utils.np_utils.to_categorical( np.random.randint(5, size=(20, 1)), 5), batch_size=20, epochs=3, callbacks=[pruning_callbacks.UpdatePruningStep()]) self.assertEqual(2, K.get_value(model.layers[0].pruning_step)) self.assertEqual(2, K.get_value(model.layers[1].pruning_step))
def testRNNLayersSingleCell_ReachesTargetSparsity(self, layer_type): model = keras.Sequential() model.add( prune.prune_low_magnitude( layer_type(10), input_shape=(3, 4), **self.params)) model.compile( loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model) model.fit( np.random.randn(*self._batch(model.input.get_shape().as_list(), 32)), np.random.randn(*self._batch(model.output.get_shape().as_list(), 32)), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity(self, 0.5, model) self._check_strip_pruning_matches_original(model, 0.5)
def testUpdatePruningStepsAndLogsSummaries(self): log_dir = tempfile.mkdtemp() pruned_model, x_train, y_train = self._pruned_model_setup() pruned_model.fit( x_train, y_train, batch_size=self._BATCH_SIZE, epochs=3, callbacks=[ pruning_callbacks.UpdatePruningStep(), pruning_callbacks.PruningSummaries(log_dir=log_dir) ]) self.assertEqual( 2, tf.keras.backend.get_value(pruned_model.layers[0].pruning_step)) self.assertEqual( 2, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step)) self._assertLogsExist(log_dir)
def testSparsityPruningMbyN_SupportedSubclassLayers(self): """Check subclass layer that is supported for m by n sparsity.""" m_by_n = (2, 4) self.params.update({'sparsity_m_by_n': m_by_n}) class SubclassLayer(tf.keras.layers.Layer): def __init__(self): super(SubclassLayer, self).__init__() self.conv1 = tf.keras.layers.Conv2D(2, 3, activation='relu', padding='same', input_shape=[7, 7, 3]) self.conv2 = tf.keras.layers.DepthwiseConv2D(3) self.flatten = keras.layers.Flatten() self.dense = layers.Dense(10, activation='sigmoid') def call(self, inputs): x = self.conv1(inputs) x = self.conv2(x) x = self.flatten(x) x = self.dense(x) return x inputs = keras.Input(shape=(7, 7, 3)) outputs = SubclassLayer()(inputs) model = keras.Model(inputs, outputs) with self.assertRaises(ValueError): model = prune.prune_low_magnitude(model, **self.params) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model) model.fit(np.random.randn( *self._batch(model.input.get_shape().as_list(), 32)), np.random.randn( *self._batch(model.output.get_shape().as_list(), 32)), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity_m_by_n(self, model, m_by_n) self._check_strip_pruning_matches_original(model, 0.5)
def testPrunesSingleLayer_ReachesTargetSparsity(self, layer_type): model = keras.Sequential() args, input_shape = self._get_params_for_layer(layer_type) if args is None: return # Test for layer not supported yet. model.add(prune.prune_low_magnitude( layer_type(*args), input_shape=input_shape, **self.params)) model.compile( loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model) model.fit( np.random.randn(*self._batch(model.input.get_shape().as_list(), 32)), np.random.randn(*self._batch(model.output.get_shape().as_list(), 32)), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity(self, 0.5, model) self._check_strip_pruning_matches_original(model, 0.5)
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 testPruneRecursivelyReachesTargetSparsity(self): internal_model = keras.Sequential( [keras.layers.Dense(10, input_shape=(10, ))]) model = keras.Sequential([ internal_model, layers.Flatten(), layers.Dense(1), ]) model.compile(loss='binary_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model) model.fit(np.random.randint(10, size=(32, 10)), np.random.randint(2, size=(32, 1)), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity(self, 0.5, model) input_data = np.random.randint(10, size=(32, 10)) self._check_strip_pruning_matches_original(model, 0.5, input_data)
def testPrunesEmbedding(self): model = keras.Sequential() model.add( prune.prune_low_magnitude(keras.layers.Embedding(input_dim=10, output_dim=3), input_shape=(5, ), **self.params)) model.add(keras.layers.Flatten()) model.add(keras.layers.Dense(1, activation='sigmoid')) model.compile(loss='binary_crossentropy', optimizer='sgd', metrics=['accuracy']) test_utils.assert_model_sparsity(self, 0.0, model) model.fit(np.random.randint(10, size=(32, 5)), np.random.randint(2, size=(32, 1)), callbacks=[pruning_callbacks.UpdatePruningStep()]) test_utils.assert_model_sparsity(self, 0.5, model) self._check_strip_pruning_matches_original(model, 0.5)
def testUpdatePruningStepsAndLogsSummaries_CustomTrainingLoop(self): log_dir = tempfile.mkdtemp() pruned_model, loss, optimizer, x_train, y_train = self._pruned_model_setup( custom_training_loop=True) unused_arg = -1 step_callback = pruning_callbacks.UpdatePruningStep() log_callback = pruning_callbacks.PruningSummaries(log_dir=log_dir) # TODO(tfmot): we need a separate API for custom training loops # that doesn't rely on users setting the model and optimizer. # # Example is currently based on callbacks.py configure_callbacks # and model.compile internals. step_callback.set_model(pruned_model) log_callback.set_model(pruned_model) pruned_model.optimizer = optimizer step_callback.on_train_begin() for _ in range(3): log_callback.on_epoch_begin(epoch=unused_arg) # only one batch given batch_size = 20 and input shape. step_callback.on_train_batch_begin(batch=unused_arg) inp = np.reshape( x_train, [self._BATCH_SIZE, 10]) # original shape: from [10]. with tf.GradientTape() as tape: logits = pruned_model(inp, training=True) loss_value = loss(y_train, logits) grads = tape.gradient(loss_value, pruned_model.trainable_variables) optimizer.apply_gradients( zip(grads, pruned_model.trainable_variables)) step_callback.on_epoch_end(batch=unused_arg) self.assertEqual( 3, tf.keras.backend.get_value(pruned_model.layers[0].pruning_step)) self.assertEqual( 3, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step)) self._assertLogsExist(log_dir)
def testUpdatePruningStepsAndLogsSummaries(self): log_dir = tempfile.mkdtemp() model = prune.prune_low_magnitude( keras_test_utils.build_simple_dense_model()) model.compile( loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) model.fit( np.random.rand(20, 10), tf.keras.utils.to_categorical(np.random.randint(5, size=(20, 1)), 5), batch_size=20, epochs=3, callbacks=[ pruning_callbacks.UpdatePruningStep(), pruning_callbacks.PruningSummaries(log_dir=log_dir) ]) self.assertEqual(2, tf.keras.backend.get_value(model.layers[0].pruning_step)) self.assertEqual(2, tf.keras.backend.get_value(model.layers[1].pruning_step)) self._assertLogsExist(log_dir)
def testPruneTrainingRaisesError_PruningStepCallbackMissing(self): model = prune.prune_low_magnitude( keras.Sequential([ layers.Dense(10, activation='relu', input_shape=(100, )), layers.Dense(2, activation='sigmoid') ]), **self.params) model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.SGD(), metrics=['accuracy']) # Throws an error since UpdatePruningStep is missing. with self.assertRaises(errors_impl.InvalidArgumentError): model.fit( np.random.rand(1000, 100), keras.utils.to_categorical(np.random.randint(2, size=(1000, 1)))) model.fit( np.random.rand(1000, 100), keras.utils.to_categorical(np.random.randint(2, size=(1000, 1))), # Works when callback is provided. callbacks=[pruning_callbacks.UpdatePruningStep()])
def prune_model(original_model, train_images, train_labels): batch_size = 256 epochs = 5 pruning_params = { 'pruning_schedule': pruning_schedule.ConstantSparsity(0.75, begin_step=0, frequency=100) } pruning_model = prune.prune_low_magnitude(original_model, **pruning_params) pruning_model.summary() callbacks = [pruning_callbacks.UpdatePruningStep()] fit_kwargs = { 'batch_size': batch_size, 'epochs': epochs, 'callbacks': callbacks, } compile_and_fit(pruning_model, train_images, train_labels, compile_kwargs={}, fit_kwargs=fit_kwargs) return pruning_model
callbacks=all_callbacks(stop_patience=yamlparameters["Training_early_stopping"], initial_lr=yamlparameters["Training_learning_rate"], lr_factor=yamlparameters["Training_lr_factor"], lr_patience=yamlparameters["Training_lr_patience"], lr_epsilon=yamlparameters["Training_lr_min_delta"], lr_cooldown=yamlparameters["Training_lr_cooldown"], lr_minimum=yamlparameters["Training_lr_minimum"], Prune_begin=experiment.get_parameter("pruning_begin_epoch"), Prune_end=experiment.get_parameter("pruning_end_epoch"), prune_lrs=[experiment.get_parameter("pruning_lr_factor_1"), experiment.get_parameter("pruning_lr_factor_2"), experiment.get_parameter("pruning_lr_factor_3")], outputDir=yamlparameters["TrainDir"]) callbacks.callbacks.append(pruning_callbacks.UpdatePruningStep()) with experiment.train(): keras_model.fit(X_train,y_train, batch_size=yamlparameters["Training_batch_size"], epochs=yamlparameters["Training_epochs"], callbacks=callbacks.callbacks, verbose=1, validation_split=yamlparameters["Training_validation_split"], shuffle=True) keras_model = strip_pruning(keras_model) keras_model.compile(optimizer=adam, loss='binary_crossentropy', metrics=['binary_accuracy']) keras_model.save(yamlparameters["TrainDir"]+"/Best_model.h5")
model.add(keras.layers.Embedding(max_features, 128, input_length=maxlen)) model.add(keras.layers.LSTM(128)) # try using a GRU instead, for fun model.add(keras.layers.Dropout(0.5)) model.add(keras.layers.Dense(1)) model.add(keras.layers.Activation("sigmoid")) model = prune.prune_low_magnitude( model, pruning_schedule.PolynomialDecay(initial_sparsity=0.3, final_sparsity=0.7, begin_step=1000, end_step=3000)) # try using different optimizers and different optimizer configs model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["accuracy"]) print_model_sparsity(model) print("Train...") model.fit(x_train, y_train, batch_size=batch_size, epochs=3, callbacks=[pruning_callbacks.UpdatePruningStep()], validation_data=(x_test, y_test)) score, acc = model.evaluate(x_test, y_test, batch_size=batch_size) print_model_sparsity(model) print("Test score:", score) print("Test accuracy:", acc)