python之Tensorflow 高级索引 : Assign a smaller tensor into a bigger one into a position based on two index tensors

我想创建一个形状为 (3,4,2,2) 的零张量，并在两个 (3,1) 张量给定的位置插入一个 (3,4) 张量。

示例代码:对数组的等效 numpy 操作如下:

# Existing arrays of required shapes 
bbox = np.arange(3*4).reshape(3,4) 
x = np.array([0,0,1]) 
y = np.array([1,1,1]) 
 
# Create zeros array and assign into it 
output = np.zeros((3,4,2,2)) 
output[np.arange(3),:,x,y] = bbox 

如何使用 Tensorflow 做类似的事情？

注意:我实际上想使用大小为 (32,125,32,32) 的张量。以上是复现的简单代码

请您参考如下方法:

以下是如何使用 tf.scatter_nd 做到这一点:

import tensorflow as tf 
import numpy as np 
 
bbox = np.arange(3 * 4).reshape(3, 4) 
x = np.array([0, 0, 1]) 
y = np.array([1, 1, 1]) 
x_size = 2 
y_size = 2 
 
# TensorFlow calculation 
with tf.Graph().as_default(), tf.Session() as sess: 
    bbox_t = tf.convert_to_tensor(bbox) 
    x_t = tf.convert_to_tensor(x) 
    y_t = tf.convert_to_tensor(y) 
    shape = tf.shape(bbox_t) 
    rows, cols = shape[0], shape[1] 
    ii, jj = tf.meshgrid(tf.range(rows), tf.range(cols), indexing='ij') 
    xx = tf.tile(tf.expand_dims(x_t, 1), (1, cols)) 
    yy = tf.tile(tf.expand_dims(y_t, 1), (1, cols)) 
    idx = tf.stack([ii, jj, xx, yy], axis=-1) 
    output = tf.scatter_nd(idx, bbox_t, [rows, cols, x_size, y_size]) 
    output_tf = sess.run(output) 
 
# Test with NumPy calculation 
output_np = np.zeros((3, 4, 2, 2)) 
output_np[np.arange(3), :, x, y] = bbox 
print(np.all(output_tf == output_np)) 
# True