减少 Keras.models.model 内存消耗

时间:2021-02-28 02:53:17

标签: python tensorflow keras

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

我主要关心的是内存需求:当批量大小为 1 时,在 tensorflow 分配器发送“分配超过系统内存的 10%”警告后,程序在加载权重时被终止。是否有任何肮脏和简单的方法可以在对源代码进行最少更改的情况下压缩模型大小?例如,将 tf/keras 默认数据类型全局设置为 float16,或者降低接近零的权重。

1 个答案:

答案 0 :(得分:0)

是的。看看: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()

另外,只是一些小技巧。尝试逐渐压缩模型。为了保持更好的性能。