Is advanced indexing available across n-dimensions in TensorFlow? - python

In PyTorch, we can use standard Pythonic indexing to apply advanced indexing across n-dimensions.
preds is a Tensor of shape [1, 3, 64, 64, 12].
a, b, c, d are 1-dimensional Tensors of the same length. In this instance that length is 9, but this is not always the case.
PyTorch example achieving the desired result:
result = preds[a, b, c, d]
result.shape
>>> [9, 12]
How can this be reproduced in TensorFlow, starting from the same 5 Tensors and creating the same output?
I have tried tf.gather whichs seem to be able to produce the same behaviour in a single dimension:
tf.shape(tf.gather(preds, a))
>>> [9, 3, 64, 64, 12]
Is it possible to extend this to eventually reach the desired output of shape [9, 12]?
I have also noted the presence of tf.gather_nd which seems like it may be relevant here but I cannot determine how I would employ it from the documentation.

Yes, gather_nd can do that
t = tf.random.uniform(shape=(1,3,64,64,12))
# i_n = indices along n-th dim
i_1 = tf.constant([0,0,0,0,0,0,0,0,0])
i_2 = tf.constant([0,1,2,1,2,2,1,0,0])
i_3 = tf.constant([0,21,15,63,22,17,21,54,39])
i_4 = tf.constant([0,16,26,51,33,45,48,29,1])
i = tf.stack([i_1, i_2, i_3, i_4], axis=1) # i.shape == (9,4)
tf.gather_nd(t, i).shape # (9,12)

Related

Slicing a tensor with a tensor of indices and tf.gather

I am trying to slice a tensor with a indices tensor. For this purpose I am trying to use tf.gather.
However, I am having a hard time understanding the documentation and don't get it to work as I would expect it to:
I have two tensors. An activations tensor with a shape of [1,240,4] and an ids tensor with the shape [1,1,120]. I want to slice the second dimension of the activations tensor with the indices provided in the third dimension of the ids tensor:
downsampled_activations = tf.gather(activations, ids, axis=1)
I have given it the axis=1 option since that is the axis in the activations tensor I want to slice.
However, this does not render the expected result and only gives me the following error:
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0,0,1] = 1 is not in [0, 1)
I have tried various combinations of the axis and batch_dims options, but to no avail so far and the documentation doesn't really help me on my path. Anybody care to explain the parameters in more detail or on the example above would be very helpful!
Edit:
The IDs are precomputed before runtime and come in through an input pipeline as such:
features = tf.io.parse_single_example(
serialized_example,
features={ 'featureIDs': tf.io.FixedLenFeature([], tf.string)}
They are then reshaped into the previous format:
feature_ids_raw = tf.decode_raw(features['featureIDs'], tf.int32)
feature_ids_shape = tf.stack([batch_size, (num_neighbours * 4)])
feature_ids = tf.reshape(feature_ids_raw, feature_ids_shape)
feature_ids = tf.expand_dims(feature_ids, 0)
Afterwards they have the previously mentioned shape (batch_size = 1 and num_neighbours = 30 -> [1,1,120]) and I want to use them to slice the activations tensor.
Edit2: I would like the output to be [1,120,4]. (So I would like to gather the entries along the second dimension of the activations tensor in accordance with the IDs stored in my ids tensor.)
You can use :
downsampled_activations =tf.gather(activations , tf.squeeze(ids) ,axis = 1)
downsampled_activations.shape # [1,120,4]
In most cases, the tf.gather method needs 1d indices, and that is right in your case, instead of indices with 3d (1,1,120), a 1d is sufficient (120,). The method tf.gather will look at the axis( = 1) and return the element at each index provided by the indices tensor.
tf.gather Gather slices from params axis axis according to indices.
Granted that the documentation is not the most expressive, and the emphasis should be placed on the slices (since you index slices from the axis and not elements, which is what I suppose you mistakenly took it for).
Let's take a much smaller example:
activations_small = tf.convert_to_tensor([[[1, 2, 3, 4], [11, 22, 33, 44]]])
print(activations_small.shape) # [1, 2, 4]
Let's picture this tensor:
XX 4 XX 44 XX XX
XX 3 XX 33 X XX
XXX 2 XX 22XX XX
X-----X-----+X XX
| 1 | 11 | XX
+-----+-----+X
tf.gather(activations1, [0, 0], axis=1) will return
<tf.Tensor: shape=(1, 2, 4), dtype=int32, numpy=
array([[[1, 2, 3, 4],
[1, 2, 3, 4]]], dtype=int32)>
What tf.gather did was to look from axis 1, and picks up index 0 (ofc, two times i.e. [0, 0]). If you were to run tf.gather(activations1, [0, 0, 0, 0, 0], axis=1).shape, you'd get TensorShape([1, 5, 4]).
Your Error
Now let's try to trigger the error that you're getting.
tf.gather(activations1, [0, 2], axis=1)
InvalidArgumentError: indices[1] = 2 is not in [0, 2) [Op:GatherV2]
What happened here was that when tf.gather looks from axis 1 perspective, there's no item (column if you will) with index = 2.
I guess this is what the documentation is hinting at by
param:<indices> The index Tensor. Must be one of the following types: int32, int64. Must be in range [0, params.shape[axis]).
Your (potential) solution
From the dimensions of indices, and that of the expected result from your question, I am not sure if the above was very obvious to you.
tf.gather(activations, indices=[0, 1, 2, 3], axis=2) or anything with indices within the range of indices in [0, activations.shape[2]) i.e. [0, 4) would work. Anything else would give you the error that you're getting.
There's a verbatim answer below in case that's your expected result.

Tensorflow - indexing according to batch position

I'm working on masking r-cnn and I have a problem with indexing the masks according to labels.
Here's what I want to achieve: I have a tensor (?,28,28,c), where ? is unknown batch_size, "28x28" are 2d coordinates and c stands for different labels, then I have a list of indices (basically my label predictions) (?,) of int32. Now I want to extract the masks for a given label according to batch index -> make it a (?,28,28,1) tensor.
I tried self.masks_sigmoids = tf.gather(self.final_conv, self.label_predictions, axis=3), but the shape remained the same.
I also looked at tf.gather_nd here http://www.riptutorial.com/tensorflow/example/29069/how-to-use-tf-gather-nd, and I guess this is the right path, but I don't know how to incorporate that I want the indices according to batch index (in numpy (b_i,:,:,c_i))
I also get a feeling that my question is somewhat similar to Batched 4D tensor Tensorflow indexing, though my problem seems less complicated. However, that question is old in terms of the quick development of tensorflow, so I'm asking for a possibly better, more clear solution. EDIT: Even a dirty solution might beneficial as I didn't get the question in the linked SO (already wrote a comment asking to clarify the question), thus I don't get much from the only answer. It might be beneficial for the community as well, because this question is simpler, which means it would demonstrate the solution more clearly.
Solution 1: more generic
You can look at the answer here, it's basically the same problem as yours, with different dimensions.
The solution described there is to create a [?, 28, 28, 4]-shaped tensor indices where indices[i, x, y, :] = [i, x, y, self.label_predictions[i]], and then use tf.gather_nd:
self.masks_sigmoids = tf.gather_nd(self.final_conv, indices=indices)
Building the indices is not very elegant, as shown in this answer (with one more dimension for you), but easy in itself.
Solution 2: A bit more elegant and adapted to your problem
This solution is very similar to the first one, but avoids creating the [x, y] part of indices. The idea is to use the slicing capabilities of gather_nd to avoid writing [x, y] in indices for each (i, x, y), by transposing the data before gathering it. I'll put the whole code here, including how to create indices and how to test:
import numpy as np
import tensorflow as tf
N_CHANNELS = 5
pl=tf.placeholder(dtype=tf.int32, shape=(None, 28, 28, N_CHANNELS))
# Indices we'll use. batch_size = 4 here.
label_predictions = tf.constant([0, 2, 0, 3])
# Indices of shape [?, 2], with indices[i] = [i, self.label_predictions[i]],
# which is easy to do with tf.range() and tf.stack()
indices = tf.stack([tf.range(tf.size(label_predictions)), label_predictions], axis=-1)
# [[0, 0], [1, 2], [2, 0], [3, 3]]
transposed = tf.transpose(pl, perm=[0, 3, 1, 2])
gathered = tf.gather_nd(transposed, indices) # Should be of shape (4, 2, 3)
result = tf.expand_dims(gathered, -1)
initial_value = np.arange(4*28*28*N_CHANNELS).reshape((4, 28, 28, N_CHANNELS))
sess = tf.InteractiveSession()
res = sess.run(result, feed_dict={pl: initial_value})
# print(res)
print("checking validity")
for i in range(4):
for x in range(28):
print(x)
for y in range(28):
assert res[i, x, y, 0] == initial_value[i, x, y, indices[i, 1].eval()]
print("All assertions passed")

How to replicate numpy.choose() in tensorflow?

I'm trying to efficiently replicate numpy's ndarray.choose() method.
Here's a numpy example of what I'm looking for:
b = np.arange(15).reshape(3, 5)
c = np.array([1,0,4])
c.choose(b.T) # trying to replicate in tensorflow
-> array([ 1, 5, 14])
The best I've been able to do with this is generate a batch_size square matrix (which is huge if batch size is huge) and take the diagonal of it:
tf_b = tf.constant(b)
tf_c = tf.constant(c)
sess.run(tf.diag_part(tf.gather(tf.transpose(tf_b), tf_c)))
-> array([ 1, 5, 14])
Is there a way to do this that is just linear in the first dimension (instead of squared)?
Yeah, there's an easier way to do this. Flatten your b array to 1-d, so it's [0, 1, 2, ..., 13, 14]. Take an array of indices that are in the range of the number of 'choices' you are taking (3 in your case). That will be [0, 1, 2]. Multiply this range by the second dimension of your original shape, which is the number of options for each choice (5 in your case). That gives you [0, 5, 10]. Then add your indices to this to obtain [1, 5, 14]. Now you're good to call tf.gather().
Here is some code that I've taken from here that does a similar thing for RNN outputs. Yours will be slightly different, but the idea is the same.
index = tf.range(0, batch_size) * max_length + (length - 1)
flat = tf.reshape(output, [-1, out_size])
relevant = tf.gather(flat, index)
return relevant
In a big picture, the operation is pretty straightforward. You use the range operation to get the index of the beginning of each row, then add the index of where you are in each row. I think doing it in 1D is easiest, so that's why we flatten it.

Tensorflow multiple scalar multiplication

I have a 3d Tensor with [batch_size,x,y] and a vector [batch_size].
I want to scalar multiply the i-th matrix [x,y] with the i-th entry of the given vector.
Is there a build in function in Tensorflow or do i have to use the tf.while_loop?
You can do this with broadcasting. You need to reshape the vector first.
a = tf.constant([[[1,1],[2,2]],[[3,3],[4,4]]])
b = tf.constant([2,3])
c = tf.reshape(b, [-1,1,1])
d = a * c
>>> sess.run(d)
array([[[ 2, 2],
[ 4, 4]],
[[ 9, 9],
[12, 12]]], dtype=int32)
I don't if there is built in function, but you also don't need to use while loop. You can do basic array manipulation. e.g.:
a=tf.random_uniform([3,5,8])
b=tf.random_uniform([3])
c=tf.expand_dims(tf.expand_dims(b, -1),1)
c=tf.tile(c,[1,5,8])
d=tf.multiply(a,c)
sess=tf.Session()
sess.run([a,b,c,d])
It should work.

What is the TensorFlow equivalent of this numpy array permutation?

I need to permute elements of a tensor in TF according to a given indexing. From 2 arrays a and b(indices), I need to compute a new array that permutes the elements in a according to the indices in b. For indices that are empty, it should fill with NA (or equivalent).
For example,
a = [10, 20, 30]
b = [-1, 0, 3]
output = [ 10, 20, NA, NA, 30]
I need to code the equivalent of what happens to the following numpy arrays but for TF tensors.
a = np.array([10,20,30])
b = np.array([-1,0,3])
mini = abs(np.min(b))
maxi = abs(np.max(b))
output = np.zeros(maxi+mini+1)
for ai,bi in zip(a,b):
output[bi+mini]= ai
How do I do this with TensorFlow tensors?
African or European?
If you know that the indices are strictly increasing, tf.sparse_to_dense does what you want.
If the indices are distinct but in increasing order, you can use tf.sparse_reorder to fix the order and then use tf.sparse_tensor_to_dense.
If there are duplicates and you want matching values to add, use tf.unsorted_segment_sum.
If there are duplicates and you want the last entry to win (corresponding exactly to your Python loop), use tf.dynamic_stitch.
Apologies for the zoo of options. The ops were all added for different reasons, so the overall design is not particularly clean.
I found a way of achieving this, I'm posting my answer here in case it helps anyone else.
The scatter_nd function in TensorFlow is very handy in this situation.
The following code permutes elements in the input tensor I according to the transformation given in tensor T. scatter_nd is used to create the new tensor according to this permutation.
sess = tf.InteractiveSession()
I = tf.constant([10,20,30])
T = tf.constant([-1,0,3])
T = T - tf.reduce_min(T)
T_shape = int(T.get_shape()[0])
T = tf.reshape(T, [T_shape,1])
O_shape = tf.reduce_max(T)+1
O = tf.scatter_nd(T, I, [O_shape])
print(sess.run([I,T,O]))
sess.close()
This code performs the following task:
Given
Input = [10, 20, 30]
Transformation = [-1, 0, 3]
Computes
Output = [10, 20, 0, 0, 30]

Categories