Slice pytorch tensor using coordinates tensor without loop - pytorch

I have a tensor T with dimension (d1 x d2 x d3 x ... dk) and a tensor I with dimension (p x q). Here, I contains coordinates of T but q < k, each column of I corresponds to a dimension of T. I have another tensor V of dimension p x di x ...dj where sum([di, ..., dj]) = k - q. (di, .., dj) corresponds to missing dimensions from I. I need to perform T[I] = V
A specific example of such problem using numpy array posted here[1].
The solution[2] uses fancy indexing[3] which relies on numpy.index_exp. In case of pytorch such option is not available. Is there any alternative way to mimic this in pytorch without using loops or casting tensors to numpy array?
Below is a demo:
import torch
t = torch.randn((32, 16, 60, 64)) # tensor
i0 = torch.randint(0, 32, (10, 1)).to(dtype=torch.long) # indexes for dim=0
i2 = torch.randint(0, 60, (10, 1)).to(dtype=torch.long) # indexes for dim=2
i = torch.cat((i0, i2), 1) # indexes
v = torch.randn((10, 16, 64)) # to be assigned
# t[i0, :, i2, :] = v ?? Obviously this does not work
[1] Slice numpy array using list of coordinates
[2] https://stackoverflow.com/a/42538465/6422069
[3] https://numpy.org/doc/stable/reference/generated/numpy.s_.html

After some discussion in the comments, we arrived at the following solution:
import torch
t = torch.randn((32, 16, 60, 64)) # tensor
# indices
i0 = torch.randint(0, 32, (10,)).to(dtype=torch.long) # indexes for dim=0
i2 = torch.randint(0, 60, (10,)).to(dtype=torch.long) # indexes for dim=2
v = torch.randn((10, 16, 64)) # to be assigned
t[(i0, slice(None), i2, slice(None))] = v

Related

Multiply a [3, 2, 3] by a [3, 2] tensor in pytorch (dot product along dimension)

Given the following tensors x and y with shapes [3,2,3] and [3,2]. I want to multiply the tensors along the 2nd dimension, this is expected to be a kind of dot product and scaling along the axis and return a [3,2,3] tensor.
import torch
a = [[[0.2,0.3,0.5],[-0.5,0.02,1.0]],[[0.01,0.13,0.06],[0.35,0.12,0.0]], [[1.0,-0.3,1.0],[1.0,0.02, 0.03]] ]
b = [[1,2],[1,3],[0,2]]
x = torch.FloatTensor(a) # shape [3,2,3]
y = torch.FloatTensor(b) # shape [3,2]
The expected output :
Expected output shape should be [3,2,3]
#output = [[[0.2,0.3,0.5],[-1.0,0.04,2.0]],[[0.01,0.13,0.06],[1.05,0.36,0.0]], [[0.0,0.0,0.0],[2.0,0.04, 0.06]] ]
I have tried the two below but none of them is giving the desired output and output shape.
torch.matmul(x,y)
torch.matmul(x,y.unsqueeze(1).shape)
What is the best way to fix this?
This is just broadcasted multiply. So you can insert a unitary dimension on the end of y to make it a [3,2,1] tensor and then multiply by x. There are multiple ways to insert unitary dimensions.
# all equivalent
x * y.unsqueeze(2)
x * y[..., None]
x * y[:, :, None]
x * y.reshape(3, 2, 1)
You could also use torch.einsum.
torch.einsum('abc,ab->abc', x, y)

multiplying each element of a matrix by a vector (or array)

Say I have a an input array of size (64,100)
t = torch.randn((64,100))
Now say I want to multiply each of the 6400 elements of t with 6400 separate vectors each of size 256 to produce a tensor of size [64, 100, 256]. This is what I am doing currently -
import copy
def clones(module, N):
"Produce N identical layers."
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
linears = clones(nn.Linear(1,256, bias=False), 6400)
idx = 0
t_final = []
for i in range(64):
t_bs = []
for j in range(100):
t1 = t[i, j] * linears[idx].weight.view(-1)
idx += 1
t_bs.append(t1)
t_bs = torch.cat(t_bs).view(1, 100, 256)
t_final.append(t_bs)
t_final = torch.cat(t_final)
print(t_final.shape)
Output: torch.Size([64, 100, 256])
Is there a faster and cleaner way of doing the same thing? I tried torch.matmul and torch.dot but couldn't do any better.
It seems broadcast is what you are looking for.
t = torch.randn((64,100)).view(6400, 1)
weights = torch.randn((6400, 256))
output = (t * weights).view(64, 100, 256)
You don't actually need to clone your linear layer if you really want to multiply tenor t with the same weight of linear layer for 6400 times. rather you can do the following:
t = torch.randn((64,100)).unsqueeze(-1)
w = torch.rand((256)).view(1,1,256).repeat(64, 100, 1)
#or
w = torch.stack(6400*[torch.rand((256))]).view(64,100,256)
result = t*w # shape: [64, 100, 256]
However, If your want to keep the same structure you currently have, then you can do something following:
t = torch.randn((64,100)).unsqueeze(-1)
w = torch.stack([linears[i].weight for i in range(len(linears))]).view(64,100,256)
result = t*w # shape: [64, 100, 256]

How to find the mean and the covariance of a 2d activation map (pytorch)

I have a tensor of shape [h, w], which consists of a normalized, 2-dimensional activation map. Considering this to be some distribution, I want to find the mean and the covariance within this activation map in pytorch. Is there an efficient way to do that?
You can use the following code, where activation_map is a tensor of shape (h,w), with non-negative elements, and is normalised (activation_map.sum() is 1):
activation_map = torch.tensor(
[[0.2, 0.1, 0.0],
[0.1, 0.2, 0.4]])
h, w = activation_map.shape
range_h = torch.arange(h)
range_w = torch.arange(w)
idxs = torch.stack([
range_w[None].repeat(h, 1),
range_h[:, None].repeat(1, w)
])
map_flat = activation_map.view(-1)
idxs_flat = idxs.reshape(2, -1).T
mean = (map_flat[:, None] * idxs_flat).sum(0)
mats = idxs_flat[:, :, None] # idxs_flat[:, None, :]
second_moments = (map_flat[:, None, None] * mats).sum(0)
covariance = second_moments - mean[:, None] # mean[None]
# mean:
# tensor([1.1000, 0.7000])
# covariance:
# tensor([[0.6900, 0.2300],
# [0.2300, 0.2100]])
One way for the covariance matrix:
h,w = 3,5
def cov(X):
X = X/np.sqrt(X.size(0) - 1)
return X.T # X
x = torch.randn(h,w)
print(x)
c = cov(x)
print(c)
Out:
tensor([[-1.5029e-01, -2.0626e-01, -7.7845e-01, -1.6811e+00, 5.0312e-01],
[ 4.4658e-01, -1.8570e+00, -6.2250e-01, -1.0989e+00, 1.6159e+00],
[ 6.8612e-01, -4.2650e-02, -9.5685e-01, -1.7947e-03, 2.1187e-01]])
tensor([[ 0.3464, -0.4138, -0.4088, -0.1197, 0.3957],
[-0.4138, 1.7464, 0.6787, 1.1938, -1.5568],
[-0.4088, 0.6787, 0.9545, 0.9972, -0.8001],
[-0.1197, 1.1938, 0.9972, 2.0169, -1.3110],
[ 0.3957, -1.5568, -0.8001, -1.3110, 1.4546]])
The mean() should be trivial just refer the documentation.

Loop over tensor dimension 0 (NoneType) with second tensor values

I have a tensor a, I'd like to loop over the rows and index values based on another tensor l. i.e. l suggests the length of the vector I need.
sess = tf.InteractiveSession()
a = tf.constant(np.random.rand(3,4)) # shape=(3,4)
a.eval()
Out:
array([[0.35879311, 0.35347166, 0.31525201, 0.24089784],
[0.47296348, 0.96773956, 0.61336239, 0.6093023 ],
[0.42492552, 0.2556728 , 0.86135674, 0.86679779]])
l = tf.constant(np.array([3,2,4])) # shape=(3,)
l.eval()
Out:
array([3, 2, 4])
Expected output:
[array([0.35879311, 0.35347166, 0.31525201]),
array([0.47296348, 0.96773956]),
array([0.42492552, 0.2556728 , 0.86135674, 0.86679779])]
The tricky part is the fact that a could have None as first dimension since it's what is usually defined as batch size through placeholder.
I can not just use mask and condition as below since I need to compute the variance of each row individually.
condition = tf.sequence_mask(l, tf.reduce_max(l))
a_true = tf.boolean_mask(a, condition)
a_true
Out:
array([0.35879311, 0.35347166, 0.31525201, 0.47296348, 0.96773956,
0.42492552, 0.2556728 , 0.86135674, 0.86679779])
I also tried to use tf.map_fn but can't get it to work.
elems = (a, l)
tf.map_fn(lambda x: x[0][:x[1]], elems)
Any help will be highly appreciated!
TensorArray object can store tensors of different shapes. However, it is still not that simple. Take a look at this example that does what you want using tf.while_loop() with tf.TensorArray and tf.slice() function:
import tensorflow as tf
import numpy as np
batch_data = np.array([[0.35879311, 0.35347166, 0.31525201, 0.24089784],
[0.47296348, 0.96773956, 0.61336239, 0.6093023 ],
[0.42492552, 0.2556728 , 0.86135674, 0.86679779]])
batch_idx = np.array([3, 2, 4]).reshape(-1, 1)
x = tf.placeholder(tf.float32, shape=(None, 4))
idx = tf.placeholder(tf.int32, shape=(None, 1))
n_items = tf.shape(x)[0]
init_ary = tf.TensorArray(dtype=tf.float32,
size=n_items,
infer_shape=False)
def _first_n(i, ta):
ta = ta.write(i, tf.slice(input_=x[i],
begin=tf.convert_to_tensor([0], tf.int32),
size=idx[i]))
return i+1, ta
_, first_n = tf.while_loop(lambda i, ta: i < n_items,
_first_n,
[0, init_ary])
first_n = [first_n.read(i) # <-- extracts the tensors
for i in range(batch_data.shape[0])] # that you're looking for
with tf.Session() as sess:
res = sess.run(first_n, feed_dict={x:batch_data, idx:batch_idx})
print(res)
# [array([0.3587931 , 0.35347167, 0.315252 ], dtype=float32),
# array([0.47296348, 0.9677396 ], dtype=float32),
# array([0.4249255 , 0.2556728 , 0.86135674, 0.8667978 ], dtype=float32)]
Note
We still had to use batch_size to extract elements one by one from first_n TensorArray using read() method. We can't use any other method that returns Tensor because we have rows of different sizes (except TensorArray.concat method but it will return all elements stacked in one dimension).
If TensorArray will have less elements than index you pass to TensorArray.read(index) you will get InvalidArgumentError.
You can't use tf.map_fn because it returns a tensor that must have all elements of the same shape.
The task is simpler if you only need to compute variances of the first n elements of each row (without actually gather elements of different sizes together). In this case we could directly compute variance of sliced tensor, put it to TensorArray and then stack it to tensor:
n_items = tf.shape(x)[0]
init_ary = tf.TensorArray(dtype=tf.float32,
size=n_items,
infer_shape=False)
def _variances(i, ta, begin=tf.convert_to_tensor([0], tf.int32)):
mean, varian = tf.nn.moments(
tf.slice(input_=x[i], begin=begin, size=idx[i]),
axes=[0]) # <-- compute variance
ta = ta.write(i, varian) # <-- write variance of each row to `TensorArray`
return i+1, ta
_, variances = tf.while_loop(lambda i, ta: i < n_items,
_variances,
[ 0, init_ary])
variances = variances.stack() # <-- read from `TensorArray` to `Tensor`
with tf.Session() as sess:
res = sess.run(variances, feed_dict={x:batch_data, idx:batch_idx})
print(res) # [0.0003761 0.06120085 0.07217039]

No N-dimensional tranpose in PyTorch

PyTorch's torch.transpose function only transposes 2D inputs. Documentation is here.
On the other hand, Tensorflow's tf.transpose function allows you to transpose a tensor of N arbitrary dimensions.
Can someone please explain why PyTorch does not/cannot have N-dimension transpose functionality? Is this due to the dynamic nature of the computation graph construction in PyTorch versus Tensorflow's Define-then-Run paradigm?
It's simply called differently in pytorch. torch.Tensor.permute will allow you to swap dimensions in pytorch like tf.transpose does in TensorFlow.
As an example of how you'd convert a 4D image tensor from NHWC to NCHW (not tested, so might contain bugs):
>>> img_nhwc = torch.randn(10, 480, 640, 3)
>>> img_nhwc.size()
torch.Size([10, 480, 640, 3])
>>> img_nchw = img_nhwc.permute(0, 3, 1, 2)
>>> img_nchw.size()
torch.Size([10, 3, 480, 640])
Einops supports verbose transpositions for arbitrary number of dimensions:
from einops import rearrange
x = torch.zeros(10, 3, 100, 100)
y = rearrange(x, 'b c h w -> b h w c')
x2 = rearrange(y, 'b h w c -> b c h w') # inverse to the first
(and the same code works for tensorfow as well)

Resources