如何在 Tensorflow 中冻结部分张量？答案

【问题标题】：How to freeze part of the tensor in Tensorflow?如何在 Tensorflow 中冻结部分张量？
【发布时间】：2021-08-30 23:48:27
【问题描述】：

我想在 Tensorflow 中实现剪枝算法。为了做到这一点，我不仅需要在修剪阶段掩盖权重，而且在随后的训练期间保持这个张量的一些条目冻结为零。

像我一样（[0, 0] 和 [1, 1] 条目被屏蔽）：

[[0 1.3 2], [1.34 0 2.3]]

在几批之后，我希望仍然有零位置[0, 0] 和[1, 1]。

这里有一个解决方案How to stop gradient for some entry of a tensor in tensorflow。但似乎它仅适用于 TensorFlow v1，因为在调用 fit 方法后更新了变量的屏蔽条目。

可以创建 Optimizer 的特殊类，该类将重新定义 apply_gradients 方法，其中梯度将在每次反向传递后乘以掩码，但这种解决方案似乎相当不方便，并且重新定义了多个优化器- Adam，RMSProp，随便。

【问题讨论】：

【解决方案1】：

Tensorflow Model Optimization 是您要找的吗？

它已经有 pruning 的实现，其工作方式类似于您描述的屏蔽技术 (pruning_impl.py)。

【讨论】：

【解决方案2】：

也许，有一个更聪明的解决方案，但实际的方法是做到这一点，以及在tfmot 中如何做到这一点是将权重乘以在每次转发之后的 0 和 1 掩码通过。

假设一个有张量：

[[0 1.3 2], [1.34 0 2.3]]

需要保留[0, 0] 和[1, 1] 位置处的零。

然后，对于每个操作，都需要将此张量乘以以下掩码：

[[0 1 1], [1 0 1]]

这将完成这项工作，因为渐变不会流过这些条目。

【讨论】：