def load_gru_pruned(weights_file, debug=False): epochs = 100 batch_size = 32 num_train_samples = 7352 end_step = np.ceil(1.0 * num_train_samples / batch_size).astype( np.int32) * epochs # print('End step: ' + str(end_step)) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=2000, end_step=end_step, frequency=100) } pruned_model = tf.keras.Sequential([ sparsity.prune_low_magnitude( GRU(100, return_sequences=True, input_shape=(128, 9)), **pruning_params), sparsity.prune_low_magnitude(GRU(64), **pruning_params), Dropout(0.15), sparsity.prune_low_magnitude(Dense(64, activation='tanh'), **pruning_params), sparsity.prune_low_magnitude(Dense(6, activation='softmax'), **pruning_params) ]) # Loading Model Weights pruned_model.load_weights(weights_file) if debug: pruned_model.summary() return pruned_model
def sparsePrune(logger, model, X_train, Y_train, X_test, Y_test, num_train_samples, batch_size, epochs, initSparse, endSparse): end_step = np.ceil(1.0 * num_train_samples / batch_size).astype( np.int32) * epochs # TODO determine how to limit this pruning to retain 90% of the network weights / size new_pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=initSparse, final_sparsity=endSparse, begin_step=0, end_step=end_step, frequency=100) } new_model = sparsity.prune_low_magnitude(model, **new_pruning_params) new_model.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy']) callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=None, profile_batch=0) ] new_model.fit(X_train, Y_train, batch_size=batch_size, epochs=epochs, verbose=1, callbacks=callbacks, validation_data=(X_test, Y_test)) score = new_model.evaluate(X_test, Y_test, verbose=0) logger.info('Sparsely Pruned Network Experiment-Results') logger.info('Test loss:', score[0]) logger.info('Test accuracy:', score[1]) return new_model
def go(batch_size, epochs, dataset): num_classes = 10 x_train, y_train, x_test, y_test, input_shape = get_data(dataset, num_classes) print('x_train shape:', x_train.shape) print(x_train.shape[0], 'train samples') print(x_test.shape[0], 'test samples') num_train_samples = x_train.shape[0] end_step = np.ceil(1.0 * num_train_samples / batch_size).astype(np.int32) * epochs print('End step: ' + str(end_step)) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=2000, end_step=end_step, frequency=100) } model = train(get_model(input_shape, num_classes, pruning_params=pruning_params), x_train, y_train, batch_size, epochs, x_test, y_test, pruning=True) score = model.evaluate(x_test, y_test, verbose=0) print('Test loss:', score[0]) print('Test accuracy:', score[1]) keras_file = "mnist_optimized.h5" print('Saving model to: ', keras_file) # Save removing pruning apparatus tf.keras.models.save_model(sparsity.strip_pruning(model), keras_file, include_optimizer=False)
def make_model(n_vert, n_feat, n_class=2): n_aggregators = 4 n_filters = 8 n_propagate = 8 x = keras.layers.Input(shape=(n_vert, n_feat)) n = keras.layers.Input(shape=(1, ), dtype='int32') inputs = [x, n] pruning_params = { 'pruning_schedule': prune.PolynomialDecay(initial_sparsity=0.30, final_sparsity=0.8, begin_step=1000, end_step=2734, frequency=100) } v = inputs v = [GarNet(4, 8, 8, input_format='xn', name='gar_1')(v), n] v = [GarNet(4, 8, 8, input_format='xn', name='gar_2')(v), n] v = GarNet(4, 8, 8, collapse='mean', input_format='xn', name='gar_3')(v) # pruning each layer# #v = [prune.prune_low_magnitude(GarNet(4,8,8, input_format='xn', name='gar_1'), **pruning_params)(v), n] #v = [prune.prune_low_magnitude(GarNet(4,8,8, input_format='xn', name='gar_2'), **pruning_params)(v), n] #v = prune.prune_low_magnitude(GarNet(n_aggregators, n_filters, n_propagate, collapse='mean', input_format='xn', name='gar_3'), **pruning_params)(v) if n_class == 2: v = keras.layers.Dense(1, activation='sigmoid')(v) else: v = keras.layers.Dense(1, activation='softmax')(v) outputs = v return keras.Model(inputs=inputs, outputs=outputs)
def set_pruning_params(final_sparsity, begin_step, frequency, end_step): pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0, final_sparsity=final_sparsity, begin_step=begin_step, end_step=end_step, frequency=frequency) } return pruning_params
def train(): (x_train, y_train), (x_test, y_test) = input_fn() num_train_samples = x_train.shape[0] end_step = np.ceil(1.0 * num_train_samples / BATCH_SIZE).astype( np.int32) * EPOCHS print('End step: ' + str(end_step)) print('Train input shape: ', x_train.shape) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.20, final_sparsity=0.95, begin_step=500, end_step=end_step, frequency=100) } pruned_model = get_model(input_shape=x_train.shape[1:], pruning_params=pruning_params) pruned_model.compile(loss=tf.keras.losses.sparse_categorical_crossentropy, optimizer='adam', metrics=['accuracy']) pruned_model.summary() initial_ws = save_weights(pruned_model) logdir = 'logs' # 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 = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=logdir, profile_batch=0), ResetCallback(initial_ws, 2) ] pruned_model.fit(x_train, y_train, batch_size=BATCH_SIZE, epochs=EPOCHS, verbose=1, callbacks=callbacks, validation_data=(x_test, y_test)) score = pruned_model.evaluate(x_test, y_test, verbose=0) print('Test loss:', score[0]) print('Test accuracy:', score[1])
def compile(self, x_train: np.array, batch_size: int, n_epochs: int, **compile_kwargs) -> None: end_step = calculate_last_train_step(x_train, batch_size, n_epochs) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay( initial_sparsity=self._initial_sparsity, final_sparsity=self._final_sparsity, begin_step=self._begin_step, end_step=end_step, frequency=self._frequency) } self._model = sparsity.prune_low_magnitude(self._model, **pruning_params) self._model.compile(**compile_kwargs)
def get_pruning_params(num_train_samples, initial_sparsity, final_sparsity, begin_step, frequency, batch_size, pruning_epochs): end_step = (np.ceil(num_train_samples / batch_size).astype(np.int32) * pruning_epochs) return { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=initial_sparsity, final_sparsity=final_sparsity, begin_step=begin_step, end_step=end_step, frequency=frequency) }
def get_pruning_model(model, begin_step, end_step): import tensorflow as tf if tf.__version__.startswith('2'): # model pruning API is not supported in TF 2.0 yet raise Exception('model pruning is not fully supported in TF 2.x, Please switch env to TF 1.x for this feature') pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.0, final_sparsity=0.7, begin_step=begin_step, end_step=end_step, frequency=100) } pruning_model = sparsity.prune_low_magnitude(model, **pruning_params) return pruning_model
def apply_pruning_to_dense(layer): end_step = np.ceil(1.0 * num_train_examples / BATCH_SIZE).astype( np.int32) * EPOCHS pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.0, final_sparsity=sparsity_target, begin_step=0, end_step=end_step, frequency=100) } if isinstance(layer, tf.keras.layers.Dense): return sparsity.prune_low_magnitude( layer, **pruning_params ) # note: can pass a whole model instead of individual layers to prune_low_magnitude if desired return layer
def pruneFunction(layer): # pruning_params = {'pruning_schedule': sparsity.ConstantSparsity(0.75, begin_step=2000, frequency=100)} pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.40, final_sparsity=0.50, begin_step=1000, end_step=8000, frequency=100) } if isinstance(layer, tf.keras.layers.Conv2D): return tfmot.sparsity.keras.prune_low_magnitude( layer, **pruning_params) if isinstance(layer, tf.keras.layers.Dense) and layer.name != 'output': return tfmot.sparsity.keras.prune_low_magnitude( layer, **pruning_params) return layer
def make_pruning(model, train_dataset, validation_dataset, n_step, v_step): end_step = np.ceil(1.0 * n_step / config.batch_size).astype(np.int32) * config.p_epochs pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=config.initial_sparsity, final_sparsity=config.final_sparsity, begin_step=config.p_begin_step, end_step=end_step, frequency=config.p_frequency) } p_model = sparsity.prune_low_magnitude(model, **pruning_params) model_setup(p_model) callbacks = callbacks_init() p_model.fit(train_dataset, epochs= config.p_epochs, verbose=1, callbacks=callbacks, validation_data=validation_dataset, steps_per_epoch= n_step, validation_steps= v_step) p_model = sparsity.strip_pruning(p_model) return p_model
def setup_model(sparsity_to, sparsity_from=0, weights=None): pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=sparsity_from, final_sparsity=sparsity_to, begin_step=PRUNING_START, end_step=PRUNING_END, frequency=PRUNING_FREQ) } m = tf.keras.Sequential([ tf.keras.layers.Conv2D(input_shape=(28, 28, 1), filters=6, kernel_size=(5, 5), padding="same", activation="relu"), tf.keras.layers.MaxPool2D(pool_size=(2, 2)), tf.keras.layers.Conv2D(filters=16, kernel_size=(5, 5), padding="valid", activation="relu"), tf.keras.layers.MaxPool2D(pool_size=(2, 2)), tf.keras.layers.Flatten(), sparsity.prune_low_magnitude( tf.keras.layers.Dense(units=120, activation="relu"), **pruning_params), sparsity.prune_low_magnitude( tf.keras.layers.Dense(units=84, activation="relu"), **pruning_params), tf.keras.layers.Dense(units=10, activation="softmax") ]) m.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.01), loss=tf.keras.losses.categorical_crossentropy, metrics=['accuracy']) if weights is not None: m.set_weights(weights) return m
f.write(model_path) # Pruning setup # Hint: modify the epochs to let model recover epochs = 10 end_step = train_generator.__len__() * epochs # Define pruning paramaters # Hint1: pruned model needs steps to recover # Hint2: initial sparsity too large will lead to low acc # TODO Compare result with final sparsity 0.25, 0.5, 0.75 pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.10, final_sparsity=0.75, begin_step=0, end_step=end_step, frequency=200) } # Assign pruning paramaters pruned_model = sparsity.prune_low_magnitude(model, **pruning_params) # Print the converted model pruned_model.summary() pruned_model.compile(loss='binary_crossentropy', optimizer=tf.keras.optimizers.SGD(lr=0.001), metrics=['acc']) callbacks = [
# Model reconstruction from JSON file with open('./model/KERAS_mnist_mlp%s.json'%num_neurons, 'r') as f: model = model_from_json(f.read()) # Load weights into the new model model.load_weights('./model/KERAS_mnist_mlp%s_weights.h5'%num_neurons) model.summary() end_step = np.ceil(1.0 * num_train_samples / batch_size).astype(np.int32) * epochs print(end_step) new_pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=initial_sparsity, final_sparsity=final_sparsity, begin_step=0, end_step=end_step, frequency=100) } pruned_model = sparsity.prune_low_magnitude(model, **new_pruning_params) pruned_model.summary() pruned_model.compile(loss='categorical_crossentropy', optimizer=optimizer, metrics=['accuracy']) callbacks = [tf.keras.callbacks.TensorBoard(log_dir=logdir, profile_batch=0), sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=logdir, profile_batch=0) ]
tf.app.flags.DEFINE_float("initial_sparsity", 0.00, "Sparsity at which pruning begins") tf.app.flags.DEFINE_float("final_sparsity", 0.50, " Sparsity at which pruning ends") tf.app.flags.DEFINE_integer("begin_step", 2000, " Step at which to begin pruning") tf.app.flags.DEFINE_integer("end_step", 8000, " Step at which to end pruning") FLAGS = tf.app.flags.FLAGS FLAGS(sys.argv) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=FLAGS.initial_sparsity, final_sparsity=FLAGS.final_sparsity, begin_step=FLAGS.begin_step, end_step=FLAGS.end_step, frequency=200) } def preprocess(): """ Prepare data """ data = FLAGS.dataset text = load_text(data) vocab = sorted(set(text)) # Creating a mapping from unique characters to indices char2idx = {u: i for i, u in enumerate(vocab)}
from tensorflow_model_optimization.sparsity import keras as sparsity """To demonstrate how to save and restore a pruned keras model, in the following example we first train the model for 10 epochs, save it to disk, and finally restore and continue training for 2 epochs. With gradual sparsity, four important parameters are begin_sparsity, final_sparsity, begin_step and end_step. The first three are straight forward. Let's calculate the end step given the number of train example, batch size, and the total epochs to train.""" import numpy as np epochs = 12 num_train_samples = x_train.shape[0] end_step = np.ceil(1.0 * num_train_samples / batch_size).astype(np.int32) * epochs print('End step: ' + str(end_step)) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=2000, end_step=end_step, frequency=100) } pruned_model = tf.keras.Sequential([ sparsity.prune_low_magnitude( l.Conv2D(32, 5, padding='same', activation='relu'), input_shape=input_shape, **pruning_params), l.MaxPooling2D((2, 2), (2, 2), padding='same'), l.BatchNormalization(), sparsity.prune_low_magnitude( l.Conv2D(64, 5, padding='same', activation='relu'), **pruning_params), l.MaxPooling2D((2, 2), (2, 2), padding='same'), l.Flatten(),
def layer_pruned_model(): #Build a pruned model layer by layer epochs = 12 (x_train, y_train), (x_test, y_test) = prepare_data() num_train_samples = x_train.shape[0] end_step = np.ceil(1.0 * num_train_samples / batch_size).astype( np.int32) * epochs print('End step: ' + str(end_step)) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=2000, end_step=end_step, frequency=100) } #build the model l = tf.keras.layers pruned_model = tf.keras.Sequential([ sparsity.prune_low_magnitude(l.Conv2D(32, 5, padding='same', activation='relu'), input_shape=input_shape, **pruning_params), l.MaxPooling2D((2, 2), (2, 2), padding='same'), l.BatchNormalization(), sparsity.prune_low_magnitude( l.Conv2D(64, 5, padding='same', activation='relu'), **pruning_params), l.MaxPooling2D((2, 2), (2, 2), padding='same'), l.Flatten(), sparsity.prune_low_magnitude(l.Dense(1024, activation='relu'), **pruning_params), l.Dropout(0.4), sparsity.prune_low_magnitude( l.Dense(num_classes, activation='softmax'), **pruning_params) ]) pruned_model.summary() logdir = tempfile.mkdtemp() print('Writing training logs to ' + logdir) # %tensorboard --logdir={logdir} # train the model pruned_model.compile(loss=tf.keras.losses.categorical_crossentropy, optimizer='adam', metrics=['accuracy']) callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=logdir, profile_batch=0) ] pruned_model.fit(x_train, y_train, batch_size=batch_size, epochs=10, verbose=1, callbacks=callbacks, validation_data=(x_test, y_test)) score = pruned_model.evaluate(x_test, y_test, verbose=0) print('Test loss:', score[0]) print('Test accuracy:', score[1]) # Save and restore checkpoint_file = './pruned_checkpoint_file.h5' # _, checkpoint_file = tempfile.mkstemp('.h5') print('Saving pruned model to: ', checkpoint_file) # saved_model() sets include_optimizer to True by default. Spelling it out here # to highlight. tf.keras.models.save_model(pruned_model, checkpoint_file, include_optimizer=True) with sparsity.prune_scope(): restored_model = tf.keras.models.load_model(checkpoint_file) restored_model.fit(x_train, y_train, batch_size=batch_size, epochs=2, verbose=1, callbacks=callbacks, validation_data=(x_test, y_test)) start_test = time.time() score = restored_model.evaluate(x_test, y_test, verbose=0) end_test = time.time() print('Test latency:', end_test - start_test) print('Test loss:', score[0]) print('Test accuracy:', score[1]) final_model = sparsity.strip_pruning(pruned_model) final_model.summary() layer_pruned_file = './layer_pruned_file.h5' # _, layer_pruned_file = tempfile.mkstemp('.h5') print('Saving pruned model to: ', layer_pruned_file) tf.keras.models.save_model(final_model, layer_pruned_file, include_optimizer=False)
validation_generator = validation_datagen.flow_from_directory( VAL_DIR, target_size=(SCALED_HEIGHT, SCALED_WIDTH), batch_size=BATCH_SIZE, class_mode='categorical') # load model and pruning parameters model = load_model('keras_baseline/baseline_keras_inception_v3.h5') epochs = 5 end_step = np.ceil(1.0 * NUM_TRAIN / BATCH_SIZE).astype(np.int32) * epochs pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=1, end_step=end_step, frequency=100) } # prune model pruned_model = sparsity.prune_low_magnitude(model, **pruning_params) pruned_model.compile(optimizer='rmsprop', loss='categorical_crossentropy', metrics=['accuracy']) # load callback steps logdir = tempfile.mkdtemp() callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=logdir, profile_batch=0)
eval_labels = to_categorical( np.load(feature_dir + 'raw_eval_labels.npy', allow_pickle=True)) ## DEFINE TRAINING PARAMETERS batch_size = 128 n_batches = len(train_features) // batch_size + 1 * ( len(train_features) % batch_size != 0) model_path = 'models' + os.sep + 'model3' + os.sep ## DEFINE PRUNING PARAMETERS pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay( initial_sparsity=0.1, final_sparsity=0.8, begin_step=n_batches * 100, # Begin pruning at epoch 100 end_step=n_batches * 300, # End pruning at epoch 300 (but continue to train for 100 more epochs) frequency=n_batches * 10) } ## GENERATE INCEPTION MODEL (SUBSYSTEM 1) model = inception_prunable(**pruning_params) print('Generated Model 3, Subsystem 1. Training it now...') ## TRAIN INCEPTION MODEL model_hist = trainp_model(model, train_features, train_labels, eval_features, eval_labels,
def create_estimator(steps=None, warmup_steps=None, model_dir=args.model_dir, num_labels=args.num_labels, max_seq_len=args.max_seq_len, learning_rate=args.learning_rate, name='bert'): def my_auc(labels, predictions): auc_metric = tf.keras.metrics.AUC(name="my_auc") auc_metric.update_state(y_true=labels, y_pred=tf.argmax(predictions, 1)) return {'auc': auc_metric} if name == 'bert': if warmup_steps is None: custom_objects = { 'BertModelLayer': bert.BertModelLayer, 'AdamW': AdamW, 'PruneLowMagnitude': PruneLowMagnitude } if args.prune_enabled: with sparsity.prune_scope(): model = tf.keras.models.load_model(h5py.File(args.keras_model_path), custom_objects=custom_objects) else: model = tf.keras.models.load_model(h5py.File(args.keras_model_path), custom_objects=custom_objects) estimator = tf.keras.estimator.model_to_estimator(model, model_dir=args.output_dir) return estimator, model input_token_ids = tf.keras.Input((max_seq_len,), dtype=tf.int32, name='input_ids') input_segment_ids = tf.keras.Input((max_seq_len,), dtype=tf.int32, name='segment_ids') input_mask = tf.keras.Input((max_seq_len,), dtype=tf.int32, name='input_mask') bert_params = bert.params_from_pretrained_ckpt(model_dir) l_bert = bert.BertModelLayer.from_params(bert_params) bert_output = l_bert(inputs=[input_token_ids, input_segment_ids], mask=input_mask) if args.pool_strategy == 'cls': first_token = tf.keras.layers.Lambda(lambda seq: seq[:, 0, :])(bert_output) pooled_output = tf.keras.layers.Dense(units=first_token.shape[-1], activation=tf.math.tanh)(first_token) dropout = tf.keras.layers.Dropout(rate=0.1)(pooled_output) elif args.pool_strategy == 'avg': seq1_tokens = tf.keras.layers.Lambda(lambda seq: seq[:,1:args.max_seq_len-1,:])(bert_output) seq2_tokens = tf.keras.layers.Lambda(lambda seq: seq[:,args.max_seq_len:2*args.max_seq_len]) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=1000, end_step=2000, frequency=100) } dense = tf.keras.layers.Dense(units=num_labels, name='label_ids') if args.prune_enabled: pruned_dense = sparsity.prune_low_magnitude( dense, **pruning_params) logits = pruned_dense(dropout) else: logits = dense(dropout) output_prob = tf.keras.layers.Softmax(name='output_prob')(logits) model = tf.keras.Model(inputs=[input_token_ids, input_segment_ids, input_mask], outputs=[logits]) model.build(input_shape=[(None, max_seq_len,), (None, max_seq_len,), (None, max_seq_len,)]) # freeze_bert_layers(l_bert) bert.load_stock_weights(l_bert, op.join(model_dir, 'bert_model.ckpt')) weight_decays = get_weight_decays(model) for k, v in weight_decays.items(): if use_weight_decay(k): weight_decays[k] = 0.01 else: del weight_decays[k] opt = create_optimizer( init_lr=learning_rate, steps=steps, weight_decays=weight_decays, warmup_steps=warmup_steps, ) model.compile( optimizer=opt, loss={"{}label_ids".format( 'prune_low_magnitude_' if args.prune_enabled else ''): tf.keras.losses.SparseCategoricalCrossentropy( from_logits=True)}, # for numerical stability metrics=[tf.keras.metrics.SparseCategoricalAccuracy()] ) model.summary() config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True config.gpu_options.per_process_gpu_memory_fraction = args.gpu_memory_fraction config.log_device_placement = False exclude_optimizer_variables = r'^((?!(iter_updates|eta_t)).)*$' ws = tf.estimator.WarmStartSettings( ckpt_to_initialize_from=op.join(args.output_dir, 'keras'), vars_to_warm_start=exclude_optimizer_variables ) estimator = tf.keras.estimator.model_to_estimator(keras_model=model, config=tf.estimator.RunConfig( model_dir=args.output_dir, session_config=config, )) estimator._warm_start_settings = ws return estimator, model raise NotImplemented("* available models: [ bert, ]")
def prune_Conv1D(final_sparsity, initial_sparsity=0.0, begin_step=0, frequency=100, version=""): # Set up some params nb_epoch = 50 # number of epochs to train on batch_size = 1024 # training batch size num_train_samples = X_train.shape[0] end_step = np.ceil(1.0 * num_train_samples / batch_size).astype( np.int32) * nb_epoch print("End step: ", end_step) pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=initial_sparsity, final_sparsity=final_sparsity, begin_step=begin_step, end_step=end_step, frequency=100) } l = tf.keras.layers dr = 0.5 # dropout rate (%) pruned_model = tf.keras.Sequential([ sparsity.prune_low_magnitude( l.Conv1D(128, 3, padding='valid', activation="relu", name="conv1", kernel_initializer='glorot_uniform', input_shape=in_shape), **pruning_params), sparsity.prune_low_magnitude( l.Conv1D(128, 3, padding='valid', activation="relu", name="conv2", kernel_initializer='glorot_uniform'), **pruning_params), l.MaxPool1D(2), sparsity.prune_low_magnitude( l.Conv1D(64, 3, padding='valid', activation="relu", name="conv3", kernel_initializer='glorot_uniform'), **pruning_params), sparsity.prune_low_magnitude( l.Conv1D(64, 3, padding='valid', activation="relu", name="conv4", kernel_initializer='glorot_uniform'), **pruning_params), l.Dropout(dr), sparsity.prune_low_magnitude( l.Conv1D(32, 3, padding='valid', activation="relu", name="conv5", kernel_initializer='glorot_uniform'), **pruning_params), sparsity.prune_low_magnitude( l.Conv1D(32, 3, padding='valid', activation="relu", name="conv6", kernel_initializer='glorot_uniform'), **pruning_params), l.Dropout(dr), l.MaxPool1D(2), l.Flatten(), sparsity.prune_low_magnitude( l.Dense(128, activation='relu', kernel_initializer='he_normal', name="dense1"), **pruning_params), sparsity.prune_low_magnitude( l.Dense(len(classes), kernel_initializer='he_normal', name="dense2"), **pruning_params), l.Activation('softmax') ]) pruned_model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=["accuracy"]) pruned_model.summary() callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=logdir, profile_batch=0) ] history = pruned_model.fit(X_train, Y_train, batch_size=batch_size, epochs=nb_epoch, verbose=1, validation_data=(X_val, Y_val), callbacks=callbacks) score = pruned_model.evaluate(X_test, Y_test, verbose=0) print("Test loss: ", score) #Save the model pruned_model = sparsity.strip_pruning(pruned_model) pruned_model.summary() # Save the model architecture print_model_to_json( pruned_model, './model/Conv1D-{}.json'.format(str(final_sparsity) + version)) # Save the weights pruned_model.save_weights( './model/Conv1D-{}.h5'.format(str(final_sparsity) + version))
def train(cfg): epochs = cfg['epochs'] save_dir = cfg['save_dir'] if not os.path.exists(save_dir): os.mkdir(save_dir) shape = (int(cfg['height']), int(cfg['width']), 3) n_class = int(cfg['class_number']) batch_size = int(cfg['batch_size']) if cfg['model'] == 'mymodel': from model.my_model import MyModel model = MyModel(shape, n_class).build() if cfg['model'] == 'v2': from model.mobilenet_v2 import MyModel model = MyModel(shape, n_class).buildRaw() train_generator, validation_generator, count1, count2 = generate(batch_size, shape[:2], cfg['train_dir'], cfg['eval_dir']) print(count1, count2) earlystop = EarlyStopping(monitor='val_acc', patience=4, verbose=0, mode='auto') checkpoint = ModelCheckpoint(filepath=os.path.join("save", 'prune_e_{epoch:02d}_{val_loss:.3f}_{val_acc:.3f}.h5'), monitor='val_acc', save_best_only=False, save_weights_only=False) reduce_lr = ReduceLROnPlateau(monitor='val_acc', factor=0.1, patience=2, verbose=1, min_lr=1e-7) model_path = r'./save/v2' # x_train, y_train = train_generator.next() # num_train_samples = batch_size # x_test, y_test = validation_generator.next() loaded_model = tf.keras.models.load_model(os.path.join(model_path,'e_06_0.20_1.00.h5')) score = loaded_model.evaluate_generator(validation_generator, count2//batch_size) print('original Test loss:', score[0]) print('original Test accuracy:', score[1]) end_step = np.ceil(1.0 * count1 / batch_size).astype(np.int32) * epochs print(end_step) new_pruning_params = {'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=0, end_step=end_step, frequency=100)} new_pruned_model = sparsity.prune_low_magnitude(loaded_model, **new_pruning_params) #new_pruned_model.summary() opt = Adam(lr=float(0.0001)) new_pruned_model.compile(loss=tf.keras.losses.categorical_crossentropy, optimizer=opt, metrics=['acc']) #现在我们开始训练和修剪模型。 #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 logdir = "./save/log" callbacks = [earlystop,checkpoint,reduce_lr, sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=logdir, profile_batch=0)] # new_pruned_model.fit(x_train, y_train, # batch_size=batch_size, # epochs=epochs, # verbose=1, # callbacks=callbacks, # validation_data=(x_test, y_test)) new_pruned_model.fit_generator(train_generator, validation_data=validation_generator, steps_per_epoch=100,#count1 // batch_size, validation_steps=count2 // batch_size, epochs=epochs, callbacks=callbacks) score = new_pruned_model.evaluate_generator(validation_generator, count2//batch_size) print('Test loss:', score[0]) print('Test accuracy:', score[1]) final_model = sparsity.strip_pruning(new_pruned_model) new_pruned_keras_file = "save/pruned_model.h5" tf.keras.models.save_model(final_model, new_pruned_keras_file, include_optimizer=False)