Batched matrices from 1D tensors

Hi guys, Let’s say we want to generate Rotation Matrices from Batch of angles ( Bx1), we could use the torch.cos() and torch.sin() to convert (Bx1) angles to (Bx1) tensors of cos and sin.

How I can stack the given angles to produce the batched version of Rotation Matrices.
Here is an example:

[ [ [cos(ang(0)), -sin(ang(0))], [sin(ang(0)), cos(ang(0))] ], [ [cos(ang(1)), -sin(ang(1))], [sin(ang(1)), cos(ang(1))] ], .... ]

Hi @Paran0, this code does the work.

import torch

n = 10 #from 0 upto 9 degrees
thetas = np.arange(n, dtype = np.float32)

cosine = torch.from_numpy(np.cos(np.deg2rad(thetas)))
sine = torch.from_numpy(np.sin(np.deg2rad(thetas)))

batch = torch.zeros((n,2,2), dtype = cosine.dtype)

batch[:,0,0] = cosine
batch[:,0,1] = -sine
batch[:,1,0] = sine
batch[:,1,1] = cosine
print(batch.shape)
1 Like

Yeah, that’s great, I tried this method too, any suggestion for performance improvements? for example, somehow we omit the pre-allocations …

Yes, you can do that.!
But your tensor shape will be (2,2, N).

batch = np.asarray([[cosine, -sine],
 					[sine, cosine]])