如何使用批量大小在自定义 TensorFlow 层中创建张量

Question

我正在创建一个自定义 TF 层，我需要在其中创建一个像这样的张量

class MyLayer(Layer):
  def __init__(self, config, **kwargs):
    super(MyLayer, self).__init__(**kwargs)
    ....

  def call(self, x):
    B, T, C = x.shape.as_list()
    ...
    ones = tf.ones((B, T, C))
    ...
    # output projection
    y = ...
    return y

现在的问题是评估层时 B（批量大小）为 None，这导致 tf.ones 失败并出现以下错误：


ValueError: in user code:

    <ipython-input-69-f3322a54c05c>:29 call  *
        ones = tf.ones((B, T, C))
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py:201 wrapper  **
        return target(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/array_ops.py:3080 ones
        shape = ops.convert_to_tensor(shape, dtype=dtypes.int32)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/profiler/trace.py:163 wrapped
        return func(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py:1535 convert_to_tensor
        ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/constant_op.py:356 _tensor_shape_tensor_conversion_function
        "Cannot convert a partially known TensorShape to a Tensor: %s" % s)

    ValueError: Cannot convert a partially known TensorShape to a Tensor: (None, 8, 128)

我怎样才能让它工作？

Answer 1

如果你只想得到一个与x形状相同的张量，那么你可以使用tf.ones_like。像这样：

class MyLayer(Layer):

  ....

  def call(self, x):
    ones = tf.ones_like(x)

    ...

    # output projection
    y = ...
    return y

直到运行时才需要知道 x 的形状。

然而，一般来说，我们可能需要在运行前知道输入的形状，在这种情况下，我们可以在我们的层中实现 build() 方法，该方法将 input_shape 作为参数，并且是当我们编译我们的模型时调用。

从文档中复制的示例 here:

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__init__()
        self.units = units

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

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

如何使用批量大小在自定义 TensorFlow 层中创建张量

How to use batch size to create a tensor within a custom TensorFlow Layer

python

tensorflow

keras-layer

tf.keras

tensorflow2.0