-
Notifications
You must be signed in to change notification settings - Fork 14
/
model_pruning_example.py
94 lines (73 loc) · 3.64 KB
/
model_pruning_example.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import shutil
from pathlib import Path
import numpy as np
from keras import datasets, utils, callbacks, optimizers, losses
from keras.preprocessing.image import ImageDataGenerator
import ridurre
from example.cifar_10_resnet import resnet
from ridurre import model_complexity
TRAIN_LOGS_FOLDER_PATH = Path("./train_logs")
if TRAIN_LOGS_FOLDER_PATH.is_dir():
shutil.rmtree(str(TRAIN_LOGS_FOLDER_PATH))
TRAIN_LOGS_FOLDER_PATH.mkdir()
# Creating ResNet50 model
model = resnet.resnet_v1(input_shape=(32, 32, 3), depth=20, num_classes=10)
def compile_model(my_model):
my_model.compile(optimizer=optimizers.Adam(lr=0.001), loss=losses.categorical_crossentropy, metrics=["accuracy"])
compile_model(model)
# Loading data
(x_train, y_train), (x_test, y_test) = datasets.cifar10.load_data()
# Data Transform
x_train = x_train.astype(np.float32) / 255.0
y_train = utils.to_categorical(y_train)
x_train_mean = np.mean(x_train, axis=0)
x_train -= x_train_mean
x_test = x_test.astype(np.float32) / 255.0
y_test = utils.to_categorical(y_test)
x_test -= x_train_mean
print("Train shape: X {0}, y: {1}".format(x_train.shape, y_train.shape))
print("Test shape: X {0}, y: {1}".format(x_test.shape, y_test.shape))
# Data Augmentation with Data Generator
data_generator = ImageDataGenerator(horizontal_flip=True, vertical_flip=True, rotation_range=20)
# Create callbacks
tensorboard_callback = callbacks.TensorBoard(log_dir=str(TRAIN_LOGS_FOLDER_PATH))
model_complexity_param = model_complexity.ModelParametersCallback(TRAIN_LOGS_FOLDER_PATH, verbose=1)
model_checkpoint_callback = callbacks.ModelCheckpoint(str(TRAIN_LOGS_FOLDER_PATH) + "/model_{epoch:02d}.h5",
save_best_only=False,
save_weights_only=False,
verbose=1)
callbacks = [tensorboard_callback, model_complexity_param, model_checkpoint_callback]
# Train the model
FIRST_TRAIN_EPOCHS = 20
BATCH_SIZE = 32
STEPS_PER_EPOCH = len(x_train) // BATCH_SIZE
model.fit_generator(data_generator.flow(x_train, y_train, BATCH_SIZE),
epochs=FIRST_TRAIN_EPOCHS,
validation_data=(x_test, y_test),
callbacks=callbacks,
steps_per_epoch=STEPS_PER_EPOCH)
# Prune the model
def finetune_model(my_model, initial_epoch, finetune_epochs):
my_model.fit_generator(data_generator.flow(x_train, y_train, BATCH_SIZE),
epochs=finetune_epochs,
validation_data=(x_test, y_test),
callbacks=callbacks,
initial_epoch=initial_epoch,
verbose=1,
steps_per_epoch=STEPS_PER_EPOCH)
pruning = ridurre.KMeansFilterPruning(0.9,
compile_model,
finetune_model,
nb_finetune_epochs=5,
maximum_pruning_percent=0.85,
maximum_prune_iterations=10,
nb_trained_for_epochs=FIRST_TRAIN_EPOCHS)
model, last_epoch_number = pruning.run_pruning(model)
# Train again for a reasonable number of epochs (no always necessary)
SECOND_TRAIN_EPOCHS = 20
model.fit_generator(data_generator.flow(x_train, y_train, BATCH_SIZE),
epochs=last_epoch_number + SECOND_TRAIN_EPOCHS,
validation_data=(x_test, y_test),
callbacks=callbacks,
steps_per_epoch=STEPS_PER_EPOCH,
initial_epoch=last_epoch_number)