TensorFlow:沿着一个特定的轴应用卷积

Question

我有两个形状为(batch_size, height, width, 1)的张量A和B，我想沿着width轴进行卷积，即A[0, 0]与B[0, 0]卷积，A[0, 1]与B[0, 1]卷积，A[3, 6]与B[3, 6]卷积，等等。我曾试图通过tf.nn.conv1d和tf.map_fn的组合来实现这一点，但我不断收到与输入形状、AutoGraph等相关的错误。

如何有效地将这两个张量沿特定轴进行卷积？

编辑：添加了非功能代码来说明的想法

# Create tensors
A, B = tf.random.normal(shape=(2, 1, 8, 512, 1))
# Reshape tensors suitable for `conv1d`
A = tf.transpose(A, [1, 0, 2, 3])
B = tf.transpose(B, [1, 2, 0, 3])
# Define convolution function
conv_fn = partial(tf.nn.conv1d, padding="SAME", stride=1)
# Apply map
AB = tf.map_fn(lambda x: conv_fn(x[0], x[1]), (A, B))

结果(为简洁起见，省略了矢量值(：

ValueError: The two structures don't have the same nested structure.
First structure: type=tuple str=(<tf.Tensor: shape=(8, 1, 512, 1), dtype=float32, numpy=
array([[omitted]], dtype=float32)>, <tf.Tensor: shape=(8, 512, 1, 1), dtype=float32, numpy=
array([[omitted]], dtype=float32)>)
Second structure: type=EagerTensor str=tf.Tensor(
[[omitted]], shape=(1, 512, 1), dtype=float32)
More specifically: Substructure "type=tuple str=(<tf.Tensor: shape=(8, 1, 512, 1), dtype=float32, numpy=
array([[omitted]], dtype=float32)>, <tf.Tensor: shape=(8, 512, 1, 1), dtype=float32, numpy=
array([[omitted]], dtype=float32)>)" is a sequence, while substructure "type=EagerTensor str=tf.Tensor(
[[omitted]], shape=(1, 512, 1), dtype=float32)" is not

Answer 1

检查有关张量形状的文档：

要执行的可调用项。它接受一个论点与elems相同(可能嵌套(的结构。其输出必须如果提供了与fn_output_signature相同的结构；否则，它必须具有与elems相同的结构。

所以尝试一个更灵活的输出签名：

import tensorflow as tf
from functools import partial
A, B = tf.random.normal(shape=(2, 1, 8, 512, 1))
# Reshape tensors suitable for `conv1d`
A = tf.transpose(A, [1, 0, 2, 3])
B = tf.transpose(B, [1, 2, 0, 3])
# Define convolution function
conv_fn = partial(tf.nn.conv1d, padding="SAME", stride=1)
# Apply map
AB = tf.map_fn(lambda x: conv_fn(x[0], x[1]), (A, B), fn_output_signature = tf.TensorSpec((None)))
print(AB.shape)
# (8, 1, 512, 1)

您也可以考虑使用tf.while_loop。