The group SE(3) is the group of 3 dimensional rotations and translations. This library aims to create SE(3) equivariant convolutional neural networks.
The code is separated in two parts:
import torch
from se3cnn.image.convolution import SE3Convolution
size = 32 # space size
scalar_field = torch.randn(1, 1, size, size, size) # [batch, _, x, y, z]
Rs_in = [(1, 0)] # 1 scalar field
Rs_out = [(1, 1)] # 1 vector field
conv = SE3Convolution(Rs_in, Rs_out, size=5)
# conv.weight.size() == [2] (2 radial degrees of freedom)
vector_field = conv(scalar_field) # [batch, vector component, x, y, z]
# vector_field.size() == [1, 3, 28, 28, 28]
from functools import partial
import torch
from se3cnn.point.radial import CosineBasisModel
from se3cnn.point.kernel import Kernel
from se3cnn.point.operations import Convolution
from se3cnn.util.plot import plot_sh_signal
import matplotlib.pyplot as plt
# Radial model: R -> R^d
# Projection on cos^2 basis functions followed by a fully connected network
RadialModel = partial(CosineBasisModel, max_radius=3.0, number_of_basis=3, h=100, L=1, act=torch.relu)
# kernel: composed on a radial part that contains the learned parameters
# and an angular part given by the spherical hamonics and the Clebsch-Gordan coefficients
K = partial(Kernel, RadialModel=RadialModel)
# Use the kernel to define a convolution operation
C = partial(Convolution, K)
Rs_in = [(1, 0)] # one scalar
Rs_out = [(1, l) for l in range(10)]
conv = C(Rs_in, Rs_out)
n = 3 # number of points
features = torch.ones(1, n)
geometry = torch.randn(n, 3)
features = conv(features, geometry)
plt.figure(figsize=(4, 4))
plot_sh_signal(features[:, 0], n=50)
plt.gca().view_init(azim=0, elev=45)
se3cnn
contains the libraryse3cnn/SO3.py
defines all the needed mathematical functionsse3cnn/image
contains the code specific to voxelsse3cnn/point
contains the code specific to pointsse3cnn/non_linearities
non linearities working for both point and voxel code
examples
simple scripts and experiments
- install pytorch
pip install git+https://github.com/AMLab-Amsterdam/lie_learn
pip install git+https://github.com/mariogeiger/se3cnn
Install with
python setup.py install
@misc{mario_geiger_2019_3348277,
author = {Mario Geiger and
Tess Smidt and
Wouter Boomsma and
Maurice Weiler and
Michał Tyszkiewicz and
Jes Frellsen and
Benjamin K. Miller},
title = {mariogeiger/se3cnn: Point cloud support},
month = jul,
year = 2019,
doi = {10.5281/zenodo.3348277},
url = {https://doi.org/10.5281/zenodo.3348277}
}