1

我正在尝试在 TF2.0 中实现一个简单的 capsnet 模型。

到目前为止，我已经添加了一些 conv2d 层和一个 reshape 层，但我现在需要添加一个 squash 函数。问题是，由于我正在压缩整个向量，所以tf.norm()我NaN必须使用自定义的 squash 函数。我以前从未编写过自定义层，我基本上只是使用教程中的模板并在call().

由于我是在模型中完成这一切的keras.models.Sequential，所以我不确定如何在前几层之后获得输出，所以我决定让 squash 函数在模型中成为它自己的层。我觉得这可能是完全错误的，所以我正在寻找一些关于解决这个问题的最佳方法的意见。

我应该keras.Model为此使用 a ，还是应该使用新的急切执行功能来手动将张量通过层？如果可以使用SquashLayer()我已经实现的，那么我应该通过什么作为参数，以便获得正确的输出以传递到下一层？

class SquashLayer(tf.keras.layers.Layer):
    def __init__(self, output_units):
        super(SquashLayer, self).__init__()
        self.output_units = output_units

    def build(self, input_shape):
        self.kernel = self.add_variable(
          'kernel', [input_shape[-1], self.output_units])

    def call(self, input):
        squared_norm = tf.reduce_sum(tf.square(input), axis=-1, keepdims=True)
        safe_norm = tf.sqrt(squared_norm + 1e-7)
        squash_factor = squared_norm / (1. + squared_norm)
        unit_vector = input / safe_norm
        return squash_factor * unit_vector

model = keras.models.Sequential([
    keras.layers.InputLayer(input_shape=(28, 28, 1)),
    keras.layers.Conv2D(filters=256, kernel_size=9, strides=1, padding='valid', activation=tf.nn.relu, name='conv1'),
    keras.layers.Conv2D(filters=256, kernel_size=9, strides=2, padding='valid', activation=tf.nn.relu, name='conv2'),
    keras.layers.Reshape((-1, caps1_n_caps, caps1_n_dims)),
    SquashLayer()
    ])
4

0 回答 0