README.md 4.12 KB
Newer Older
Jia Deng's avatar
Jia Deng committed
1
# LieTorch: Tangent Space Backpropagation
zachteed's avatar
zachteed committed
2
3
4
5


## Introduction

Jia Deng's avatar
Jia Deng committed
6
The LieTorch library generalizes PyTorch to 3D transformation groups. Just as `torch.Tensor` is a multi-dimensional matrix of scalar elements, `lietorch.SE3` is a multi-dimensional matrix of SE3 elements. We support common tensor manipulations such as indexing, reshaping, and broadcasting. Group operations can be composed into computation graphs and backpropagation is automatically peformed in the tangent space of each element. For more details, please see our paper:
zachteed's avatar
zachteed committed
7
8
9

<center><img src="lietorch.png" width="480" style="center"></center>

Jia Deng's avatar
Jia Deng committed
10
[Tangent Space Backpropagation for 3D Transformation Groups](https://arxiv.org/pdf/2103.12032.pdf)  
zachteed's avatar
zachteed committed
11
12
Zachary Teed and Jia Deng, CVPR 2021

Jia Deng's avatar
Jia Deng committed
13
14
```
@inproceedings{teed2021tangent,
Jia Deng's avatar
Jia Deng committed
15
  title={Tangent Space Backpropagation for 3D Transformation Groups},
Jia Deng's avatar
Jia Deng committed
16
  author={Teed, Zachary and Deng, Jia},
Jia Deng's avatar
Jia Deng committed
17
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
Jia Deng's avatar
Jia Deng committed
18
19
20
21
22
  year={2021},
}
```


zachteed's avatar
zachteed committed
23
24
25
## Installation


Zach Teed's avatar
Zach Teed committed
26
### Installing (from source):
zachteed's avatar
zachteed committed
27

Zach Teed's avatar
Zach Teed committed
28
29
30
Requires torch >= 2 and CUDA >= 11. Tested up to torch==2.7 and CUDA 12. Make sure PyTorch and CUDA major versions match. 

```bash
zachteed's avatar
zachteed committed
31
32
33
git clone --recursive https://github.com/princeton-vl/lietorch.git
cd lietorch

Zach Teed's avatar
Zach Teed committed
34
35
python3 -m venv .venv
source .venv/bin/activate
zachteed's avatar
zachteed committed
36

Zach Teed's avatar
Zach Teed committed
37
38
39
40
41
42
43
# install requirements
pip install torch torchvision torchaudio wheel

# optional: specify GPU architectures
export TORCH_CUDA_ARCH_LIST="7.5;8.6;8.9;9.0"

# install lietorch
Zach Teed's avatar
Zach Teed committed
44
pip install --no-build-isolation .
zachteed's avatar
zachteed committed
45
46
```

Zach Teed's avatar
Zach Teed committed
47
48
49
50
### Installing (with pip)
```bash
# optional: specify GPU architectures
export TORCH_CUDA_ARCH_LIST="7.5;8.6;8.9;9.0"
zachteed's avatar
zachteed committed
51

Zach Teed's avatar
Zach Teed committed
52
pip install --no-build-isolation git+https://github.com/princeton-vl/lietorch.git
zachteed's avatar
zachteed committed
53
```
Zach Teed's avatar
Zach Teed committed
54
55
56
57
58


To run the examples, you will need these additional libraries
```bash
pip install opencv-python open3d scipy pyyaml
zachteed's avatar
zachteed committed
59
60
```

Zach Teed's avatar
Zach Teed committed
61
62
63
### Running Tests

After building, you can run the tests
64
```bash
Zach Teed's avatar
Zach Teed committed
65
./run_tests.sh
66
67
68
69
```



zachteed's avatar
zachteed committed
70
71
72
73
74
75
76
77
78
79
80
## Overview

LieTorch currently supports the 3D transformation groups. 

| Group  | Dimension | Action |
| -------| --------- | ------------- |
| SO3    | 3  | rotation |
| RxSO3  | 4  | rotation + scaling |
| SE3    | 6  | rotation + translation |
| Sim3   | 7  | rotation + translation + scaling |

81
Each group supports the following differentiable operations:
zachteed's avatar
zachteed committed
82
83
84
85
86
87

| Operation | Map | Description |
| -------| --------| ------------- |
| exp    | g -> G | exponential map |
| log    | G -> g | logarithm map |
| inv    | G -> G | group inverse |
Zach Teed's avatar
Zach Teed committed
88
| mul    | G x G -> G | group multiplication |
zachteed's avatar
zachteed committed
89
90
| adj    | G x g -> g | adjoint |
| adjT   | G x g*-> g* | dual adjoint |
91
92
93
94
95
96
97
| act    | G x R^3 -> R^3 | action on point (set) |
| act4   | G x P^3 -> P^3 | action on homogeneous point (set) |
| matrix | G -> R^{4x4} | convert to 4x4 matrix
| vec    | G -> R^D | map to Euclidean embedding vector |
| InitFromVec | R^D -> G | initialize group from Euclidean embedding


zachteed's avatar
zachteed committed
98
99

&nbsp;
100
### Simple Example:
zachteed's avatar
zachteed committed
101
102
103
104
105
106
107
108
109
Compute the angles between all pairs of rotation matrices

```python
import torch
from lietorch import SO3

phi = torch.randn(8000, 3, device='cuda', requires_grad=True)
R = SO3.exp(phi)

110
# relative rotation matrix, SO3 ^ {8000 x 8000}
zachteed's avatar
zachteed committed
111
112
dR = R[:,None].inv() * R[None,:]

113
# 8000x8000 matrix of angles
zachteed's avatar
zachteed committed
114
115
116
117
118
119
120
ang = dR.log().norm(dim=-1)

# backpropogation in tangent space
loss = ang.sum()
loss.backward()
```

121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

### Converting between Groups Elements and Euclidean Embeddings
We provide differentiable `FromVec` and `ToVec` functions which can be used to convert between LieGroup elements and their vector embeddings. Additional, the `.matrix` function returns a 4x4 transformation matrix.
```python

# random quaternion
q = torch.randn(1, 4, requires_grad=True)
q = q / q.norm(dim=-1, keepdim=True)

# create SO3 object from quaternion (differentiable w.r.t q)
R = SO3.InitFromVec(q)

# 4x4 transformation matrix (differentiable w.r.t R)
T = R.matrix()

# map back to quaterion (differentiable w.r.t R)
q = R.vec()

```


zachteed's avatar
zachteed committed
142
143
144
145
146
147
148
149
## Examples
We provide real use cases in the examples directory
1. Pose Graph Optimization
2. Deep SE3/Sim3 Registrtion
3. RGB-D SLAM / VO

### Acknowledgements
Many of the Lie Group implementations are adapted from [Sophus](https://github.com/strasdat/Sophus).