def get_model(input_shape=(28, 28, 1), pruning_params=None): assert pruning_params is not None, "You must specify pruning params." pruned_model = tf.keras.Sequential([ sparsity.prune_low_magnitude( l.Conv2D(32, 5, padding='same', activation='relu'), **pruning_params), l.MaxPooling2D((2, 2), (2, 2), padding='same'), sparsity.prune_low_magnitude( l.Conv2D(64, 5, padding='same', activation='relu'), **pruning_params), l.MaxPooling2D((2, 2), (2, 2), padding='same'), 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), sparsity.prune_low_magnitude( l.Dense(num_classes, activation='softmax'), **pruning_params) ]) pruned_model.build(input_shape=(None, ) + input_shape) return pruned_model
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 pruned_lstm(vocab_size, embedding_dim, rnn_units, batch_size): """ Function that builds the pruned LSTM model using the pruning parameters Args: vocab_size: Length of the vocabulary in chars embedding_dim: The embedding dimension rnn_units: Number of RNN units batch_size: Batch size Returns: LSTM model """ pruned_model = tf.keras.Sequential([ tf.keras.layers.Embedding(vocab_size, embedding_dim, batch_input_shape=[batch_size, None]), sparsity.prune_low_magnitude( tf.keras.layers.LSTM(rnn_units, return_sequences=True, stateful=True), **pruning_params), sparsity.prune_low_magnitude( tf.keras.layers.LSTM(rnn_units, return_sequences=True, stateful=True), **pruning_params), sparsity.prune_low_magnitude(tf.keras.layers.Dense(vocab_size), **pruning_params) ]) return pruned_model
def __init__(self, nb_funcs, nb_variables, model=None, **kwargs): super(PrunedRegression, self).__init__() if not model is None: self.xi = sparsity.prune_low_magnitude( PrunableLayer(starting_vector=model.xi.layer.w), **kwargs) else: self.xi = sparsity.prune_low_magnitude( PrunableLayer(nb_funcs=nb_funcs, nb_variables=nb_variables, starting_vector=starting_vector), **kwargs)
def initialize_sparse_model(trained_model, pruned_model_with_mask, pm): """ Given a filename (or a model) with weights and a pruned model with its mask, returns a new model with weights in filename and pruned with mask """ model = clone_model(trained_model) model.set_weights(trained_model.get_weights()) sparcity = 1 - pm sprasity_sched = ConstantSparsity( sparcity, 0, # Do sparcity calculation in the first step end_step=0, frequency=10000000) prunned_model_layers = [] for i, layer in enumerate(pruned_model_with_mask.layers): if isinstance(layer, pruning_wrapper.PruneLowMagnitude): l_weights = model.layers[i].get_weights() l_weights[0] = l_weights[0] * layer.pruning_vars[0][1].numpy() model.layers[i].set_weights(l_weights) prunned_model_layers.append( prune_low_magnitude(model.layers[i], sprasity_sched)) else: prunned_model_layers.append(model.layers[i]) prunned_model = Sequential(prunned_model_layers) prunned_model.compile(optimizer=optimizers.SGD(lr=0), loss='sparse_categorical_crossentropy', metrics='accuracy') return prunned_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 get_model(input_shape, num_classes, pruning_params=None): l = tf.keras.layers layers = [ l.Conv2D(32, 5, padding='same', activation='relu', input_shape=input_shape), l.MaxPooling2D((2, 2), (2, 2), padding='same'), l.BatchNormalization(), l.Conv2D(64, 5, padding='same', activation='relu'), l.MaxPooling2D((2, 2), (2, 2), padding='same'), l.Flatten(), l.Dense(1024, activation='relu', activity_regularizer=tf.keras.regularizers.l1(0.01)), l.Dropout(0.4), l.Dense(num_classes, activation='softmax') ] if pruning_params is not None: # Add pruning in the right places for i in [0, 3]: layers[i] = sparsity.prune_low_magnitude(layers[i], **pruning_params) model = tf.keras.Sequential(layers) return model
def prune_model(grid_params, job_id, JOB_RUNS_FOLDER, TRAIN_BASELINE_MODELS_JOB_ID, prune_thresh, pruning_training_params, X_train, y_train): model_for_pruning = build_mnist_model(**grid_params) model_for_pruning.load_weights(base_weights_path(job_id, JOB_RUNS_FOLDER, TRAIN_BASELINE_MODELS_JOB_ID)).expect_partial() epochs_prune = pruning_training_params['epochs'] batch_size_prune = pruning_training_params['batch_size'] val_split_prune = pruning_training_params['validation_split'] pruning_params = {'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay( initial_sparsity=0, final_sparsity=prune_thresh, begin_step=0, end_step=np.ceil(len(X_train) / batch_size_prune).astype(np.int32) * epochs_prune) } model_for_pruning = prune_low_magnitude(model_for_pruning, **pruning_params) # `prune_low_magnitude` requires a recompile. model_for_pruning.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy']) model_for_pruning.fit(X_train, y_train, verbose=False, batch_size=batch_size_prune, epochs=epochs_prune, validation_split=val_split_prune, callbacks=[ tfmot.sparsity.keras.UpdatePruningStep(), ]) return model_for_pruning
def pruned_nn(pruning_params, first_layer=300, second_layer=100): """ Function to define a 300 100 dense fully connected architecture for MNIST classification Arguments: pruning_params=dictionary for the desired pruning rate. see the used_func.py for detailed function first_layer= Desired unit size in first hidden layer second_layer= Desired unit size in second hidden layer returns a model """ #Using 2 different pruning rate, because the layer connected #to the output pruned as half of the rest of the network prun = pruning_params[0] prun_out = pruning_params[1] pruned_model = Sequential() pruned_model.add( sparsity.prune_low_magnitude( Dense(units=first_layer, activation='relu', kernel_initializer=tf.keras.initializers.glorot_normal(), input_shape=(784, )), **prun)) pruned_model.add( sparsity.prune_low_magnitude( Dense(units=second_layer, activation='relu', kernel_initializer=tf.keras.initializers.glorot_normal()), **prun)) pruned_model.add( sparsity.prune_low_magnitude(Dense(units=10, activation='softmax'), **prun_out)) pruned_model.compile( loss=tf.keras.losses.categorical_crossentropy, # 0.0012 learning rate with Adam optimizer is used in the original paper optimizer=tf.keras.optimizers.Adam(0.0012), metrics=['accuracy']) return pruned_model
def prune_and_initilize(trained_model, pm, initial_weights, layers_to_prune=None): sparcity = 1 - pm sprasity_sched = ConstantSparsity( sparcity, 0, # Do sparcity calculation in the first step end_step=0, # Do it only once frequency=10000000) model = clone_model(trained_model) model.set_weights(trained_model.get_weights()) if is_pruned(model): model = strip_pruning(model) if layers_to_prune is None: layers_to_prune = get_default_layers(model) # prunned_model_layers = [] # for layer in model.layers: # if layer.name in layers_to_prune: # prunned_model_layers.append(prune_low_magnitude(layer, sprasity_sched)) # else: # prunned_model_layers.append(layer) # trained_pruned_model = Sequential(prunned_model_layers) trained_pruned_model = apply_wrapper_to_layer(model, layers_to_prune, prune_low_magnitude, sprasity_sched, clone=False) # Calculates mask initialize_pruned_model(trained_pruned_model) model.load_weights(initial_weights) prunned_model_layers = [] for i, layer in enumerate(trained_pruned_model.layers): if isinstance(layer, pruning_wrapper.PruneLowMagnitude): l_weights = model.layers[i].get_weights() l_weights[0] = l_weights[0] * layer.pruning_vars[0][1].numpy() model.layers[i].set_weights(l_weights) prunned_model_layers.append( prune_low_magnitude(model.layers[i], sprasity_sched)) else: prunned_model_layers.append(model.layers[i]) untrained_prunned_model = Sequential(prunned_model_layers) untrained_prunned_model.compile(optimizer=optimizers.SGD(lr=0), loss='sparse_categorical_crossentropy', metrics='accuracy') return untrained_prunned_model
def prune_model(self,model,target_sparsity): """ Extends the ANN's model with priuning layers. Args: model (): The model of the ANN. target_sparsity (float): Target constant sparsity of the network. Returns: model (): The model of the ANN. Notes: Only constant sparsity active. """ if True: model = sparsity.prune_low_magnitude(model,sparsity.ConstantSparsity(target_sparsity,0,frequency=10)) else: model = sparsity.prune_low_magnitude(model,sparsity.PolynomialSparsity(0,target_sparsity,0,self["epochs"], frequency=10)) return 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
def pruned_lstm(vocab_size, embedding_dim, rnn_units, batch_size): """ """ pruned_model = tf.keras.Sequential([ tf.keras.layers.Embedding(vocab_size, embedding_dim, batch_input_shape=[batch_size, None]), sparsity.prune_low_magnitude( tf.keras.layers.LSTM(rnn_units, return_sequences=True, stateful=True), **pruning_params), sparsity.prune_low_magnitude( tf.keras.layers.LSTM(rnn_units, return_sequences=True, stateful=True), **pruning_params), sparsity.prune_low_magnitude( tf.keras.layers.Dense(vocab_size), **pruning_params) ]) return pruned_model
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 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 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 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 get_prunned_model(trained_model, pm=0.5, X_train=None, y_train=None, layers_to_prune=None): """ Given a filename with weights, and a list with layers to prune, returns a pruned model with correct mask X_train, y_train are necesary to get the mask calculated by keras (Needs a fit) pm = 1 - sparcity as mentioned in paper """ model = clone_model(trained_model) model.set_weights(trained_model.get_weights()) sparcity = 1 - pm sprasity_sched = ConstantSparsity( sparcity, 0, # Do sparcity calculation in the first step end_step=0, # Do it only once frequency=10000000) if layers_to_prune is None: layers_to_prune = get_default_layers(model) prunned_model_layers = [] for layer in model.layers: if layer.name in layers_to_prune: prunned_model_layers.append( prune_low_magnitude(layer, sprasity_sched)) else: prunned_model_layers.append(layer) pruned_model = Sequential(prunned_model_layers) del model # This is necesary to make keras calculate the mask, learning rate is 0 initialize_pruned_model(pruned_model) return pruned_model
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(), 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()
def run(network_type, epochs, pruning_epochs, tensorboard_log_dir, model_dir_path, _log, _run): if network_type == 'mlp' or network_type == 'mlp_regression': # don't use gpu for mlps os.environ['CUDA_VISIBLE_DEVICES'] = '' elif network_type == 'cnn': os.environ['CUDA_VISIBLE_DEVICES'] = '3' else: # use gpu for cnn_vggs os.environ['CUDA_VISIBLE_DEVICES'] = '3' assert network_type in ('mlp', 'mlp_regression', 'cnn', 'cnn_vgg') _log.info('Emptying model directory...') if model_dir_path.exists(): shutil.rmtree(model_dir_path) Path(model_dir_path).mkdir(parents=True) _log.info('Loading data...') (X_train, y_train), (X_test, y_test) = load_data() metrics = {} unpruned_model_path, pruned_model_path = get_two_model_paths() unpruned_model = create_model() _log.info('Training unpruned model...') metrics['unpruned'] = train_model(unpruned_model, X_train, y_train, X_test, y_test, unpruned_model_path, epochs=epochs, is_pruning=False) _log.info('Unpruned model sparsity: {}'.format( get_sparsity(unpruned_model))) save_weights(unpruned_model, unpruned_model_path) pruning_params = get_pruning_params(X_train.shape[0]) pruned_model = sparsity.prune_low_magnitude(unpruned_model, **pruning_params) pruning_callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=tensorboard_log_dir, profile_batch=0) ] _log.info('Training pruned model...') metrics['pruned'] = train_model(pruned_model, X_train, y_train, X_test, y_test, pruned_model_path, epochs=pruning_epochs, is_pruning=True, callbacks=pruning_callbacks) _log.info('Pruned model sparsity: {}'.format(get_sparsity(pruned_model))) save_weights(pruned_model, pruned_model_path) ex.add_source_file(__file__) with open(model_dir_path / 'metrics.json', 'w') as f: json.dump(metrics, f, cls=NumpyEncoder) return metrics
def pruing_resnet50(num_classes, batch_size=None, use_l2_regularizer=True, rescale_inputs=False, batch_norm_decay=0.9, batch_norm_epsilon=1e-5): """Instantiates the ResNet50 architecture. Args: num_classes: `int` number of classes for image classification. batch_size: Size of the batches for each step. use_l2_regularizer: whether to use L2 regularizer on Conv/Dense layer. rescale_inputs: whether to rescale inputs from 0 to 1. batch_norm_decay: Moment of batch norm layers. batch_norm_epsilon: Epsilon of batch borm layers. Returns: A Keras model instance. """ #input_shape = (224, 224, 3) input_shape = (28, 28, 1) img_input = layers.Input(shape=input_shape, batch_size=batch_size) if rescale_inputs: # Hub image modules expect inputs in the range [0, 1]. This rescales these # inputs to the range expected by the trained model. x = layers.Lambda(lambda x: x * 255.0 - backend.constant( imagenet_preprocessing.CHANNEL_MEANS, shape=[1, 1, 3], dtype=x.dtype), name='rescale')(img_input) else: x = img_input if backend.image_data_format() == 'channels_first': x = layers.Permute((3, 1, 2))(x) bn_axis = 1 else: # channels_last bn_axis = 3 block_config = dict(use_l2_regularizer=use_l2_regularizer, batch_norm_decay=batch_norm_decay, batch_norm_epsilon=batch_norm_epsilon) x = layers.ZeroPadding2D(padding=(3, 3), name='conv1_pad')(x) x = sparsity.prune_low_magnitude( layers.Conv2D( 64, (7, 7), strides=(2, 2), padding='valid', use_bias=False, kernel_initializer='he_normal', kernel_regularizer=_gen_l2_regularizer(use_l2_regularizer), name='conv1'))(x) x = layers.BatchNormalization(axis=bn_axis, momentum=batch_norm_decay, epsilon=batch_norm_epsilon, name='bn_conv1')(x) x = layers.Activation('relu')(x) x = layers.MaxPooling2D((3, 3), strides=(2, 2), padding='same')(x) x = pruing_conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1), **block_config) x = pruing_identity_block(x, 3, [64, 64, 256], stage=2, block='b', **block_config) x = pruing_identity_block(x, 3, [64, 64, 256], stage=2, block='c', **block_config) x = pruing_conv_block(x, 3, [128, 128, 512], stage=3, block='a', **block_config) x = pruing_identity_block(x, 3, [128, 128, 512], stage=3, block='b', **block_config) x = pruing_identity_block(x, 3, [128, 128, 512], stage=3, block='c', **block_config) x = pruing_identity_block(x, 3, [128, 128, 512], stage=3, block='d', **block_config) x = pruing_conv_block(x, 3, [256, 256, 1024], stage=4, block='a', **block_config) x = pruing_identity_block(x, 3, [256, 256, 1024], stage=4, block='b', **block_config) x = pruing_identity_block(x, 3, [256, 256, 1024], stage=4, block='c', **block_config) x = pruing_identity_block(x, 3, [256, 256, 1024], stage=4, block='d', **block_config) x = pruing_identity_block(x, 3, [256, 256, 1024], stage=4, block='e', **block_config) x = pruing_identity_block(x, 3, [256, 256, 1024], stage=4, block='f', **block_config) x = pruing_conv_block(x, 3, [512, 512, 2048], stage=5, block='a', **block_config) x = pruing_identity_block(x, 3, [512, 512, 2048], stage=5, block='b', **block_config) x = pruing_identity_block(x, 3, [512, 512, 2048], stage=5, block='c', **block_config) x = layers.GlobalAveragePooling2D()(x) x = sparsity.prune_low_magnitude( layers.Dense( num_classes, kernel_initializer=initializers.RandomNormal(stddev=0.01), kernel_regularizer=_gen_l2_regularizer(use_l2_regularizer), bias_regularizer=_gen_l2_regularizer(use_l2_regularizer), name='fc1000'))(x) # A softmax that is followed by the model loss must be done cannot be done # in float16 due to numeric issues. So we pass dtype=float32. x = layers.Activation('softmax', dtype='float32')(x) # Create model. return models.Model(img_input, x, name='pruing_resnet50')
def pruing_conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2), use_l2_regularizer=True, batch_norm_decay=0.9, batch_norm_epsilon=1e-5): """A block that has a conv layer at shortcut. Note that from stage 3, the second conv layer at main path is with strides=(2, 2) And the shortcut should have strides=(2, 2) as well Args: input_tensor: input tensor kernel_size: default 3, the kernel size of middle conv layer at main path filters: list of integers, the filters of 3 conv layer at main path stage: integer, current stage label, used for generating layer names block: 'a','b'..., current block label, used for generating layer names strides: Strides for the second conv layer in the block. use_l2_regularizer: whether to use L2 regularizer on Conv layer. batch_norm_decay: Moment of batch norm layers. batch_norm_epsilon: Epsilon of batch borm layers. Returns: Output tensor for the block. """ filters1, filters2, filters3 = filters if backend.image_data_format() == 'channels_last': bn_axis = 3 else: bn_axis = 1 conv_name_base = 'res' + str(stage) + block + '_branch' bn_name_base = 'bn' + str(stage) + block + '_branch' x = sparsity.prune_low_magnitude( layers.Conv2D( filters1, (1, 1), use_bias=False, kernel_initializer='he_normal', kernel_regularizer=_gen_l2_regularizer(use_l2_regularizer), name=conv_name_base + '2a'))(input_tensor) x = layers.BatchNormalization(axis=bn_axis, momentum=batch_norm_decay, epsilon=batch_norm_epsilon, name=bn_name_base + '2a')(x) x = layers.Activation('relu')(x) x = sparsity.prune_low_magnitude( layers.Conv2D( filters2, kernel_size, strides=strides, padding='same', use_bias=False, kernel_initializer='he_normal', kernel_regularizer=_gen_l2_regularizer(use_l2_regularizer), name=conv_name_base + '2b'))(x) x = layers.BatchNormalization(axis=bn_axis, momentum=batch_norm_decay, epsilon=batch_norm_epsilon, name=bn_name_base + '2b')(x) x = layers.Activation('relu')(x) x = sparsity.prune_low_magnitude( layers.Conv2D( filters3, (1, 1), use_bias=False, kernel_initializer='he_normal', kernel_regularizer=_gen_l2_regularizer(use_l2_regularizer), name=conv_name_base + '2c'))(x) x = layers.BatchNormalization(axis=bn_axis, momentum=batch_norm_decay, epsilon=batch_norm_epsilon, name=bn_name_base + '2c')(x) shortcut = sparsity.prune_low_magnitude( layers.Conv2D( filters3, (1, 1), strides=strides, use_bias=False, kernel_initializer='he_normal', kernel_regularizer=_gen_l2_regularizer(use_l2_regularizer), name=conv_name_base + '1'))(input_tensor) shortcut = layers.BatchNormalization(axis=bn_axis, momentum=batch_norm_decay, epsilon=batch_norm_epsilon, name=bn_name_base + '1')(shortcut) x = layers.add([x, shortcut]) x = layers.Activation('relu')(x) return x
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 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)
# 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 = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir='./', profile_batch=0) ] print('[INFO] Start pruning process...')
end_step = np.ceil(1.0 * 50000 / 128).astype(np.int32) * epochs pruning_params = { 'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50, final_sparsity=0.90, begin_step=0, end_step=end_step, frequency=100) } model = tensorflow.keras.Sequential([ # 1st Convolutional Layer sparsity.prune_low_magnitude( (tensorflow.keras.layers.Conv2D(64, (3, 3), padding='same', activation='relu', input_shape=x_train.shape[1:], strides=(1, 1))), **pruning_params), tensorflow.keras.layers.BatchNormalization(scale=True), tensorflow.keras.layers.Dropout(0.2), # 2nd Convolutional Layer sparsity.prune_low_magnitude( (tensorflow.keras.layers.Conv2D(128, (3, 3), padding='same', activation='relu', input_shape=x_train.shape[1:], strides=(1, 1))), **pruning_params), tensorflow.keras.layers.BatchNormalization(scale=True), tensorflow.keras.layers.Dropout(0.2),
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 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)
#pruning_params = { # 'pruning_schedule': ConstantSparsity(0.75, begin_step=20, frequency=100) # } tfd = tfp.distributions input_shape = x1.shape init = tf.global_variables_initializer() with tf.Session() as sess: sess.run(init) pruned_model = tf.keras.Sequential([ sparsity.prune_low_magnitude( tf.keras.layers.Dense(10, activation='relu', kernel_initializer='glorot_uniform'), **pruning_params), tfp.layers.DistributionLambda(lambda t: tfd.Normal(loc=t, scale=1)) ]) negloglik = lambda x, rv_x: -rv_x.log_prob(x) callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir='D:\Python\logs2', profile_batch=0) ] pruned_model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.0001), loss=negloglik)
def run(network_type, batch_size, epochs, pruning_epochs, tensorboard_log_dir, model_dir_path, _log, _run, num_iters): assert network_type in ('mlp', 'cnn') _log.info('Emptying model directory...') if model_dir_path.exists(): shutil.rmtree(model_dir_path) Path(model_dir_path).mkdir(parents=True) _log.info('Loading data...') (X_train, y_train), (X_test, y_test) = load_data() metrics = {} unpruned_model_path, _ = get_two_model_paths() pruned_model_paths = get_pruned_model_paths() init_model_path = get_init_model_path() unpruned_model = create_model() save_weights(unpruned_model, init_model_path) _log.info('Training unpruned model...') metrics['unpruned'] = train_model(unpruned_model, X_train, y_train, X_test, y_test, unpruned_model_path, epochs=epochs) _log.info('Unpruned model sparsity: {}'.format( get_sparsity(unpruned_model))) save_weights(unpruned_model, unpruned_model_path) begin_step = np.ceil(X_train.shape[0] / batch_size).astype( np.int32) * (pruning_epochs - 1) for i in range(num_iters): pruned_model_path = pruned_model_paths[i] sparsity_ratio = 1.0 - 0.8**(i + 1) pruning_params = get_pruning_params(X_train.shape[0], initial_sparsity=sparsity_ratio / 2.0, final_sparsity=sparsity_ratio, begin_step=begin_step) pruned_model = sparsity.prune_low_magnitude(unpruned_model, **pruning_params) pruning_callbacks = [ sparsity.UpdatePruningStep(), sparsity.PruningSummaries(log_dir=tensorboard_log_dir, profile_batch=0) ] _log.info('Training pruned model...') metrics[f'iter{i+1}'] = train_model(pruned_model, X_train, y_train, X_test, y_test, pruned_model_path, epochs=pruning_epochs, callbacks=pruning_callbacks) _log.info('Pruned model sparsity: {}'.format( get_sparsity(pruned_model))) save_weights(pruned_model, pruned_model_path) ex.add_source_file(__file__) with open(model_dir_path / 'metrics.json', 'w') as f: json.dump(metrics, f, cls=NumpyEncoder) return metrics