activation functions

examples/tetris_point.py

@@ -3,8 +3,9 @@
 
 import torch
 
+from se3cnn.blocks import GatedBlock
 from se3cnn.convolution import SE3PointConvolution
-from se3cnn.blocks.point_gated_block import GatedBlock
+from se3cnn.non_linearities import rescaled_act
 from se3cnn.point_kernel import SE3PointKernel
 from se3cnn.point_radial import CosineBasisModel
 from se3cnn.SO3 import rand_rot
@@ -40,14 +41,14 @@ def __init__(self, num_classes):
         features = [(1,), (2, 2, 2, 1), (4, 4, 4, 4), (6, 4, 4, 0), (64,)]
         self.num_features = len(features)
 
-        RadialModel = partial(CosineBasisModel, 3.0, 3, 10, 100)
+        sp = rescaled_act.Softplus(beta=5)
+
+        RadialModel = partial(CosineBasisModel, max_radius=3.0, number_of_basis=3, h=100, L=50, act=sp)
         Kernel = partial(SE3PointKernel, RadialModel=RadialModel)
         Convolution = partial(SE3PointConvolution, Kernel)
 
-        # note: torch.randn(100000).sigmoid().mul(1.85).pow(2).mean() == 1
-        # note: torch.randn(100000).tanh().mul(1.6).pow(2).mean() == 1
         self.layers = torch.nn.ModuleList([
-            GatedBlock(features[i], features[i+1], lambda x: x.tanh().mul(1.6), lambda x: x.sigmoid().mul(1.85), Convolution)
+            GatedBlock(features[i], features[i+1], sp, rescaled_act.sigmoid, Convolution)
             for i in range(len(features) - 1)
         ])
         self.layers += [AvgSpacial(), torch.nn.Linear(64, num_classes)]

se3cnn/non_linearities/rescaled_act.py

@@ -0,0 +1,24 @@
+# pylint: disable=no-member, missing-docstring, invalid-name
+import torch
+
+
+class Softplus:
+    def __init__(self, beta):
+        x = torch.randn(100000, dtype=torch.float64)
+        self.factor = torch.nn.functional.softplus(x, beta).pow(2).mean().rsqrt().item()
+        self.beta = beta
+
+    def __call__(self, x):
+        return torch.nn.functional.softplus(x, self.beta).mul(self.factor)
+
+
+def sigmoid(x):
+    return x.sigmoid().mul(1.84623)
+
+
+def tanh(x):
+    return x.tanh().mul(1.59254)
+
+
+def relu(x):
+    return x.relu().mul(2 ** 0.5)

se3cnn/point_radial.py

@@ -6,12 +6,12 @@
 
 
 class FiniteElementModel(torch.nn.Module):
-    def __init__(self, position, basis, Model, out_dim):
+    def __init__(self, out_dim, position, basis, Model):
         '''
+        :param out_dim: output dimension
         :param position: tensor [i, ...]
         :param basis: scalar function: tensor [a, ...] -> [a]
         :param Model: Class(d1, d2), trainable model: R^d1 -> R^d2
-        :param out_dim: output dimension
         '''
         super().__init__()
         self.register_buffer('position', position)
@@ -66,17 +66,13 @@ def forward(self, x):
         return x
 
 
-def GaussianBasisModel(max_radius, number_of_basis, number_of_hidden_layers, number_of_hidden_neurons, out_dim):
-    radii = torch.linspace(0, max_radius, steps=number_of_basis)
-    sigma = 0.45 * (radii[1] - radii[0])
-    basis = lambda x: x.pow(2).div(2 * sigma ** 2).neg().exp()
-    Model = partial(FC, h=number_of_hidden_neurons, L=number_of_hidden_layers, act=lambda x: x.relu().mul(2 ** 0.5))
-    return FiniteElementModel(radii, basis, Model, out_dim)
+def FiniteElementFCModel(out_dim, position, basis, h, L, act):
+    Model = partial(FC, h=h, L=L, act=act)
+    return FiniteElementModel(out_dim, position, basis, Model)
 
 
-def CosineBasisModel(max_radius, number_of_basis, number_of_hidden_layers, number_of_hidden_neurons, out_dim):
+def CosineBasisModel(out_dim, max_radius, number_of_basis, h, L, act):
     radii = torch.linspace(0, max_radius, steps=number_of_basis)
     step = radii[1] - radii[0]
     basis = lambda x: x.div(step).add(1).relu().sub(2).neg().relu().add(1).mul(math.pi / 2).cos().pow(2)
-    Model = partial(FC, h=number_of_hidden_neurons, L=number_of_hidden_layers, act=lambda x: x.relu().mul(2 ** 0.5))
-    return FiniteElementModel(radii, basis, Model, out_dim)
+    return FiniteElementFCModel(out_dim, radii, basis, h, L, act)