降低Keras.models.model内存消耗

Question

我正在探索将MaskRCNN模型(https://github.com/matterport/Mask_RCNN)部署到移动设备。该模型基于tensorflow-1.x和Keras-2构建，具有一些自定义代码，但主要依赖于keras.models.model作为其核心。

我主要关心的是内存需求:当批处理大小= 1时，在tensorflow分配器发送"allocation over 10% of system memory“警告后，程序在加载权重的过程中被终止。有没有什么既脏又简单的方法可以压缩模型的大小，而对源代码的改动最小？比如，将tf/keras默认数据类型全局设置为float16，或者丢弃接近于零的权重。

Answer 1

是。查看：Pruning in Keras example

这个想法相对简单。也就是说，在训练过程中，经过一定次数的迭代后，网络中的一些权重接近于0，然后将它们设置为0。重复，直到网络参数达到目标稀疏度。这样，在模型训练完成后，模型中的许多权重都为0。然后，当我们使用zip进行压缩时。在推理过程中，减少了大量的计算量。

下面是一个示例：

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()

另外，只是一些小费。尝试逐步压缩模型。以保持更好的性能。