在 Keras 自定义层中获取批量大小并使用张量流操作（tf.Variable）

Question

我想编写一个带有张量流操作的 Keras 自定义层，需要批量大小作为输入。显然我在每个角落都在挣扎。

假设一个非常简单的层：（1）获取批量大小（2）根据批量大小创建一个 tf.Variable（我们称之为 my_var），然后使用一些 tf.random 操作来改变 my_var（3）最后，返回输入乘以我的_var

到目前为止我尝试过的：

class TestLayer(Layer):

    def __init__(self, **kwargs):

        self.num_batch = None
        self.my_var = None

        super(TestLayer, self).__init__(**kwargs)

    def build(self, input_shape):

        self.batch_size = input_shape[0]

        var_init = tf.ones(self.batch_size, dtype = x.dtype)
        self.my_var = tf.Variable(var_init, trainable=False, validate_shape=False)

        # some tensorflow random operations to alter self.my_var

        super(TestLayer, self).build(input_shape)  # Be sure to call this at the end

    def call(self, x):

        return self.my_var * x

    def compute_output_shape(self, input_shape):

        return input_shape

现在创建一个非常简单的模型：

# define model
input_layer = Input(shape = (2, 2, 3), name = 'input_layer')
x = TestLayer()(input_layer)

# connect model
my_mod = Model(inputs = input_layer, outputs = x)
my_mod.summary()

不幸的是，无论我在代码中尝试/更改什么，我都会遇到多个错误，其中大多数都带有非常神秘的回溯（ValueError：无法将部分已知的 TensorShape 转换为 Tensor：或 ValueError：不支持 None 值。）。

有什么一般性建议吗？提前致谢。

Answer 1

如果要创建 size 变量，则需要指定批量大小batch_size。此外，如果要打印摘要，则摘要tf.Variable必须具有固定形状 ( validatate_shape=True)，并且必须可广播才能成功乘以输入：

import tensorflow as tf
from tensorflow.keras.layers import Layer, Input
from tensorflow.keras.models import Model

class TestLayer(Layer):

    def __init__(self, **kwargs):
        self.num_batch = None
        self.my_var = None
        super(TestLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.batch_size = input_shape[0]
        var_init = tf.ones(self.batch_size, dtype=tf.float32)[..., None, None, None]
        self.my_var = tf.Variable(var_init, trainable=False, validate_shape=True)
        super(TestLayer, self).build(input_shape)  # Be sure to call this at the end

    def call(self, x):
        res = self.my_var * x
        return res

    def compute_output_shape(self, input_shape):
        return input_shape

# define model
input_layer = Input(shape=(2, 2, 3), name='input_layer', batch_size=10)
x = TestLayer()(input_layer)

# connect model
my_mod = Model(inputs=input_layer, outputs=x)
my_mod.summary()

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_layer (InputLayer)     (10, 2, 2, 3)             0         
_________________________________________________________________
test_layer (TestLayer)       (10, 2, 2, 3)             0         
=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0