如何在张量流中按行旋转张量?

时间:2019-09-30 18:35:33

标签: python tensorflow

我希望能够以类似于collections.deque.rotate的方式旋转张量。能够分批广播将是理想的。例如,给定旋转数组foo,我想通过bar中的每一行来旋转foo的每一行。

例如

def tf_rotate(to_rotate, rotate_by):
    pass #How to do this?

with tf.Session('') as sesh:
    foo = tf.constant([1,2,3])
    bar = tf.constant(
        [[1,  2,  3],
         [4,  5,  6],
         [7,  8,  9]])
    print(
        sesh.run(
            tf_rotate(bar, foo)
    ))

应该给出:

[[3  1  2]
 [5  6  4]
 [7  8  9]]

此操作作为for循环非常简单,但是在tensorflow中实现算法时,这些事情变得很复杂。我们如何在张量流中对数组/张量执行circular permutation

1 个答案:

答案 0 :(得分:1)

可能是这样的:

def rotate(matrix, shifts):
    """"requested rotate function - assumes matrix shape is mxn and shifts shape is m"""

    # get shape of the input matrix
    shape = tf.shape(matrix)

    # compute and stack the meshgrid to get the index matrix of shape (2,m,n)
    ind = tf.stack(tf.meshgrid(tf.range(shape[0]), tf.range(shape[1]), indexing='ij'))

    # reshape it to (m,n,2)
    ind = tf.transpose(ind, [1,2,0])

    # add the value from shifts to the corresponding row and devide modulo shape[1]
    # this will effectively introduce the desired shift, but at the level of indices
    shifted_ind = tf.mod(tf.transpose(tf.transpose(ind[:,:,1]) + shifts), shape[1])

    # convert the shifted indices to the right shape
    new_ind = tf.transpose(tf.stack([ind[:,:,0], shifted_ind]) , [1,2,0]) 

    # return the resliced tensor
    return tf.gather_nd(matrix, new_ind)

示例:

In [36]: matrix = np.tile(10*np.arange(10), [7,1])                                                                                        

In [37]: matrix                                                                                                                           
Out[37]: 
array([[ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90],
       [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90]])

In [38]: shifts = np.random.randint(low=0, high=matrix.shape[1], size=matrix.shape[0], dtype=int32)                                       

In [39]: shifts                                                                                                                           
Out[39]: array([5, 7, 6, 4, 5, 6, 6], dtype=int32)

In [40]: res = rotate(tf.constant(matrix), tf.constant(shifts))                                                                           

In [41]: sess.run(res)                                                                                                                    
Out[41]: 
array([[50, 60, 70, 80, 90,  0, 10, 20, 30, 40],
       [70, 80, 90,  0, 10, 20, 30, 40, 50, 60],
       [60, 70, 80, 90,  0, 10, 20, 30, 40, 50],
       [40, 50, 60, 70, 80, 90,  0, 10, 20, 30],
       [50, 60, 70, 80, 90,  0, 10, 20, 30, 40],
       [60, 70, 80, 90,  0, 10, 20, 30, 40, 50],
       [60, 70, 80, 90,  0, 10, 20, 30, 40, 50]])

不确定它是否比循环更好。