Reducing Keras.models.model memory consumption

user_5101551 0
Viewed 417

I am exploring deploying a MaskRCNN model (https://github.com/matterport/Mask_RCNN) to a mobile device. The model is built on tensorflow-1.x and Keras-2, having some custom code but mostly rely on keras.models.model as its core.

My major concern is the memory requirement: with batch size = 1, the program gets killed during loading the weights, after tensorflow allocator sending "allocation over 10% of system memory" warnings. Is there any dirty and simple ways that I can compress the model size with minimal change to the source code? Like, setting the tf/keras default data type to float16 globally, or drop near-zero weights.

keras tensorflow python
1 Answers

Yes. Check it out: Pruning in Keras example

The idea is relatively simple. That is, in the training process, after a certain number of iterations, some weights are close to 0 in the net, then just set them to be 0. Repeat until Net parameters reach the target sparsity. In this way, after the model training is completed, many weights in the model are 0. Then, when we use zip for compression. In the inference process, a lot of calculations are reduced.

The following is an example:

import tensorflow_model_optimization as tfmot

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

# Compute end step to finish pruning after 2 epochs.
batch_size = 128
epochs = 2
validation_split = 0.1 # 10% of training set will be used for validation set. 

num_images = train_images.shape[0] * (1 - validation_split)
end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs

# Define model for pruning.
pruning_params = {
      'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.50,
                                                               final_sparsity=0.80,
                                                               begin_step=0,
                                                               end_step=end_step)
}

model_for_pruning = prune_low_magnitude(model, **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.summary()

Plus, just small tips. Try to compress the model gradually. To preserve a better performance.

user_128384030
Related