Блог

Я пытаюсь реализовать простую модель capsnet в TF2.0.

На данный момент я добавил несколько слоев conv2d и слой reshape, но теперь мне нужно добавить функцию squash. Проблема в том, что tf.norm() отправит меня на NaN посадку, поскольку я сжимаю целые векторы, поэтому мне приходится использовать пользовательскую функцию squash. Я никогда раньше не писал пользовательский слой, и я в основном просто использовал шаблон из руководства и добавил математическую функцию под call() .

Поскольку я делаю все это внутри keras.models.Sequential модели, я не был уверен, как получить результат после первых нескольких слоев, поэтому я просто решил сделать функцию squash своим собственным слоем в модели. Я чувствую, что это, вероятно, совершенно неверно, поэтому я ищу информацию о наилучшем способе решения этой проблемы.

Должен ли я вообще использовать keras.Model для этого, или я должен использовать новую функцию eager execution, чтобы просто передавать тензоры через слои вручную? Если можно использовать 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()
    ])