change init scale

april-tools · Jul 18, 2024 · 1324738 · 1324738
1 parent 1b104af
commit 1324738
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/src/layers.py b/src/layers.py
@@ -24,7 +24,11 @@ def __init__(
         assert weight1.num_folds == num_folds
         assert weight2.num_folds == num_folds
         super().__init__(
-            num_input_units, num_output_units, arity=1, num_folds=num_folds, semiring=semiring
+            num_input_units,
+            num_output_units,
+            arity=1,
+            num_folds=num_folds,
+            semiring=semiring,
         )
         self._in_shape = (weight1.shape[1], weight2.shape[1])
         self.weight1 = weight1

diff --git a/src/models.py b/src/models.py
@@ -184,7 +184,7 @@ def __init__(
             backend="torch", semiring="complex-lse-sum", fold=True, optimize=True
         )
         # Use a different optimization rule for the dense-kronecker pattern
-        self._pipeline._compiler._optimization_registry['layer_shatter'].add_rule(
+        self._pipeline._compiler._optimization_registry["layer_shatter"].add_rule(
             apply_dense_product, signature=DenseKroneckerPattern
         )
         self._circuit, self._int_sq_circuit = self._build_circuits(
@@ -332,7 +332,7 @@ def categorical_layer_factory(
             num_channels,
             num_categories=input_layer_kwargs["num_categories"],
             logits_factory=lambda shape: Parameter.from_leaf(
-                TensorParameter(*shape, initializer=NormalInitializer(0.0, 3e-1))
+                TensorParameter(*shape, initializer=NormalInitializer(0.0, 1e-1))
             ),
         )
 
@@ -347,7 +347,7 @@ def gaussian_layer_factory(
                 TensorParameter(*shape, initializer=NormalInitializer(0.0, 1.0))
             ),
             stddev_factory=lambda shape: Parameter.from_sequence(
-                TensorParameter(*shape, initializer=NormalInitializer(0.0, 3e-1)),
+                TensorParameter(*shape, initializer=NormalInitializer(0.0, 1e-1)),
                 ExpParameter(shape),
                 ClampParameter(shape, vmin=1e-5),
             ),
@@ -367,7 +367,7 @@ def dense_layer_factory(
             num_output_units,
             weight_factory=lambda shape: Parameter.from_unary(
                 ExpParameter(shape),
-                TensorParameter(*shape, initializer=NormalInitializer(0.0, 3e-1)),
+                TensorParameter(*shape, initializer=NormalInitializer(0.0, 1e-1)),
             ),
         )
 
@@ -413,7 +413,7 @@ def categorical_layer_factory(
             num_channels,
             num_categories=input_layer_kwargs["num_categories"],
             logits_factory=lambda shape: Parameter.from_leaf(
-                TensorParameter(*shape, initializer=NormalInitializer(0.0, 3e-1))
+                TensorParameter(*shape, initializer=NormalInitializer(0.0, 1e-1))
             ),
         )
 
@@ -428,7 +428,7 @@ def gaussian_layer_factory(
                 TensorParameter(*shape, initializer=NormalInitializer(0.0, 1.0))
             ),
             stddev_factory=lambda shape: Parameter.from_sequence(
-                TensorParameter(*shape, initializer=NormalInitializer(0.0, 3e-1)),
+                TensorParameter(*shape, initializer=NormalInitializer(0.0, 1e-1)),
                 ExpParameter(shape),
                 ClampParameter(shape, vmin=1e-5),
             ),