python - 应用@tf.custom_gradient 后渐变不流动

Question

我正在实施Lei 等人。(2015) 的 Rationalizing Neural Predictions神经网络为NLP问题生成提取摘要。

该架构需要一个采样层，该层从文本输入中选择可用单词的子集。由于采样不可微分，我使用@tf.custom_gradient，将 0 的“人工”梯度分配给采样层。

问题是采样层之上的层不再训练（看起来网络下方的所有梯度流基本上都无效了）。

我的问题是如何规避这种影响并让网络忽略或跳过不可微层？

我尝试了一个硬 sigmoid 激活（例如alpha = 60）来近似选择 0 和 1 之间，但它还不够好（更不用说优雅了）。

这是具有自定义渐变的图层

@tf.custom_gradient
def sampling(x):
    def grad(dy):
        return x*0+0
    dist = tfd.Bernoulli(probs=x,dtype = tf.float32)  
    return dist.sample(),grad

所以，这就是发生的事情：

1. 神经网络第 1 层 # 不训练
2. 神经网络第 2 层 # 不训练
3. 具有自定义梯度的层 # 不应该训练
4. 下一层 # 训练好
5. 下一层 # 训练好
6. 下一层 # 训练好
7. 输出

那么我如何让 1）和 2）仍然训练？

标签： pythontensorflow