README.md 4.29 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
26
27
## Installation


### Requirements: 
 * Cuda >= 10.1 (with nvcc compiler)
28
 * PyTorch >= 1.7
zachteed's avatar
zachteed committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

We recommend installing within a virtual enviornment. Make sure you clone using the `--recursive` flag. If you are using Anaconda, the following command can be used to install all dependencies
```
git clone --recursive https://github.com/princeton-vl/lietorch.git
cd lietorch

conda create -n lie_env
conda activate lie_env
conda install scipy pyyaml pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch
```

To run the examples, you will need OpenCV and Open3D. Depending on your operating system, OpenCV and Open3D can either be installed with pip or may need to be built from source
```
pip install opencv-python open3d
```

45
### Installing (from source)
zachteed's avatar
zachteed committed
46
47
48
49
50
51
52
53

Clone the repo using the `--recursive` flag and install using `setup.py` (may take up to 10 minutes)
```
git clone --recursive https://github.com/princeton-vl/lietorch.git
python setup.py install
./run_tests.sh
```

54
55
56
57
58
59
60
61
### Installing (pip)
You can install the library directly using pip
```bash
pip install git+https://github.com/princeton-vl/lietorch.git
```



zachteed's avatar
zachteed committed
62
63
64
65
66
67
68
69
70
71
72
## 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 |

73
Each group supports the following differentiable operations:
zachteed's avatar
zachteed committed
74
75
76
77
78
79
80
81
82

| Operation | Map | Description |
| -------| --------| ------------- |
| exp    | g -> G | exponential map |
| log    | G -> g | logarithm map |
| inv    | G -> G | group inverse |
| mul   | G x G -> G | group multiplication |
| adj    | G x g -> g | adjoint |
| adjT   | G x g*-> g* | dual adjoint |
83
84
85
86
87
88
89
| 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
90
91

&nbsp;
92
### Simple Example:
zachteed's avatar
zachteed committed
93
94
95
96
97
98
99
100
101
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)

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

105
# 8000x8000 matrix of angles
zachteed's avatar
zachteed committed
106
107
108
109
110
111
112
ang = dR.log().norm(dim=-1)

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

113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

### 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
134
135
136
137
138
139
140
141
## 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).