如何实现神经网络剪枝?

How to implement neural network pruning?

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

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

如果您添加掩码,则只有一部分权重会参与计算,因此您的模型将被修剪。例如,自回归模型使用掩码来屏蔽掉引用未来数据的权重,以便时间步长 t 的输出仅取决于时间步长 0, 1, ..., t-1.

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

根据评论中的讨论,这是一种修剪神经网络层(权重矩阵)的方法。该方法本质上所做的是 selects k% 基于它们的范数的最小权重(矩阵的元素),并将它们设置为零。这样,相应的矩阵可以被视为稀疏矩阵,我们可以执行 dense-sparse 矩阵乘法,如果修剪了足够的权重,它会更快。

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

虽然上面的方法修剪单个连接(权重),但下面的方法从权重矩阵修剪整个神经元。即,方法 select 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 数据集上进行实验。