python - TensorFlow 2 中自定义层的 call() 中的 tf.while_loop()

Question

我想编写一个自定义层，该层应用密集层，然后将一些指定的函数应用于该计算的输出。我想指定应用于列表中各个输出的函数，以便我可以轻松更改它们。

我正在尝试应用 tf.while_loop 内的函数，但我不知道如何访问和写入 dense_output_nodes 的各个元素。

dense_output_nodes[i] = ... 不起作用，因为它告诉我

TypeError: 'Tensor' object does not support item assignment

所以我之前尝试过tf.unstack，即下面的代码，但现在使用 hidden_​​1 = ArithmeticLayer(unit_types=['id', 'sin', 'cos'])(输入)，我收到错误

TypeError: list indices must be integers or slices, not Tensor

因为显然 TensorFlow 将 i 从 tf.constant 转换为 tf.Tensor。

到目前为止，我真的在努力寻找解决这个问题的方法。有什么方法可以让它发挥作用吗？ 或者我应该将整个 ArithmeticLayer 构建为密集层和应用自定义函数的 Lambda 层的组合？

class ArithmeticLayer(layers.Layer):
    # u = number of units
    
    def __init__(self, name=None, regularizer=None, unit_types=['id', 'sin', 'cos']):
        self.regularizer=regularizer
        super().__init__(name=name)
        self.u = len(unit_types)
        self.u_types = unit_types

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.u),
                                 initializer='random_normal',
                                 regularizer=self.regularizer,
                                 trainable=True)
        self.b = self.add_weight(shape=(self.u,),
                                 initializer='random_normal',
                                 regularizer=self.regularizer,
                                 trainable=True)


    def call(self, inputs):
        # get the output nodes of the dense layer as a list
        dense_output_nodes = tf.matmul(inputs, self.w) + self.b
        dense_output_list = tf.unstack(dense_output_nodes, axis=1)
        
        # apply the function units
        i = tf.constant(0)
        def c(i):
            return tf.less(i, self.u)
        def b(i):
            dense_output_list[i] = tf.cond(self.u_types[i] == 'sin',
                                            lambda: tf.math.sin(dense_output_list[i]),
                                            lambda: dense_output_list[i]
                                           )
            dense_output_list[i] = tf.cond(self.u_types[i] == 'cos',
                                            lambda: tf.math.cos(dense_output_list[i]),
                                            lambda: dense_output_list[i]
                                           )
            return (tf.add(i, 1), )
        [i] = tf.while_loop(c, b, [i])
        
        final_output_nodes = tf.stack(dense_output_list, axis=1)
        return final_output_nodes

感谢您的建议!

Answer 1

如果您想在批量样本中按列应用某些函数，那么使用 tf.tensor_scatter_nd_update 应该可以解决问题。以下是在急切执行和图形模式下工作的示例:

import tensorflow as tf

class ArithmeticLayer(tf.keras.layers.Layer):
    # u = number of units
    
    def __init__(self, name=None, regularizer=None, unit_types=['id', 'sin', 'cos']):
        self.regularizer=regularizer
        super().__init__(name=name)
        self.u_types = tf.constant(unit_types)
        self.u_shape = tf.shape(self.u_types)

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.u_shape[0]),
                                 initializer='random_normal',
                                 regularizer=self.regularizer,
                                 trainable=True)
        self.b = self.add_weight(shape=(self.u_shape[0],),
                                 initializer='random_normal',
                                 regularizer=self.regularizer,
                                 trainable=True)

    def call(self, inputs):
        dense_output_nodes = tf.matmul(inputs, self.w) + self.b
        d_shape = tf.shape(dense_output_nodes)
        i = tf.constant(0)
        c = lambda i, d: tf.less(i, self.u_shape[0])

        def b(i, d):
          d = tf.cond(unit_types[i] == 'sin', 
                lambda: tf.tensor_scatter_nd_update(d, tf.stack([tf.range(d_shape[0]), tf.repeat([i], d_shape[0])], axis=1), tf.math.sin(d[:, i])), 
                lambda: d)
          d = tf.cond(unit_types[i] == 'cos', 
                lambda: tf.tensor_scatter_nd_update(d, tf.stack([tf.range(d_shape[0]), tf.repeat([i], d_shape[0])], axis=1), tf.math.cos(d[:, i])), 
                lambda: d)
          return tf.add(i, 1), d
        _, dense_output_nodes = tf.while_loop(c, b, loop_vars=[i, dense_output_nodes])

        return dense_output_nodes
x = tf.random.normal((4, 3))
inputs = tf.keras.layers.Input((3,))
arithmetic = ArithmeticLayer()
outputs = arithmetic(inputs)
model = tf.keras.Model(inputs, outputs)
model.compile(optimizer='adam', loss='mse')
model.fit(x, tf.random.normal((4, 3)), batch_size=2)

2/2 [==============================] - 3s 11ms/step - loss: 1.4259
<keras.callbacks.History at 0x7fe50728c850>