Use tf functions instead of for loops tensorflow to get slice/mask

问题

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)

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)

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

回答1:

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)

来源：https://stackoverflow.com/questions/61053317/use-tf-functions-instead-of-for-loops-tensorflow-to-get-slice-mask

标签

python

tensorflow

for-loop