如何实现神经网络剪枝？

Posted 2023-02-16

【中文标题】如何实现神经网络剪枝？

【英文标题】：How to implement neural network pruning?

【发布时间】：2019-10-11 10:24:41

【问题描述】：

我在 keras 中训练了一个模型，我正在考虑修剪我的全连接网络。我对如何修剪图层有点迷茫。

“学习权重和连接以提高效率”的作者 神经网络，说他​​们为层的阈值权重添加了一个掩码。我可以尝试做同样的事情并微调训练好的模型。但是，它如何减少模型大小和计算量？

【问题讨论】：

具体来说，你想知道如何修剪神经网络中的特定权重吗？例如，给定一个W 矩阵，您想将某些元素设置为 0？

@gorjan 我的目标是减小最终模型的大小并加快推理速度。我不确定设置W 的某些值是否会减小模型大小。我需要移除连接。据我了解，TensorRT 和 TensorFlow Lite 是这样做的吗？

您基本上不能“删除”权重。您可以做的是将某些权重设置为 0，然后将矩阵威胁为稀疏矩阵。然后，TF 对密集-稀疏/稀疏-稀疏矩阵乘法有一些小的支持，可用于加速推理。这是一个相关的***线程：***.com/questions/44859321/…

@gorjan 有道理。我认为还有比这更多的东西。让我尝试实现类似的东西。

当然！作为答案，我将发布一个给定权重矩阵w: tf.Variable 和k: int 的方法，它将根据它们的规范删除k% 最小权重（矩阵中的元素）。

【参考方案1】：

根据 cmets 中的讨论，这是一种修剪神经网络层（权重矩阵）的方法。该方法的本质是根据它们的范数选择k% 最小的权重（矩阵的元素），并将它们设置为零。这样，对应的矩阵就可以看成一个稀疏矩阵，我们可以进行密集-稀疏矩阵乘法，如果剪掉足够多的权重会更快。

def weight_pruning(w: tf.Variable, k: float) -> tf.Variable:
    """Performs pruning on a weight matrix w in the following way:

    - The absolute value of all elements in the weight matrix are computed.
    - The indices of the smallest k% elements based on their absolute values are selected.
    - All elements with the matching indices are set to 0.

    Args:
        w: The weight matrix.
        k: The percentage of values (units) that should be pruned from the matrix.

    Returns:
        The unit pruned weight matrix.

    """
    k = tf.cast(tf.round(tf.size(w, out_type=tf.float32) * tf.constant(k)), dtype=tf.int32)
    w_reshaped = tf.reshape(w, [-1])
    _, indices = tf.nn.top_k(tf.negative(tf.abs(w_reshaped)), k, sorted=True, name=None)
    mask = tf.scatter_nd_update(tf.Variable(tf.ones_like(w_reshaped, dtype=tf.float32), name="mask", trainable=False), tf.reshape(indices, [-1, 1]), tf.zeros([k], tf.float32))

    return w.assign(tf.reshape(w_reshaped * mask, tf.shape(w)))

虽然上面的方法会修剪单个连接（权重），但下面的方法会从权重矩阵中修剪整个神经元。即，该方法根据欧几里得范数选择k%最小的神经元（权重矩阵的列），并将它们设置为零。

def unit_pruning(w: tf.Variable, k: float) -> tf.Variable:
    """Performs pruning on a weight matrix w in the following way:

    - The euclidean norm of each column is computed.
    - The indices of smallest k% columns based on their euclidean norms are selected.
    - All elements in the columns that have the matching indices are set to 0.

    Args:
        w: The weight matrix.
        k: The percentage of columns that should be pruned from the matrix.

    Returns:
        The weight pruned weight matrix.

    """
    k = tf.cast(
        tf.round(tf.cast(tf.shape(w)[1], tf.float32) * tf.constant(k)), dtype=tf.int32
    )
    norm = tf.norm(w, axis=0)
    row_indices = tf.tile(tf.range(tf.shape(w)[0]), [k])
    _, col_indices = tf.nn.top_k(tf.negative(norm), k, sorted=True, name=None)
    col_indices = tf.reshape(
        tf.tile(tf.reshape(col_indices, [-1, 1]), [1, tf.shape(w)[0]]), [-1]
    )
    indices = tf.stack([row_indices, col_indices], axis=1)

    return w.assign(
        tf.scatter_nd_update(w, indices, tf.zeros(tf.shape(w)[0] * k, tf.float32))
    )

最后，这个Github repository 经历了这里解释的剪枝方法，并在 MNIST 数据集上进行了实验。

【参考方案2】：

如果您添加掩码，那么只有权重的一个子集将有助于计算，因此您的模型将被修剪。例如，自回归模型使用掩码来屏蔽引用未来数据的权重，以便时间步 t 的输出仅取决于时间步 0, 1, ..., t-1。

在您的情况下，由于您有一个简单的全连接层，因此最好使用 dropout。它在每个迭代步骤中随机关闭一些神经元，从而降低计算复杂度。然而，发明 dropout 的主要原因是为了解决过拟合问题：通过随机关闭一些神经元，可以减少神经元的相互依赖性，即避免一些神经元依赖其他神经元。此外，在每次迭代中，您的模型都会不同（不同数量的活动神经元和它们之间的不同连接），因此您的最终模型可以解释为几个不同模型的集合（集合），每个模型都专门（我们希望）在理解输入空间的特定子集。

【讨论】：

是的。但是，我的目标是加快推理速度并减小模型大小。如果我确实使用了遮罩，我仍然会存储所有层的权重，并且我仍然会执行整个 W.X + b（其中一些 W_ij 设置为 0。）

如果您的任务是减小模型大小，那么您无法使用动态掩码来实现此目的。如果掩码是静态的，则只需删除您不想学习的权重。您的网络将变得更加稀疏。

使用掩码确实可以加快计算速度。考虑一个过滤掉矩阵W 的前 3 列的掩码。然后，您可以将其实现为W[:, 3:]。这样，计算将仅在矩阵的剩余部分上进行。对于更复杂的掩码（不是连续的 ecc），您仍然可以获得一些优势，因为不会为等于 0 的权重计算梯度

但是，面具背后的原因通常不是为了加快训练速度。