Skip to content

added cuda support for all pytorch mlg operations #25

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
angadbajwa wants to merge 1 commit into
base: main
Choose a base branch
from
Open

added cuda support for all pytorch mlg operations #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 15 additions & 12 deletions pymlg/torch/se2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,21 @@ class SE2(MatrixLieGroupTorch):
matrix_size = 3

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
phi = torch.rand(N, 1, 1) * 2 * torch.pi
r = torch.randn(N, 2, 1)
C = SO2.Exp(phi)
return SE2.from_components(C, r)
return SE2.from_components(C, r).to(device)

@staticmethod
def from_components(C, r):
"""
Construct an SE(2) matrix from a rotation matrix and translation vector.
"""
T = torch.zeros(C.shape[0], 3, 3, dtype=C.dtype)

# first, confirm that both components are allocated on the same device
assert C.device == r.device, "Components must be on the same device for SE2.from_components"
T = torch.zeros(C.shape[0], 3, 3, dtype=C.dtype, device=C.device)

T[:, 0:2, 0:2] = C
T[:, 0:2, 2] = r.view(-1, 2)
Expand All @@ -45,7 +48,7 @@ def wedge(xi):
phi = xi[:, 0]
xi_r = xi[:, 1:]
Xi_phi = SO2.wedge(phi)
Xi = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype)
Xi = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype, device=xi.device)
Xi[:, 0:2, 0:2] = Xi_phi
Xi[:, 0:2, 2] = xi_r.view(-1, 2)
return Xi
Expand All @@ -72,17 +75,17 @@ def log(T):
Xi_phi = SO2.log(T[:, 0:2, 0:2])
r = T[:, 0:2, 2].unsqueeze(2)
xi_r = SE2.V_matrix_inv(SO2.vee(Xi_phi)) @ r
Xi = torch.zeros(T.shape[0], 3, 3, dtype=T.dtype)
Xi = torch.zeros(T.shape[0], 3, 3, dtype=T.dtype, device=T.device)
Xi[:, 0:2, 0:2] = Xi_phi
Xi[:, 0:2, 2] = xi_r.squeeze(2)
return Xi

@staticmethod
def odot(b):

X = torch.zeros(b.shape[0], 3, 3, dtype=b.dtype)
X = torch.zeros(b.shape[0], 3, 3, dtype=b.dtype, device=b.device)
X[:, 0:2, 0] = SO2.odot(b[:, :2]).squeeze(2)
X[:, 0:2, 1:3] = batch_eye(b.shape[0], 2, 2) * b[:, 2].unsqueeze(2)
X[:, 0:2, 1:3] = batch_eye(b.shape[0], 2, 2, device=b.device) * b[:, 2].unsqueeze(2)

return X

Expand All @@ -103,7 +106,7 @@ def left_jacobian(xi):
large_angle_mask = small_angle_mask.logical_not()
large_angle_inds = large_angle_mask.nonzero(as_tuple=True)[0]

J = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype)
J = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype, device=xi.device)

if small_angle_inds.numel():
A = (1 - 1.0 / 6.0 * phi_sq[small_angle_inds]).view(-1)
Expand Down Expand Up @@ -144,7 +147,7 @@ def adjoint(T):
r = T[:, 0:2, 2]

# build Om matrix manually (will this break the DAG?)
Om = torch.Tensor([[0, -1], [1, 0]]).repeat(T.shape[0], 1, 1)
Om = torch.Tensor([[0, -1], [1, 0]], device=T.device).repeat(T.shape[0], 1, 1)

A = torch.zeros(T.shape[0], 3, 3, dtype=T.dtype)
A[:, 0, 0] = 1
Expand All @@ -155,7 +158,7 @@ def adjoint(T):

@staticmethod
def adjoint_algebra(Xi):
A = torch.zeros(Xi.shape[0], 3, 3, dtype=Xi.dtype)
A = torch.zeros(Xi.shape[0], 3, 3, dtype=Xi.dtype, device=Xi.device)
A[:, 1, 0] = Xi[:, 1, 2]
A[:, 2, 0] = -Xi[:, 0, 2]
A[:, 1:, 1:] = Xi[:, 0:2, 0:2]
Expand All @@ -171,7 +174,7 @@ def V_matrix(phi):
large_angle_mask = small_angle_mask.logical_not()
large_angle_inds = large_angle_mask.nonzero(as_tuple=True)[0]

V = batch_eye(phi.shape[0], 2, 2, dtype=phi.dtype)
V = batch_eye(phi.shape[0], 2, 2, device=phi.device, dtype=phi.dtype)

if small_angle_inds.numel():
V[small_angle_inds] += .5 * SO2.wedge(phi[small_angle_inds])
Expand All @@ -193,7 +196,7 @@ def V_matrix_inv(phi):
large_angle_mask = small_angle_mask.logical_not()
large_angle_inds = large_angle_mask.nonzero(as_tuple=True)[0]

V_inv = batch_eye(phi.shape[0], 2, 2, dtype=phi.dtype)
V_inv = batch_eye(phi.shape[0], 2, 2, device=phi.device, dtype=phi.dtype)

if small_angle_inds.numel():
V_inv[small_angle_inds] -= .5 * SO2.wedge(phi[small_angle_inds])
Expand Down
29 changes: 16 additions & 13 deletions pymlg/torch/se23.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class SE23(MatrixLieGroupTorch):
matrix_size = 5

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
"""
Generates a random batch of SE_2(3) matricies.

Expand All @@ -39,7 +39,7 @@ def random(N=1):

C = SO3.Exp(phi)

return SE23.from_components(C, v, r)
return SE23.from_components(C, v, r).to(device)

@staticmethod
def from_components(C: torch.Tensor, v: torch.Tensor, r: torch.Tensor):
Expand All @@ -65,8 +65,11 @@ def from_components(C: torch.Tensor, v: torch.Tensor, r: torch.Tensor):
# firstly, check that batch dimension for all 3 components matches
if not (C.shape[0] == v.shape[0] == r.shape[0]):
raise ValueError("Batch dimension for SE_2(3) components don't match.")

# check that all components are on the same device
assert C.device == v.device == r.device, "Components must be on the same device for SE23.from_components."

X = batch_eye(C.shape[0], 5, 5, dtype=C.dtype)
X = batch_eye(C.shape[0], 5, 5, device=C.device, dtype=C.dtype)

X[:, 0:3, 0:3] = C
X[:, 0:3, 3] = v.squeeze(2)
Expand Down Expand Up @@ -112,7 +115,7 @@ def wedge(xi: torch.Tensor):
Xi = torch.cat(
(SO3.cross(xi_phi), xi_v, xi_r), dim=2
) # this yields a (N, 3, 5) matrix that must now be blocked with a (2, 5) batched matrix
block = torch.zeros(xi_phi.shape[0], 2, 5)
block = torch.zeros(xi_phi.shape[0], 2, 5, device=xi.device, dtype=xi.dtype)
return torch.cat((Xi, block), dim=1)

@staticmethod
Expand Down Expand Up @@ -152,13 +155,13 @@ def log(X):
Xi = torch.cat(
(SO3.cross(phi), v, r), dim=2
) # this yields a (N, 3, 5) matrix that must now be blocked with a (2, 5) batched matrix
block = torch.zeros(X.shape[0], 2, 5)
block = torch.zeros(X.shape[0], 2, 5, device=X.device, dtype=X.dtype)
return torch.cat((Xi, block), dim=1)

@staticmethod
def adjoint(X):
C, v, r = SE23.to_components(X)
O = torch.zeros(v.shape[0], 3, 3)
O = torch.zeros(v.shape[0], 3, 3, device=X.device, dtype=X.dtype)

# creating block matrix
b1 = torch.cat((C, O, O), dim=2)
Expand All @@ -168,7 +171,7 @@ def adjoint(X):

@staticmethod
def adjoint_algebra(Xi):
A = torch.zeros(Xi.shape[0], 9, 9)
A = torch.zeros(Xi.shape[0], 9, 9, dtype=Xi.dtype, device=Xi.device)
A[:, 0:3, 0:3] = Xi[:, 0:3, 0:3]
A[:, 3:6, 0:3] = SO3.wedge(Xi[:, 0:3, 3])
A[:, 3:6, 3:6] = Xi[:, 0:3, 0:3]
Expand All @@ -177,14 +180,14 @@ def adjoint_algebra(Xi):
return A

@staticmethod
def identity(N=1):
return batch_eye(N, 5, 5)
def identity(device, N=1, dtype=torch.float64):
return batch_eye(N, 5, 5, device=device, dtype=dtype)

@staticmethod
def odot(xi : torch.Tensor):
X = torch.zeros(xi.shape[0], 5, 9)
X = torch.zeros(xi.shape[0], 5, 9, dtype=xi.dtype, device=xi.device)
X[:, 0:4, 0:6] = SE3.odot(xi[:, 0:4])
X[:, 0:3, 6:9] = xi[:, 4] * batch_eye(xi.shape[0], 3, 3)
X[:, 0:3, 6:9] = xi[:, 4] * batch_eye(xi.shape[0], 3, 3, device=xi.device, dtype=xi.dtype)
return X

@staticmethod
Expand All @@ -193,7 +196,7 @@ def left_jacobian(xi):
xi_v = xi[:, 3:6]
xi_r = xi[:, 6:9]

J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype, device=xi.device)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE23._small_angle_tol
Expand Down Expand Up @@ -242,7 +245,7 @@ def left_jacobian_inv(xi):
xi_v = xi[:, 3:6]
xi_r = xi[:, 6:9]

J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype, device=xi.device)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE23._small_angle_tol
Expand Down
29 changes: 16 additions & 13 deletions pymlg/torch/se3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def _left_jacobian_Q_matrix(xi_phi, xi_rho):
return Q.squeeze_()

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
"""
Generates a random batch of SE_(3) matricies.

Expand All @@ -70,7 +70,7 @@ def random(N=1):

C = SO3.Exp(phi)

return SE3.from_components(C, r)
return SE3.from_components(C, r).to(device)

@staticmethod
def from_components(C: torch.Tensor, r: torch.Tensor):
Expand All @@ -94,8 +94,11 @@ def from_components(C: torch.Tensor, r: torch.Tensor):
# firstly, check that batch dimension for all 3 components matches
if not (C.shape[0] == r.shape[0]):
raise ValueError("Batch dimension for SE(3) components don't match.")

# then, check that both components are on the same device
assert C.device == r.device, "Components must be on the same device for SE3.from_components."

X = batch_eye(C.shape[0], 4, 4, dtype = C.dtype)
X = batch_eye(C.shape[0], 4, 4, device=C.device, dtype = C.dtype)

X[:, 0:3, 0:3] = C
X[:, 0:3, 3] = r.squeeze(2)
Expand Down Expand Up @@ -140,7 +143,7 @@ def wedge(xi: torch.Tensor):
) # this yields a (N, 3, 4) matrix that must now be blocked with a (1, 4) batched matrix

# generating a (N, 1, 4) batched matrix to append
b1 = torch.tensor([0, 0, 0, 0], dtype = xi.dtype).reshape(1, 1, 4)
b1 = torch.tensor([0, 0, 0, 0], dtype = xi.dtype, device=xi.device).reshape(1, 1, 4)
block = b1.repeat(Xi.shape[0], 1, 1)

return torch.cat((Xi, block), dim=1)
Expand Down Expand Up @@ -180,22 +183,22 @@ def log(X : torch.Tensor):
) # this yields a (N, 3, 4) matrix that must now be blocked with a (1, 4) batched matrix

# generating a (N, 1, 4) batched matrix to append
b1 = torch.tensor([0, 0, 0, 0]).reshape(1, 1, 4)
b1 = torch.tensor([0, 0, 0, 0], device=X.device).reshape(1, 1, 4)
block = b1.repeat(Xi.shape[0], 1, 1)

return torch.cat((Xi, block), dim=1)

@staticmethod
def odot(b : torch.Tensor):
X = torch.zeros(b.shape[0], 4, 6, dtype=b.dtype)
X = torch.zeros(b.shape[0], 4, 6, dtype=b.dtype, device=b.device)
X[:, 0:3, 0:3] = SO3.odot(b[0:3])
X[:, 0:3, 3:6] = b[:, 3] * batch_eye(b.shape[0], 3, 3, dtype=b.dtype)
X[:, 0:3, 3:6] = b[:, 3] * batch_eye(b.shape[0], 3, 3, dtype=b.dtype, device=b.device)
return X

@staticmethod
def adjoint(X):
C, r = SE3.to_components(X)
O = torch.zeros(r.shape[0], 3, 3)
O = torch.zeros(r.shape[0], 3, 3, device=X.device)

# creating block matrix
b1 = torch.cat((C, O), dim=2)
Expand All @@ -204,22 +207,22 @@ def adjoint(X):

@staticmethod
def adjoint_algebra(Xi):
A = torch.zeros(Xi.shape[0], 6, 6, dtype=Xi.dtype)
A = torch.zeros(Xi.shape[0], 6, 6, dtype=Xi.dtype, device=Xi.device)
A[:, 0:3, 0:3] = Xi[:, 0:3, 0:3]
A[:, 3:6, 0:3] = SO3.wedge(Xi[:, 0:3, 3])
A[:, 3:6, 3:6] = Xi[:, 0:3, 0:3]
return A

@staticmethod
def identity(N=1, dtype=torch.float32):
return batch_eye(N, 4, 4, dtype=dtype)
def identity(device, N=1, dtype=torch.float64):
return batch_eye(N, 4, 4, device=device, dtype=dtype)

@staticmethod
def left_jacobian(xi):
xi_phi = xi[:, 0:3]
xi_r = xi[:, 3:6]

J_left = batch_eye(xi.shape[0], 6, 6, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 6, 6, device=xi.device, dtype=xi.dtype)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE3._small_angle_tol
Expand Down Expand Up @@ -256,7 +259,7 @@ def left_jacobian_inv(xi):
xi_phi = xi[:, 0:3]
xi_r = xi[:, 3:6]

J_left = batch_eye(xi.shape[0], 6, 6, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 6, 6, device=xi.device, dtype=xi.dtype)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE3._small_angle_tol
Expand Down
4 changes: 2 additions & 2 deletions pymlg/torch/so2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ class SO2(MatrixLieGroupTorch):
matrix_size = 2

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
"""
Generates a random batch of SO_(2) matricies.

Expand All @@ -21,7 +21,7 @@ def random(N=1):
batch size, by default 1
"""
phi = torch.rand(N, 1) * 2 * torch.pi
return SO2.Exp(phi)
return SO2.Exp(phi).to(device)

@staticmethod
def wedge(phi):
Expand Down
Loading
Loading