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
