使用tf函数代替for循环tensorflow获取切片/掩码 [英] Use tf functions instead of for loops tensorflow to get slice/mask

问题描述

我有2个张量，我的预测张量 pred 的形状为 [批，行，列，列，深度，值] 或 [32、40、60，2，2] .我的地面真实张量 y 是形状 [batch，num_objs，vals] 或 [32，4，10] . y 张量具有与 pred 张量的一个切片相对应的值.我使用

I have 2 tensors, my prediction tensor pred is of shape [batch, rows, cols, depth, vals] or [32, 40, 60, 2, 2]. My ground truth tensor y is shape [batch, num_objs, vals] or [32, 4, 10]. The y tensor has values that correspond to a slice of the pred tensor. I get the indexes using

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

现在，我想使用 true_grid_coords 中的 vals 创建与我的 pred 张量相同形状的蒙版.我想要类似 [batch，rows，cols，depth，vals] = 1 的内容在 true_grid_coords 中有一个vals对，它们等于 pred 张量中的( rows，cols )对.这就是我现在所拥有的，它可以工作，但是速度很慢.

Now I want to create a mask with the same shape as my pred tensor using the vals from true_grid_coords. I want something like [batch, rows, cols, depth, vals] = 1 if there in one vals pair in true_grid_coords that equals a (rows, cols) pair from the pred tensor. This is what I have now, it works but it is slow.

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循环的情况下获得此掩码张量?

Is there a way to get this mask tensor without using the for loops?

推荐答案

这是您可以执行的操作:

This is how you can do that:

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)

这篇关于使用tf函数代替for循环tensorflow获取切片/掩码的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！