CODEWITHSUNDEEP
indexingtensorflowmapping
D.Giunchi
D.Giunchi

Reputation: 1960

tensorflow gather or gather_nd

I have this problem. I have two tensors, one shaped(batch_size=128, height=48, width = 48 , depth=1) that should contain indexes (from 0 to 32x32-1) and another one shaped (batch_size=128, height=32, width = 32 , depth=1) that contains the values that I should map. In this second each matrix in this batch contains its own values.

I would like to map for example, the third "index matrix" with the third "map matrix", considering that index inside each item of the batch range from 0 to 32x32. The same procedure should be applied to all the items in the batch. Since this stuff should be done in the loss function, and I see that we use batches there, how can I do this task? I thought that tf.gather could be helpful, since I've already used but in a simple case (like a constant array), but I don't know how to use it in this complex case.

Edited:

let's suppose I have:
[
   [
    [1,2,0,3],
    [4,2,4,0],
    [1,3,3,1],
    [1,2,4,8]
   ], 
   [
    [3,2,0,0],
    [4,5,4,2],
    [7,6,3,1],
    [1,5,4,8]
   ] 
]  that is a (2,4,4,1) and a tensor
[
  [
   [0.3,0.4,0.6],
   [0.9,0.2,0.5],
   [0.1,0.2,0.1]
  ] , 
  [
   [0.1,0.4,0.5],
   [0.8,0.1,0.6],
   [0.2,0.4,0.3]
  ]
]  that is a (2,3,3,1). 
The first contains the indexes of the second.
I would like an output:
[
   [
    [0.4,0.6,0.3,0.9],
    [0.2,0.6,0.2,0.3],
    [0.4,0.9,0.9,0.4],
    [0.4,0.6,0.2,0.1],
   ],
   [
    [0.8,0.5,0.1,0.1],
    [0.1,0.6,0.1,0.5],
    [0.4,0.2,0.8,0.4],
    [0.4,0.6,0.1,0.3]
   ] 
]

so the indexing should be referred to the single item of the batch. Should I also provide a derivative for this transformation?

Upvotes: 3

Views: 8300

Answers (1)

gdelab
gdelab

Reputation: 6220

If I've understood your question correctly, you'll want to use 

output = tf.gather_nd(tensor2, indices)

with indices being a matrix of shape (batch_size, 48, 48, 3) such that

indices[sample][i][j] = [i, row, col]

where(row, col) are the coordinates of the value you want to fetch in tensor2. They are a translation of the content given in tensor1, coded in 2 numbers instead of 1: 

(row, col) = (tensor1[i, j] / 32, tensor1[i, j] % 32)

To create indices dynamically, something like that should do it: 

batch_size = tf.shape(tensor1)[0]
i_mat = tf.transpose(tf.reshape(tf.tile(tf.range(batch_size), [48*48]),
                                   [48, 48, batch_size]))
# i_mat should be such that i_matrix[i, j, k, l]=i
mat_32 = tf.fill(value=tf.constant(32, dtype=tf.int32), dims=[batch_size, 48, 48])
row_mat = tf.floor_div(tensor1, mat_32)
col_mat = tf.mod(tensor1, mat_32)
indices = tf.stack([i_mat, row_mat, col_mat], axis=-1)

output = tf.gather_nd(tensor2, indices)

EDIT 2

The code above has changed a bit.

The code above considers that your input tensors are actually of shape (batch_size, 48, 48) and (batch_size, 32, 32), as opposed to (batch_size, 48, 48, 1) and (batch_size, 32, 32, 1). To correct that, use for instance 

tensor1=tf.squeeze(tensor1, axis=-1)
tensor2=tf.squeeze(tensor2, axis=-1)

before my code above, and

output = tf.expand_dims(tf.gather_nd(tensor2, indices), axis=-1)
tensor1= tf.expand_dims(tensor1, axis=-1)
tensor2= tf.expand_dims(tensor2, axis=-1)

at the end

Upvotes: 3

Related Questions