Используйте функции tf вместо тензорного потока для циклов, чтобы получить срез / маску

Question

У меня есть 2 тензора, тензор моего прогноза pred имеет форму [batch, rows, cols, depth, vals] или [32, 40, 60, 2, 2]. Мой тензор истинности y - это форма [batch, num_objs, vals] или [32, 4, 10]. Тензор y имеет значения, которые соответствуют срезу тензора pred. Я получаю индексы, используя

true_grid_coords = (y[:,:,:2] // params.grid_stride)

Теперь я хочу создать маску той же формы, что и мой тензор pred, используя vals из true_grid_coords. Я хочу что-то вроде [batch, rows, cols, depth, vals] = 1, если в одной паре vals в true_grid_coords есть пара (rows, cols) из тензора pred. Это то, что у меня сейчас, работает, но медленно.

center_coords, wh_coords, obj_scores, class_probs = DetectNet.predict_transform(y_)
center_coords_shape = center_coords.shape
true_obj_coord_mask = tf.constant(1, shape=center_coords_shape, dtype=tf.float32)
true_obj_coord_mask = tf.Variable(true_obj_coord_mask)
true_grid_coords = (y[:,:,:2] // params.grid_stride) % 60
true_2d = tf.constant(1, shape=[2,2], dtype=tf.float32)
for i in range(true_grid_coords.shape[0]):
    for j in range(true_grid_coords.shape[1]):
        grid_x = int(true_grid_coords[i, j, 0].numpy())
        grid_y = int(true_grid_coords[i, j, 1].numpy())
        true_obj_coord_mask[i, grid_y, grid_x, :, :].assign(true_2d)

Есть ли способ получить тензор маски без использования циклов for?

Answer 1

Вот как вы можете это сделать:

import tensorflow as tf

batch = 32
rows = 40
cols = 60
depth = 2
vals = 2
num_objs = 4
# Make some random data
tf.random.set_seed(0)
pred = tf.random.uniform([batch, rows, cols, depth, vals], 0, 1, tf.float32)
true_grid_coords = tf.stack([tf.random.uniform([batch, num_objs], 0, rows, tf.int32),
                             tf.random.uniform([batch, num_objs], 0, rows, tf.int32)],
                            axis=-1)
print(true_grid_coords.shape)
# (32, 4, 2)

# Make index for batch dimension
s1 = tf.shape(true_grid_coords, out_type=true_grid_coords.dtype)
b = tf.range(s1[0])
# Repeat batch index for each object
b = tf.repeat(b, s1[1])
# Concatenate with row and column indices
idx = tf.concat([tf.expand_dims(b, 1), tf.reshape(true_grid_coords, [-1, s1[2]])], axis=1)
# Make mask by scattering values
s2 = tf.shape(pred)
mask = tf.scatter_nd(idx, tf.ones_like(b, dtype=tf.float32), s2[:3])
# Tile mask across last two dimensions
mask = tf.tile(mask[..., tf.newaxis, tf.newaxis], [1, 1, 1, s2[3], s2[4]])
print(mask.shape)
# (32, 40, 60, 2, 2)