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)
