如何使用批量大小在自定義 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作為參數，並在我們編譯模型時調用。

從此處的文檔復制的示例：

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