Tensorflow：如何在不包装的情况下使用 tf.roll？

Question

我想独立移动二维张量的列或行，例如：

a = tf.constant([[1,2,3], [4,5,6]])
shift = tf.constant([2, -1])
b = shift_fn(a, shift)

这给了我：

b = [[0, 0, 1], [5, 6, 0]]

我发现tf.roll()可以做类似的事情，但会包装元素。 我如何使用它来填充零？

Answer 1

一个不太好的解决方案是首先使用tf.pad填充张量，然后在tf.map_fn tf.roll独立移动填充张量的每一行（或列）。 最后，您可以从结果中提取适当的部分。 例如：

a = tf.constant([[1,2,3], [4,5,6]])
shift = tf.constant([2, -1])

cols = a.shape[1]
paddings = tf.constant([[0, 0], [cols, cols]])
padded_a = tf.pad(a, paddings)

tf.map_fn(
    lambda x: tf.roll(x[0], x[1], axis=-1),
    elems=(padded_a, shift),
    # The following argument is required here
    fn_output_signature=a.dtype
)[:,cols:cols*2]

"""
<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
array([[0, 0, 1],
       [5, 6, 0]], dtype=int32)>
"""

或者，为了帮助节省一些 memory，填充和切片都可以在传递给tf.map_fn的fn function 内完成。

Answer 2

@tf.function
def shift(inputs, shift, axes, pad_val = 0):
    assert len(shift) == len(axes)
    axis_shift = zip(axes, shift)
    axis2shift = dict(axis_shift)  
    old_shape = inputs.shape

    for axis in axis2shift:
        pad_shape = list(inputs.shape)
        pad_shape[axis] = abs(axis2shift[axis])
        input_pad = tf.fill(pad_shape, pad_val)
        inputs = tf.concat((inputs, input_pad), axis = axis) 
    
    
    input_roll = tf.roll(inputs, shift, axes)
    ret = tf.slice(input_roll, [0 for _ in range(len(old_shape))], old_shape)

    return ret

上面的 function 实现了 shift with tf.roll like interface*（shift 和 axes 应该是列表而不是张量）。 正如@today 在之前的回答中所说，我执行了 pad-roll-slice 操作。

希望这可以帮助