在 Keras 层中使用自定义操作并使张量可训练

Use custom operation in Keras layer and make tensor trainable

我有如下操作:

def custom_operation(input):
    kernel = tf.constant([3, 4, 5], tf.float32)
    frames = tf.signal.frame(input, 3, 1)
    return tf.reduce_sum(tf.abs(frames - tf.reshape(kernel, (1, 3))), axis=-1)

custom_operation(tf.constant([1, 2, 3, 4, 5, 6, 7], tf.float32))
# <tf.Tensor: shape=(5,), dtype=float32, numpy=array([6., 3., 0., 3., 6.], dtype=float32)>

我想在 Keras 自定义层中使用它,其中 input 是层的输入,kernel 是具有可训练值的张量,而不是硬编码的 [3, 4, 5].

从 Keras 调整 Conv1D 层以调用 custom_operation 而不是 tf.nn.conv1d 似乎并不难,但我不知道如何制作 kernel 可训练。

这个怎么样:

import tensorflow as tf
from tensorflow.keras.layers import Layer


class CustomLayer(Layer):
  """``CustomLayer``."""
  def __init__(self, name="custom"):
    super().__init__(name=name)
    
  def build(self, input_shape):
    self.w = self.add_weight(
        shape=(1, 3),
        initializer="random_normal",
        trainable=True)
    
  def call(self, x):
    frames = tf.signal.frame(x, 3, 1)
    return tf.math.reduce_sum(tf.math.abs(frames - self.w), axis=-1)

测试图层。

x = tf.constant([1, 2, 3, 4, 5, 6, 7], tf.float32)

CustomLayer()(x)

# <tf.Tensor: shape=(5,),
# dtype=float32,
# numpy= array([ 6.0877113,  9.087711 , 12.087711 , 15.087711 , 18.087711 ],
# dtype=float32)>