生成与张量流中的卷积核输入等效的索引

Generate indices equivalent to convolution kernel inputs in tensorflow

我想在进入 tensorflow 中每个 2D 卷积核的数据中添加一些预处理。如果输入是大小为 (N, H, W, 5) 的 NHWC,并且卷积核的大小为 (3, 3),步长为 (1, 1) 并且填充“有效”(我将在稍后应用卷积,如前所述下面),我首先要获得大小为 (N, H-2, W-2, 5*9) 的张量,其中 dims 1 和 2 的索引表示卷积核的空间位置,dim 3 的大小5*9=45,每个索引对应一个将进入卷积的输入 cell/channel,所有内核输入都沿通道维度排列。获得此张量后,我想沿通道维度 (dim 3) 应用一些变换,然后 then 调用大小为 (1, 1) 且步幅为 (1, 1) 的内核的卷积).本质上,我想收集所有将首先进入每个卷积核的输入,对它们做一些处理,然后调用卷积(它将具有大小为 (1, 1) 的内核,因为我已经收集了所有内核输入通道轴)。我感兴趣的是 tensorflow 中是否有一个操作可以帮助我为这些 windows 生成索引,我可能稍后可以将其与 tf.gather_nd() 一起使用。欢迎任何其他 ideas/solutions。谢谢!

我在下面粘贴我的解决方案:

N = 32
H = 101
W = 101
C = 5

# input batch data
batch = tf.constant(
    value=range(N*H*W*C), 
    shape=[N, H, W, C], 
    dtype=tf.float32)

print(f"Batch shape: {batch.shape}")
print(f"Batch tensor size: {tf.math.reduce_prod(batch.shape).numpy()}")
print()

def prepare_conv_input(batch):

    i_to_stack = []
    for i in range(0, H-2, 2):

        j_to_stack = []
        for j in range(0, W-2, 2):

            inp = tf.slice(
                batch,
                begin=[0, i, j, 0],
                size=[N, 3, 3, C])

            j_to_stack.append(
                tf.reshape(
                    inp, 
                    (N, 3*3*C)))

        i_to_stack.append(
            tf.stack(j_to_stack, 1))

    return tf.stack(i_to_stack, 1)

result = prepare_conv_input(batch)

print(f"Prepared shape: {result.shape}")
print(f"Prepared tensor size: {tf.math.reduce_prod(result.shape).numpy()}")
print()

结果是:

Batch shape: (32, 101, 101, 5)
Batch tensor size: 1632160

Prepared shape: (32, 50, 50, 45)
Prepared tensor size: 3600000

由于相邻卷积之间的重叠,数据大小增加了。