In PyTorch calc Euclidean distance instead of matrix multiplication - python

Let say we have 2 matrices:
mat = torch.randn([20, 7]) * 100
mat2 = torch.randn([7, 20]) * 100
n, m = mat.shape
The simplest usual matrix multiplication looks like this:
def mat_vec_dot_product(mat, vect):
n, m = mat.shape
res = torch.zeros([n])
for i in range(n):
for j in range(m):
res[i] += mat[i][j] * vect[j]
return res
res = torch.zeros([n, n])
for k in range(n):
res[:, k] = mat_vec_dot_product(mat, mat2[:, k])
But what if I need to apply L2 norm instead of dot product? The code is next:
def mat_vec_l2_mult(mat, vect):
n, m = mat.shape
res = torch.zeros([n])
for i in range(n):
for j in range(m):
res[i] += (mat[i][j] - vect[j]) ** 2
res = res.sqrt()
return res
for k in range(n):
res[:, k] = mat_vec_l2_mult(mat, mat2[:, k])
Can we do this somehow in an optimal way using Torch or any other libraries? Cause naive O(n^3) Python code works really slow.

Use torch.cdist for L2 norm - euclidean distance
res = torch.cdist(mat, mat2.permute(1,0), p=2)
Here, I have used permute to swap dim of mat2 from 7,20 to 20,7

First of all, matrix multiplication in PyTorch has a built-in operator: #.
So, to multiply mat and mat2 you simply do:
mat # mat2
(should work, assuming dimensions agree).
Now, to compute the Sum of Squared Differences(SSD, or L2-norm of differences) which you seem to compute in your second block, you can do a simple trick.
Since the squared L2-norm ||m_i - v||^2 (where m_i is the i'th row of matrix M and v is the vector) is equal to the dot product <m_i - v, m_i-v> - from linearity of the dot product you obtain: <m_i,m_i> - 2<m_i,v> + <v,v> so you can compute the SSD of each row in M from vector v by computing once the squared L2-norm of each row, once the dot product between each row and the vector and once the L2-norm of the vector. This can be done in O(n^2).
However, for the SSD between 2 matrices you will still get O(n^3). Improvements can be made though by vectorizing the operations instead of using loops.
Here is a simple implementation for 2 matrices:
def mat_mat_l2_mult(mat,mat2):
rows_norm = (torch.norm(mat, dim=1, p=2, keepdim=True)**2).repeat(1,mat2.shape[1])
cols_norm = (torch.norm(mat2, dim=0, p=2, keepdim=True)**2).repeat(mat.shape[0], 1)
rows_cols_dot_product = mat # mat2
ssd = rows_norm -2*rows_cols_dot_product + cols_norm
return ssd.sqrt()
mat = torch.randn([20, 7])
mat2 = torch.randn([7,20])
print(mat_mat_l2_mult(mat, mat2))
The resulting matrix will have at each cell i,j the L2-norm of the difference between each row i in mat and each column j in mat2.

Related

Computing derivatives using numpy

I'm trying to implement a differential in python via numpy that can accept a scalar, a vector, or a matrix.
import numpy as np
def foo_scalar(x):
f = x * x
df = 2 * x
return f, df
def foo_vector(x):
f = x * x
n = x.size
df = np.zeros((n, n))
for mu in range(n):
for i in range(n):
if mu == i:
df[mu, i] = 2 * x[i]
return f, df
def foo_matrix(x):
f = x * x
m, n = x.shape
df = np.zeros((m, n, m, n))
for mu in range(m):
for nu in range(n):
for i in range(m):
for j in range(n):
if (mu == i) and (nu == j):
df[mu, nu, i, j] = 2 * x[i, j]
return f, df
This works fine, but it seems like there should be a way to do this in a single function, and let numpy "figure out" the correct dimensions. I could force everything into a 2-D array form with something like
x = np.array(x)
if len(x.shape) == 0:
x = x.reshape(1, 1)
elif len(x.shape) == 1:
x = x.reshape(-1, 1)
if len(f.shape) == 0:
f = f.reshape(1, 1)
elif len(f.shape) == 1:
f = f.reshape(-1, 1)
and always have 4 nested for loops, but this doesn't scale if I need to generalize to higher-order tensors.
Is what I'm trying to do possible, and if so, how?
I highly doubt there is a function to generate the second parameter returned by the function in Numpy. That being said you can play with the feature of Numpy and Python so to vectorize this and make the function faster. You first need to generate the indices and, then generate the target matrix and set it. Note that operating with N-dimensional generic arrays tends to be slow and tricky in non-trivial cases. The magic * unrolling operator is used to generate N parameters.
def foo_generic(x):
f = x ** 2
idx = np.stack(np.meshgrid(*[np.arange(e) for e in x.shape], indexing='ij'))
idx = tuple(np.concatenate((idx, idx)).reshape(2*x.ndim, -1))
df = np.zeros([*x.shape, *x.shape])
df[idx] = 2 * x.ravel()
return f, df
Note that foo_generic does not support scalar and it would be very inefficient to use it for that anyway, but you can add a condition in it to support this special case apart.
The df matrix will very quickly be huge for higher order so I strongly advise you not to use dense matrices for that since the number of zeros is huge compared to the number of values in the matrix case already. Sparse matrices fix this. In fact, for a 5x5 matrix, there are >95% of zeros. Not to mention the matrix becomes quickly huge and willing a huge matrix full of zeros is not efficient.

How can I scale a set of 2D arrays (3D array) by a 2D array in a vectorized way using NumPy?

I have a 3D matrix containing N x N covariance matrices for M channels [M x N x N]. I also have a 2D matrix of scaling factors for each channel at a series of time points [M x T]. I want to produce a 4D matrix containing a scaled version of the relevant channel's covariance at each time point. So to be clear, [M x T] * [M x N x N] -> [M x T x N x N]
Current version using for loops:
m, t, n = 4, 10, 7
channel_timeseries = np.zeros((m, t))
covariances = np.random.rand(m, n, n)
result_array = np.zeros((m, t, n, n))
# Each channel
for i, (channel_cov, channel_timeseries) in enumerate(zip(covariances, channel_timeseries)):
# Each time point
for j, time_point in enumerate(channel_timeseries):
result_array[i, j] = time_point * channel_cov
This should lead to the result array being all zeros. Replacing the initialisation of the channel_timeseries with np.ones, we should see the covariance for each channel replicated unchanged at every step of the time series.
The case which actually matters to me is one in which every channel has a scalar value at every time point and we scale the covariance matrix for the relevant channel by the value matching the correct channel and time point.
As you can see above, I can do this with a for loop and it works completely fine, but I'm working with some huge datasets and it would be better to have a vectorised solution.
Many thanks for your time.
You can use np.einsum, as b-fg said
np.einsum('mt,mno->mtno', channel_timeseries, covariances)
or Broadcasting:
channel_timeseries[:, :, None, None] * covariances[:, None, :, :]
numpy.einsum will come handy here. I have modified your code with a random channel_timeseries array, increased the arrays size, and renamed the loop variables (otherwise you overwrite the original ones!)
import numpy as np
import time
m, t, n = 40, 100, 70
channel_timeseries = np.random.rand(m, t)
covariances = np.random.rand(m, n, n)
t0 = time.time()
result_array_1 = np.zeros((m, t, n, n))
# Each channel
for i, (c_cov, c_ts) in enumerate(zip(covariances, channel_timeseries)):
# Each time point
for j, time_point in enumerate(c_ts):
result_array_1[i, j] = time_point * c_cov
t1 = time.time()
result_array_2 = np.einsum('ij,ikl->ijkl', channel_timeseries, covariances)
t2 = time.time()
print(np.array_equal(result_array_1, result_array_2)) # True
print('Time for result_array_1: ', t1-t0) # 0.07601261138916016
print('Time for result_array_2: ', t2-t1) # 0.02957916259765625
This results in a speed increase of more than 50% with numpy.einsum in my machine.

Array Assignment with Autograd ( not again :( )

I have a fairly non trivial function I want to differentiate with autograd but I'm not quite enough of a numpy wizard to figure our how to do it without array assingment.
I also apologize that I had to make this example incredibly contrived and meaningless to be able to run standalone. The actual code I'm working with is for non linear finite elements and is trying to compute the jacobian for a complex non linear system.
import autograd.numpy as anp
from autograd import jacobian
def alpha(x):
return anp.exp(-(x - 10) ** 2) / (x + 1)
def f(x):
# Matrix getting constructed
k = anp.zeros((x.shape[0], x.shape[0]))
# loop over some random 3 dimensional vectors
for element in anp.random.randint(0, x.shape[0], (x.shape[0], 3)):
# select 3 values from x
x_ijk = anp.array([[x[i] for i in element]])
norm = anp.linalg.norm(
x_ijk # anp.vstack((element, element)).transpose()
)
# make some matrix from the element
m = element.reshape(3, 1) # element.reshape(1, 3)
# alpha is an arbitrary differentiable function R -> R
alpha_value = alpha(norm)
# combine m matricies into k scaling by alpha_value
n = m.shape[0]
for i in range(n):
for j in range(n):
k[element[i], element[j]] += m[i, j] * alpha_value
return k # x
print(jacobian(f)(anp.random.rand(10)))
# And course we get an error
# k[element[i], element[j]] += m[i, j] * alpha_value
# ValueError: setting an array element with a sequence.
I don't really understand this message since no type error is happening. I assume it must be from assignment.
After writting the above I made a trivial switch to PyTorch and the code runs just fine. But I would still prefer to use autograd
#pytorch version
import torch
from torch.autograd.gradcheck import zero_gradients
def alpha(x):
return torch.exp(x)
def f(x):
# Matrix getting constructed
k = torch.zeros((x.shape[0], x.shape[0]))
# loop over some random 3 dimensional vectors
for element in torch.randint(0, x.shape[0], (x.shape[0], 3)):
# select 3 values from x
x_ijk = torch.tensor([[1. if n == e else 0 for n in range(len(x))] for e in element]) # x
norm = torch.norm(
x_ijk # torch.stack((torch.tanh(element.float() + 4), element.float() - 4)).t()
)
m = torch.rand(3, 3)
# alpha is an arbitrary differentiable function R -> R
alpha_value = alpha(norm)
n = m.shape[0]
for i in range(n):
for j in range(n):
k[element[i], element[j]] += m[i, j] * alpha_value
print(k)
return k # x
x = torch.rand(4, requires_grad=True)
print(x, '\n')
y = f(x)
print(y, '\n')
grads = []
for val in y:
val.backward(retain_graph=True)
grads.append(x.grad.clone())
zero_gradients(x)
if __name__ == '__main__':
print(torch.stack(grads))
In Autograd, and in JAX, you are not allowed to perform array indexing assignments. See the JAX gotchas for a partial explanation of this.
PyTorch allows this functionality. If you want to run your code in autograd, you'll have to find a way to remove the offending line k[element[i], element[j]] += m[i, j] * alpha_value. If you are okay with running your code in JAX (which has essentially the same syntax as autograd, but more features), then it looks like jax.ops could be helpful for performing this sort of indexing assignment.

Vectorized syntax for creating a sequence of block matrices in NumPy

I have two 3D arrays A and B with shapes (k, n, n) and (k, m, m) respectively. I would like to create a matrix C of shape (k, n+m, n+m) such that for each 0 <= i < k, the 2D matrix C[i,:,:] is the block diagonal matrix obtained by putting A[i, :, :] at the upper left n x n part and B[i, :, :] at the lower right m x m part.
Currently I am using the following to achieve this is NumPy:
C = np.empty((k, n+m, n+m))
for i in range(k):
C[i, ...] = np.block([[A[i,...], np.zeros((n,m))],
[np.zeros((m,n)), B[i,...]]])
I was wondering if there is a way to do this without the for loop. I think if k is large my solution is not very efficient.
IIUC You can simply slice and assign -
C = np.zeros((k, n+m, n+m),dtype=np.result_type(A,B))
C[:,:n,:n] = A
C[:,n:,n:] = B

Inverse of numpy.dot

I can easily calculate something like:
R = numpy.column_stack([A,np.ones(len(A))])
M = numpy.dot(R,[k,m0])
where A is a simple array and k,m0 are known values.
I want something different. Having fixed R, M and k, I need to obtain m0.
Is there a way to calculate this by an inverse of the function numpy.dot()?
Or it is only possible by rearranging the matrices?
M = numpy.dot(R,[k,m0])
is performing matrix multiplication. M = R * x.
So to compute the inverse, you could use np.linalg.lstsq(R, M):
import numpy as np
A = np.random.random(5)
R = np.column_stack([A,np.ones(len(A))])
k = np.random.random()
m0 = np.random.random()
M = R.dot([k,m0])
(k_inferred, m0_inferred), residuals, rank, s = np.linalg.lstsq(R, M)
assert np.allclose(m0, m0_inferred)
assert np.allclose(k, k_inferred)
Note that both k and m0 are determined, given M and R (assuming len(M) >= 2).

Categories

Resources