Keras 中的自定义图层 returns NaN 作为渐变。造成这种情况的潜在问题有哪些？

Question

我从事一个项目，我们尝试从几何基元重建二维图像。为此，我开发了一个自定义 Keras 层，它输出一个给定几何特征的圆锥体图像。

它的输入是一个形状为batch_size * 5的张量，其中五个数字分别是圆锥顶点的xy坐标，描述圆锥轴线的单位向量的xy坐标，以及圆锥体顶部的角度。

目标是将该层用作编码器-解码器架构中的不可训练解码器。然后我们将向神经网络提供锥体图像。预期的行为是神经网络然后应该学习类似于上述的潜在表示。

当我将这一层合并到一个更大的网络中并尝试对其进行优化时，总会有一些权重最终被更新为 NaN。即使网络像没有激活函数的双神经元隐藏层一样简单，也会发生这种情况。

我已经彻底测试了我的图层。它的输出与我的预期一致。我在实现中找不到任何微不足道的错误（但你应该被警告我对 tensorflow 和 keras 还很陌生）。我已将问题缩小到图层的自动微分。

梯度似乎等于 0.0 或 NaN。我的理解是一些数值不稳定性导致梯度发散。

问题是双重的：

这里的根本原因是什么？
我该如何解决？

下面是一个最小的工作示例，显示梯度如何最终达到 0.0 或特定值的 NaN。

import numpy as np
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Layer
import tensorflow as tf
import numpy.random as rnd

class Cones(Layer):

    def __init__(self, output_dim, **kwargs):
        super(Cones, self).__init__(**kwargs)
        self.output_dim = output_dim
        coordinates = np.zeros((self.output_dim, self.output_dim, 2))
        for i in range(self.output_dim):
           for j in range(self.output_dim):
              coordinates[i,j,:] = np.array([i,j])

        coordinates = K.constant(coordinates)
        self.coordinates = tf.Variable(initial_value=coordinates, trainable=False)
        self.smooth_sign_width = tf.Variable(initial_value=output_dim, dtype=tf.float32, trainable=False)
        self.grid_width = tf.Variable(initial_value=output_dim, dtype=tf.float32, trainable=False)


    def build(self, input_shape):
        super(Cones, self).build(input_shape)

    def call(self, x):
        center = self.grid_width*x[:,:2]
        center = K.expand_dims(center, axis=1)
        center = K.expand_dims(center, axis=1)

        direction = x[:,2:4]
        direction = K.expand_dims(direction,1)
        direction = K.expand_dims(direction,1)
        direction = K.l2_normalize(direction, axis=-1)

        aperture = np.pi*x[:,4:]
        aperture = K.expand_dims(aperture)

        u = self.coordinates - center
        u = K.l2_normalize(u, axis=-1)

        angle = K.sum(u*direction, axis=-1)
        angle = K.minimum(angle, K.ones_like(angle))
        angle = K.maximum(angle, -K.ones_like(angle))

        angle = tf.math.acos(angle)


        output = self.smooth_sign(aperture-angle)

        output = K.expand_dims(output, -1)
        return output

    def smooth_sign(self, x):
        return tf.math.sigmoid(self.smooth_sign_width*x)


    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.output_dim, self.output_dim, 1)

geom = K.constant([[0.34015268, 0.31530404, -0.6827047, 0.7306944, 0.8521315]])
image = Cones(Nx)(geom)

x0 = geom
y0 = image

with tf.GradientTape() as t:
    t.watch(x0)
    cone = Cones(Nx)(x0)
    error = cone-y0
    error_squared = error*error
    mse = tf.math.reduce_mean(error_squared)

print(t.gradient(mse, x0))

geom = K.constant([[0.742021, 0.25431857, 0.90899783, 0.4168009, 0.58542883]])
image = Cones(Nx)(geom)

x0 = geom
y0 = image

with tf.GradientTape() as t:
    t.watch(x0)
    cone = Cones(Nx)(x0)
    error = cone-y0
    error_squared = error*error
    mse = tf.math.reduce_mean(error_squared)

print(t.gradient(mse, x0))

Answer 1

首先，我回答我自己的问题并把它留在那里，以防将来对某人有所帮助。我不知道这是否是 Whosebug 上公认的礼仪。

通过评论调用函数的后续步骤，我发现问题出在 tf.math.acos。在上面的代码中，我已经遇到了 acos 的问题，这导致我将输入的值剪裁在 -1 和 1 之间。数值问题意味着有时两个单位向量的点积超出了这个范围，其中 acos 被定义。但是，通过这样做，我最终在 1 和 -1 处评估了 acos，它不可微分，因此梯度中的 NaN。

为了解决这个问题，我首先改变了我的方法来计算两个向量之间的角度，使用 this scicomp stack exchange answer。然后，我剪裁了我执行计算的范围以避免 sqrt 在 0 处的不可微性。更准确地说，每当我有 c > 1.95 时，我将角度四舍五入到 pi，并且每当我有 c < 0.05 时，我将角度四舍五入为 0。

Keras 中的自定义图层 returns NaN 作为渐变。造成这种情况的潜在问题有哪些？

A custom layer in Keras returns NaN as a gradient. What are some potential issues causing this?

gradient

neural-network

keras

tensorflow

tensorflow2.0