Python 向张量的一个维度添加偏移量，该偏移量是另一个张量的输入

我正在寻找一种有效的方法，用TensorFlow对下面的示例进行编码。下面，它是用numpy以最愚蠢的方式编码的。
其思想是使用另一个张量值在高度维度上偏移大小张量（

batch_size

，

height

，

width

，

channels

）。换句话说：

tensor_2[i, j, k, l] = tensor_1[i, j + tensor_offset[i, j, k, l], k, l]

以下是我使用的代码：

import numpy as np 
import time

begin = time.time()
b, h, w ,c = 5, 256, 512, 20
offset = np.random.rand(b, h, w , c).astype(int)
image = np.ones((b, h, w ,c))
label = np.ones((b, h, w ,c))
label_offset = np.zeros ((b, h, w,c ))

loss = 0
count = 0

for i in range(b):
    for j in range (h):
        for k in range (w):
            for l in range (c): 
                offset_ = j + np.int(offset [i,j,k,l])
                if offset_ > 255:
                    pass
                else:
                    label_offset[i,j,k,l] = label [i,offset_,k,l]
                    loss =+ label_offset[i,j,k,l]*np.log(image [i,j,k,l])
                    count=+1


loss = loss/count
end = time.time()
print ('duree:', end - begin)

你可以使用

索引[i，j，k，l]=[i，j+张量偏移量[i，j，k，l]，k，l]

（因此

索引的维度为5）

您可以这样构建它：

import numpy as np
import tensorflow as tf

b, h, w ,c = 2, 11, 13, 7
final_shape = [b, h, w, c]
offset = np.random.randint(low=0, high=h, size=final_shape)
image = np.random.randint(low=0, high=1000, size=final_shape)

input_tensor = tf.constant(image)

m1 = tf.transpose(tf.reshape(tf.tile(tf.range(b), multiples=[h * w * c]), [h, w, c, b]), perm=[3, 0, 1, 2])
m2 = tf.transpose(tf.reshape(tf.tile(tf.range(h), multiples=[b * w * c]), [b, w, c, h]), perm=[0, 3, 1, 2]) + offset
not_too_big = tf.less(m2, h)
m2_safe = tf.mod(m2, h)  # Makes sure we don't go too far in the original array
m3 = tf.transpose(tf.reshape(tf.tile(tf.range(w), multiples=[b * h * c]), [b, h, c, w]), perm=[0, 1, 3, 2])
m4 = tf.reshape(tf.tile(tf.range(c), multiples=[b * h * w]), [b, h, w, c])  # No transposition needed here

indices = tf.stack([m1, m2_safe, m3, m4], axis=-1)
tmp = tf.gather_nd(input_tensor, indices)
output = tf.multiply(tmp, tf.cast(not_too_big, tmp.dtype))  # Sets all the values corresponding to j+offset>h to 0

编辑：这适用于我的换位。
@gdelab
我对您的代码进行了以下改进：

def tensor_offset (input_tensor, offset_tensor, batch, nbcl):

    b, h, w ,c = batch, 256,512,nbcl

    m = tf.reshape(tf.tile(tf.range(b), multiples=[w*h*c]), [h,w,c, b])
    m1 = tf.reshape(tf.tile(tf.range(h), multiples=[w*b*c]), [b,w,c,h])
    m2 = tf.reshape(tf.tile(tf.range(w), multiples=[h*b*c]), [b,h,c,w])
    m2 = m2 +tf.transpose(tf.cast(offset_tensor,tf.int32),perm=[0, 1, 3, 2])
    m3 = tf.reshape(tf.tile(tf.range(c), multiples=[h*b*w]), [b,h,w,c])

    indices = tf.stack([tf.transpose(m,perm=[3,0,1,2]), tf.transpose(m1,perm=[0, 3, 1, 2]), tf.transpose(m2,perm=[0, 1, 3, 2]),m3], axis=-1)
    paddings = tf.constant([[0, 0], [0, 0], [0,100], [0,0]])
    output = tf.gather_nd(tf.pad(input_tensor, paddings), indices)

    return output    

非常感谢您的回答，但它给出了以下错误：对于输入形状为[5]、[4]的“Tile_4”（op:“Tile”），形状必须为秩1，但为秩4。我猜这是因为范围（b）和倍数的形状不一样。听起来不错，我没有正确使用瓷砖。它和编辑一起工作吗？谢谢，它工作起来很有魅力，但我有点困惑。使用numpy和for循环大约需要8秒钟，使用扁平化矩阵和列表需要7秒钟，使用tensorflow上的代码需要10秒钟。我希望tensorflow上的代码会快得多。对不起，我的gpu已经饱和了。需要3秒钟。非常感谢你救了我一天@它工作不正常吗？你为什么不接受我的答案？看来我们最终得到了相同的指数解决方案。填充是处理过大偏移的好方法！