python - 具有 None 类型形状张量的 TensorFlow 2.0 层

Question

我尝试在 tf.keras.layers.Lambda() 中调用下面的函数遵循 TF 2.0。 inputs和outputs张量将是具有 3 个颜色 channel 的相同维度的两个图像。我的目标是从 outputs 中提取一个掩码张量，将其应用到 inputs张量，然后返回结果张量。展平张量的动机是由于 tf.tensor_scatter_nd_update() 的限制功能。当我构建模型时，它无法初始化 updates自 indices.shape[0]是 None值(value)。如果我用两个 tf.constant() 在模型外部调用该层要初始化的张量x ，它在急切执行中运行得非常好(因为x张量已经定义了值)。不幸的是，当我用 tf.keras.layers.Lambda() 调用这个函数时我收到以下错误:

TypeError: can't multiply sequence by non-int of type 'NoneType'

@tf.function
def applyMask(x):
  # Extract Tensors
  inputs = x[0]
  outputs = x[1]

  # Flatten the Outputs Tensor and Extract Mask Indices 
  outputs = tf.reshape(outputs,(tf.size(outputs),))
  indices = tf.where(outputs==1.)
  indices = tf.cast(indices, tf.int32)

  # Construct Updates Tensor from Mask Indices
  updates = tf.constant([1.]*indices.shape[0])

  # Flatten Input Tensor and Apply Mask
  out_dim = inputs.shape
  inputs = tf.reshape(inputs,(tf.size(inputs),))
  tensor = tf.tensor_scatter_nd_update(inputs, indices, updates)

  # Reconstruct Input Into Tensor
  tensor = tf.reshape(tensor, out_dim)
  return tensor

Answer 1

没必要这么复杂。简单地做，

inp1 = Input(shape=(None, None, 3)) # Inputs
inp2 = Input(shape=(None, None, 3)) # Outputs

out = Lambda(lambda x: tf.where(tf.equal(x[1], 1), x[1], x[0]))([inp1, inp2])

您甚至可以设置高度和宽度 无，只要将并行样本传递给inp1并且inp2 完全相同(形状方面)，tf.where 可以正常工作。