What is smart way to get batched gather? - python

I have two matrices, A and B, with shapes (n, m, k) and (n, m) respectively. n is the batch size, m is the amount of data in a batch, and k is the feature size.
Each element of B is an index less than m (specifically B = torch.randint(high=m, shape=(n,m))).
I want to implement [A[i][B[i]] for i in range(n)] in a smarter way.
Is there a better way in pytorch to implement this without doing for loop?

You can use
a[torch.arange(n)[:, None], b]
An example:
>>> n, m, k = 3, 2, 5
>>> a = torch.arange(30).view(n, m, k)
>>> b = torch.randint(high=m, size=(n,m))
# first indexer (of shape (n, 1))
>>> torch.arange(n)[:, None]
tensor([[0],
[1],
[2]])
# second indexer
>>> b
tensor([[1, 0],
[0, 1],
[1, 1]])
The indexers have the shape (3, 1) and (3, 2) respectively so they'll be broadcasted to (3, 2) to effectively have
tensor([[0, 0],
[1, 1],
[2, 2]])
and
tensor([[1, 0],
[0, 1],
[1, 1]])
which says: for the first row, take 1st (k,) array and put the result and take 0th (k,) array and put the result. This fills in a (m, k) array in the output which is repeated n times for each row,
to get
>>> a[torch.arange(n)[:, None], b]
tensor([[[ 5, 6, 7, 8, 9],
[ 0, 1, 2, 3, 4]],
[[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19]],
[[25, 26, 27, 28, 29],
[25, 26, 27, 28, 29]]])
comparing with list comprehension:
>>> [a[i][b[i]] for i in range(n)]
[tensor([[5, 6, 7, 8, 9],
[0, 1, 2, 3, 4]]),
tensor([[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19]]),
tensor([[25, 26, 27, 28, 29],
[25, 26, 27, 28, 29]])]

Related

Is there a vectorized way of appending/indexing values from nd_array 2 based off component r/c values in nd_array 1?

I'm looking to vectorize some code that I have, but am unsure how to approach it or if it's even possible. Here's what I have so far:
arr_row = [0, 1, 2, ..., n] # Array of size n with random integers
arr_col = [0, 1, 2, ..., n] # Array of size n with random integers
for i in range(n):
r = arr_row[i]
c = arr_col[i]
result = n_by_n_table[r,c]
So briefly about the above. I have two arrays that are n elements, I iterate through i to n getting the item at the row/col location obtained from my row/col arrs. I get this result from a guaranteed n x n matrix. I am hoping that I could vectorized this for a performance improvement. Here's what I had in mind:
vectorized = nd_array((n, 3))
vectorized[:, 0] = [0, 1, 2, ..., n] # Likely using np.rand func or something here.
vectorized[:, 1] = [0, 1, 2, ..., n] # Same as above
vectorized[:, 2] = n_by_n_table[vectorized[:,0], vectorized[:,1]]
Is this possible and any concerns doing the above?
example data
arr_row = [5,1,0,3,2,4]
arr_col = [1,4,2,0,5,3]
n_by_n_table = np.array(range(36)).reshape(6,6)
n_by_n_table :
[[ 0, 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10, 11],
[12, 13, 14, 15, 16, 17],
[18, 19, 20, 21, 22, 23],
[24, 25, 26, 27, 28, 29],
[30, 31, 32, 33, 34, 35]]
solution
n_by_n_table[arr_row,arr_col]
result:
array([31, 10, 2, 18, 17, 27])

Pytorch gather question (3D Computer Vision)

I have N groups of C-dimension points. In each groups there are M points. So, there is a tensor of (N, M, C). Let's call it features.
I calculated the maximum element and the index through M dimension, to find the maximum points for each C dimension (a max pooling operation), resulting max tensor (N, 1, C) and index tensor (N, 1, C).
I have another tensor of shape (N, M, 3) storing the geometric coordinates of those N*M high-dimention points. Now, I want to use the index from the maximum points in each C dimension, to get the coordinates of all those maximum points.
For example, N=2, M=4, C=6.
The coordinate tensor, whose shape is (2, 4, 3):
[[[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
[8, 7, 6]]
[11, 12, 13]
[14, 15, 16]
[17, 18, 19]
[18, 17, 16]]]
The indices tensor, whose shape is (2, 1, 6):
[[[0, 1, 2, 1, 2, 3]]
[[1, 2, 3, 2, 1, 0]]]
For example, the first element in indices is 0, I want to grab [1, 2, 3] from the coordinate tensor out. For the second element (1), I want to grab [4, 5, 6] out. For the third element in the next dimension (3), I want to grab [18, 17, 16] out.
The result tensor will look like:
[[[1, 2, 3] # 0
[4, 5, 6] # 1
[7, 8, 9] # 2
[4, 5, 6] # 1
[7, 8, 9] # 2
[8, 7, 6]] # 3
[[14, 15, 16] # 1
[17, 18, 19] # 2
[18, 17, 16] # 3
[17, 18, 19] # 2
[14, 15, 16] # 1
[11, 12, 13]]]# 0
whose shape is (2, 6, 3).
I tried to use torch.gather but I can not get it worked. I wrote a naive algorithm enumerating all N groups, but indeed it is slow, even using TorchScript's JIT. So, how to write this efficiently in pytorch?
You can use integer array indexing combined with broadcasting semantics to get your result.
import torch
x = torch.tensor([
[[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
[8, 7, 6]],
[[11, 12, 13],
[14, 15, 16],
[17, 18, 19],
[18, 17, 16]],
])
i = torch.tensor([[[0, 1, 2, 1, 2, 3]],
[[1, 2, 3, 2, 1, 0]]])
# rows is shape [2, 1], cols is shape [2, 6]
rows = torch.arange(x.shape[0]).type_as(i).unsqueeze(1)
cols = i.squeeze(1)
# y is [2, 6, ...]
y = x[rows, cols]

Extracting blocks from block diagonal PyTorch tensor

I have a tensor of shape (m*n, m*n) and I want to extract a tensor of size (n, m*n) containing the m blocks of size n*n that are on the diagonal. For example:
>>> a
tensor([[1, 2, 0, 0],
[3, 4, 0, 0],
[0, 0, 5, 6],
[0, 0, 7, 8]])
I want to have a function extract(a, m, n) that will output:
>>> extract(a, 2, 2)
tensor([[1, 2, 5, 6],
[3, 4, 7, 8]])
I've thought of using some kind of slicing, because the blocks can be expressed by:
>>> for i in range(m):
... print(a[i*m: i*m + n, i*m: i*m + n])
tensor([[1, 2],
[3, 4]])
tensor([[5, 6],
[7, 8]])
You can take advantage of reshape and slicing:
import torch
import numpy as np
def extract(a, m, n):
s=(range(m), np.s_[:], range(m), np.s_[:]) # the slices of the blocks
a.reshape(m, n, m, n)[s] # reshaping according to blocks and slicing
return a.reshape(m*n, n).T # reshape to desired output format
Example:
a = torch.arange(36).reshape(6,6)
a
tensor([[ 0, 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10, 11],
[12, 13, 14, 15, 16, 17],
[18, 19, 20, 21, 22, 23],
[24, 25, 26, 27, 28, 29],
[30, 31, 32, 33, 34, 35]])
extract(a, 3, 2)
tensor([[ 0, 6, 14, 20, 28, 34],
[ 1, 7, 15, 21, 29, 35]])
extract(a, 2, 3)
tensor([[ 0, 6, 12, 21, 27, 33],
[ 1, 7, 13, 22, 28, 34],
[ 2, 8, 14, 23, 29, 35]])
You can achieve this for a block diagonal matrix (of equally sized square blocks of width n) with torch.nonzero():
>>> n = 2
>>> a[a.nonzero(as_tuple=True)].view(n,n,-1)
tensor([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])

numpy map 2D array values

I'm trying to map values of 2D numpy array, i.e. to iterate (efficiently) over rows and append values based on row index.
One of approaches I have tried is:
source = misc.imread(fname) # Load some image
img = np.array(source, dtype=np.float64) / 255 # Cast and normalize values
w, h, d = tuple(img.shape) # Get dimensions
img = np.reshape(img, (w * h, d)) # Flatten 3D to 2D
# The actual problem:
# Map (R, G, B) pixels to (R, G, B, X, Y) to preserve position
img_data = ((px[0], px[1], px[2], idx % w, int(idx // w)) for idx, px in enumerate(img))
img_data = np.fromiter(img_data, dtype=tuple) # Get back to np.array
but the solution raises: ValueError: cannot create object arrays from iterator
Can anyone suggest how to perform efficiently this absurdly simple operation in numpy? It's out of my mind how intricate is this library... And why that code consumes a few gigs of memory for 7000x5000 px?
Thanks
maybe np.concatenate and np.indices:
np.concatenate((np.arange(40).reshape((4,5,2)), *np.indices((4,5,1))), axis=-1)[:,:,:-1]
Out[264]:
array([[[ 0, 1, 0, 0],
[ 2, 3, 0, 1],
[ 4, 5, 0, 2],
[ 6, 7, 0, 3],
[ 8, 9, 0, 4]],
[[10, 11, 1, 0],
[12, 13, 1, 1],
[14, 15, 1, 2],
[16, 17, 1, 3],
[18, 19, 1, 4]],
[[20, 21, 2, 0],
[22, 23, 2, 1],
[24, 25, 2, 2],
[26, 27, 2, 3],
[28, 29, 2, 4]],
[[30, 31, 3, 0],
[32, 33, 3, 1],
[34, 35, 3, 2],
[36, 37, 3, 3],
[38, 39, 3, 4]]])
the [:,:,:-1] strips an 'extra' 0 entry, maybe there's a better way

What does transpose(3, 0, 1, 2) mean?

What does this mean?
data.transpose(3, 0, 1, 2)
Also, if data.shape == (10, 10, 10), why do I get ValueError: axes don't match array?
Let me discuss in terms of Python3.
I use the transpose function in python as data.transpose(3, 0, 1, 2)
This is wrong as this operation requires 4 dimensions, while you only provide 3 (as in (10,10,10)). Reproducible as:
>>> a = np.arange(60).reshape((1,4,5,3))
>>> b = a.transpose((2,0,1))
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
ValueError: axes don't match array
You can either add another dimension simply by reshaping (10,10,10) to (1,10,10,10) if the image batch is 1. This can be done as:
w,h,c = original_image.shape #10,10,10
modified_img = np.reshape((1,w,h,c)) #(1,10,10,10)
what does it mean of 3, 0, 1, 2.
For 2D numpy arrays, transpose for an array (matrix) operates just as the names say. But for higher dimensional arrays like yours, it basically works as moveaxis.
>>> a = np.arange(60).reshape((4,5,3))
>>> b = a.transpose((2,0,1))
>>> b.shape
(3, 4, 5)
>>> c = np.moveaxis(a,-1,0)
>>> c.shape
(3, 4, 5)
>>> b
array([[[ 0, 3, 6, 9, 12],
[15, 18, 21, 24, 27],
[30, 33, 36, 39, 42],
[45, 48, 51, 54, 57]],
[[ 1, 4, 7, 10, 13],
[16, 19, 22, 25, 28],
[31, 34, 37, 40, 43],
[46, 49, 52, 55, 58]],
[[ 2, 5, 8, 11, 14],
[17, 20, 23, 26, 29],
[32, 35, 38, 41, 44],
[47, 50, 53, 56, 59]]])
>>> c
array([[[ 0, 3, 6, 9, 12],
[15, 18, 21, 24, 27],
[30, 33, 36, 39, 42],
[45, 48, 51, 54, 57]],
[[ 1, 4, 7, 10, 13],
[16, 19, 22, 25, 28],
[31, 34, 37, 40, 43],
[46, 49, 52, 55, 58]],
[[ 2, 5, 8, 11, 14],
[17, 20, 23, 26, 29],
[32, 35, 38, 41, 44],
[47, 50, 53, 56, 59]]])
As evident, both methods work the same.
The operation converts from (samples, rows, columns, channels) into (samples, channels, rows, cols),maybe opencv to pytorch.
Have a look at numpy.transpose
Use transpose(a, argsort(axes)) to invert the transposition of tensors
when using the axes keyword argument.
Transposing a 1-D array returns an unchanged view of the original
array.
e.g.
>>> x = np.arange(4).reshape((2,2))
>>> x
array([[0, 1],
[2, 3]])
>>>
>>> np.transpose(x)
array([[0, 2],
[1, 3]])
You specified too many values in the transpose
>>> a = np.arange(8).reshape(2,2,2)
>>> a.shape (2, 2, 2)
>>> a.transpose([2,0,1])
array([[[0, 2],
[4, 6]],
[[1, 3],
[5, 7]]])
>>> a.transpose(3,0,1,2) Traceback (most recent call last): File "<interactive input>", line 1, in <module> ValueError: axes don't match array
>>>
From the python documentation on np.transpose, the second argument of the np.transpose function is axes, which is a list of ints, optional
by default and reverse the dimensions, otherwise permute the axes
according to the values given.
Example :
>>> x = np.arange(9).reshape((3,3))
>>> x
array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
>>> np.transpose(x, (0,1))
array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
>>> np.transpose(x, (1,0))
array([[0, 3, 6],
[1, 4, 7],
[2, 5, 8]])
The thing is you have taken a 3 dimensional matrix and applied a 4 dimensional transpose.
your command is to convert a 4d matrix(batch,rows,cols,channel) to another 4d matrix (rows,cols,channel,batch) but you need a command to convert 3d matrix.so remove 3 and write
data.transpose(2, 0, 1).
For all i, j, k, l, the following holds true:
arr[i, j, k, l] == arr.transpose(3, 0, 1, 2)[l, i, j, k]
transpose(3, 0, 1, 2) reorders the array dimensions from (a, b, c, d) to (d, a, b, c):
>>> arr = np.zeros((10, 11, 12, 13))
>>> arr.transpose(3, 0, 1, 2).shape
(13, 10, 11, 12)

Categories

Resources