se3cnn

Author: mariogeiger

README.md

@@ -7,18 +7,18 @@ This library aims to create SE(3) equivariant convolutional neural networks.
 
 ## Hierarchy
 
-- `se3_cnn` contains the library
-  - `se3_cnn/convolution.py` defines `SE3Convolution` the main class of the library
-  - `se3_cnn/blocks` defines ways of introducing non linearity in an equivariant way
-  - `se3_cnn/batchnorm.py` equivariant batch normalization
-  - `se3_cnn/groupnorm.py` equivariant group normalization
-  - `se3_cnn/dropout.py` equivariant dropout
+- `se3cnn` contains the library
+  - `se3cnn/convolution.py` defines `SE3Convolution` the main class of the library
+  - `se3cnn/blocks` defines ways of introducing non linearity in an equivariant way
+  - `se3cnn/batchnorm.py` equivariant batch normalization
+  - `se3cnn/groupnorm.py` equivariant group normalization
+  - `se3cnn/dropout.py` equivariant dropout
 - `experiments` contains experiments made with the library
 - `examples` simple scripts
 
 ## Dependencies
 
-- [pytorch](https://pytorch.org) 
+- [pytorch](https://pytorch.org)
 - [lie_learn](https://github.com/AMLab-Amsterdam/lie_learn) is required to compute the irreducible representations of SO(3)
 - scipy
 

examples/example.py

@@ -16,7 +16,7 @@
 
 # The class GatedBlock inherit from the class torch.nn.Module.
 # It contains one convolution, some ReLU and multiplications
-from se3_cnn.blocks import GatedBlock
+from se3cnn.blocks import GatedBlock
 
 
 class AvgSpacial(nn.Module):

examples/plots/kernels.py

@@ -4,10 +4,10 @@
 import matplotlib.pyplot as plt
 from matplotlib import cm
 from mpl_toolkits.mplot3d import Axes3D  # pylint: disable=W
-from se3_cnn.SO3 import irr_repr, spherical_harmonics
-from se3_cnn.basis_kernels import _basis_transformation_Q_J
-from se3_cnn.util.cache_file import cached_dirpklgz
-from se3_cnn.SO3 import compose
+from se3cnn.SO3 import irr_repr, spherical_harmonics
+from se3cnn.basis_kernels import _basis_transformation_Q_J
+from se3cnn.util.cache_file import cached_dirpklgz
+from se3cnn.SO3 import compose
 
 
 def beta_alpha(n):
@@ -87,9 +87,9 @@ def main():
     # f(r^-1 x)
 
     f = np.einsum(
-        "ij,zjkba,kl->zilba", 
-        irr_repr(args.order_out, args.alpha, args.beta, args.gamma), 
-        f, 
+        "ij,zjkba,kl->zilba",
+        irr_repr(args.order_out, args.alpha, args.beta, args.gamma),
+        f,
         irr_repr(args.order_in, -args.gamma, -args.beta, -args.alpha)
     )
     # rho_out(r) f(r^-1 x) rho_in(r^-1)

examples/plots/kernels_cutaway.py

@@ -1,6 +1,6 @@
 # pylint: disable=C,R,E1101
-from se3_cnn import basis_kernels
-from se3_cnn import SO3
+from se3cnn import basis_kernels
+from se3cnn import SO3
 import numpy as np
 import matplotlib.pyplot as plt
 from functools import partial

examples/tetris.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from se3_cnn.blocks import GatedBlock
+from se3cnn.blocks import GatedBlock
 
 
 class AvgSpacial(nn.Module):
@@ -86,8 +86,8 @@ def __init__(self):
         super(SE3Net, self).__init__()
         features = [
             (1,),
-            (2, 2, 2, 2), 
-            (4, 4, 4, 4), 
+            (2, 2, 2, 2),
+            (4, 4, 4, 4),
             (16,)
         ]
         common_block_params = {

experiments/datasets/modelnet/modelnet.py

@@ -8,7 +8,7 @@
 from functools import partial
 
 from scipy.ndimage import affine_transform
-from se3_cnn.SO3 import rot
+from se3cnn.SO3 import rot
 
 
 
@@ -43,7 +43,7 @@ class ModelNet(torch.utils.data.Dataset):
     def __init__(self, root_dir, dataset, mode, size, classes, transform=None, target_transform=None):
         '''
         :param root: directory to store dataset in
-        :param dataset: 
+        :param dataset:
         :param mode: dataset to load: 'train', 'validation', 'test' or 'train_full'
                      the validation set is split from the train set, the full train set can be accessed via 'train_full'
                     :param transform: transformation applied to image in __getitem__
@@ -91,7 +91,7 @@ def __len__(self):
 
 
 class AddZAxis(object):
-    ''' add z-axis as second channel to volume 
+    ''' add z-axis as second channel to volume
         the scale of the z-axis can be set freely
         if the volume tensor does not contain a channel dimension, add it
         the z axis is assumed to be the last axis of the volume tensor

experiments/scripts/MRI/networks/MICCAI2012/SE3BasicModel/SE3BasicModel.py

@@ -4,8 +4,8 @@
 
 from functools import partial
 
-from se3_cnn.blocks import GatedBlock
-from se3_cnn import basis_kernels
+from se3cnn.blocks import GatedBlock
+from se3cnn import basis_kernels
 
 
 class network(nn.Module):

experiments/scripts/MRI/networks/MICCAI2012/SE3UNet/SE3UNet.py

@@ -4,8 +4,8 @@
 
 from functools import partial
 
-from se3_cnn.blocks import GatedBlock
-from se3_cnn import basis_kernels
+from se3cnn.blocks import GatedBlock
+from se3cnn import basis_kernels
 from experiments.util.arch_blocks import Merge
 
 

experiments/scripts/MRI/networks/MRBrainS/SE3BasicModel/SE3BasicModel.py

@@ -4,8 +4,8 @@
 
 from functools import partial
 
-from se3_cnn.blocks import GatedBlock
-from se3_cnn import basis_kernels
+from se3cnn.blocks import GatedBlock
+from se3cnn import basis_kernels
 
 
 class network(nn.Module):

experiments/scripts/MRI/networks/MRBrainS/SE3UNet/SE3UNet.py

@@ -4,7 +4,7 @@
 
 from functools import partial
 
-from se3_cnn import basis_kernels
+from se3cnn import basis_kernels
 from experiments.util.arch_blocks import NonlinearityBlock
 from experiments.util.arch_blocks import SkipSumBlock
 
@@ -31,7 +31,7 @@ def __init__(self, output_size, args):
             'batch_norm_momentum': 0.01,
         }
 
-        features = [(3,), # in 
+        features = [(3,), # in
                     [(6,6,6,4), (16,16,16,12)], # level 1 (enc and dec)
                     ( 32, 32, 32, 24), # level 2 (enc and dec)
                     ( 64, 64, 64, 48), # level 3 (bridge)