Merge pull request #9 from Ipsedo/develop

change init weight std
Ipsedo · Jul 15, 2024 · bda657e · bda657e
2 parents 6d962ab + 569b327
commit bda657e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/kan/networks/conv.py b/kan/networks/conv.py
@@ -42,8 +42,8 @@ def __init__(
             )
         )
 
-        xavier_normal_(self._w_b)
-        normal_(self._c, 0, 1e-1)
+        xavier_normal_(self._w_b, gain=1e-3)
+        normal_(self._c, 0, 1e-3)
 
         self._in_channels = in_channels
         self._kernel_size = kernel_size

diff --git a/kan/networks/linear.py b/kan/networks/linear.py
@@ -27,14 +27,18 @@ def __init__(
             th.ones(self.__act_fun.get_size(), out_features, in_features)
         )
 
-        xavier_normal_(self._w_b)
-        normal_(self._c, 0, 1e-1)
+        xavier_normal_(self._w_b, gain=1e-3)
+        normal_(self._c, 0, 1e-3)
 
     def forward(self, x: th.Tensor) -> th.Tensor:
         assert len(x.size()) == 2
 
         return th.sum(
-            self._w_s * th.einsum("bai,aoi->boi", self.__act_fun(x), self._c)
+            th.einsum(
+                "boi,oi->boi",
+                th.einsum("bai,aoi->boi", self.__act_fun(x), self._c),
+                self._w_s,
+            )
             + th.einsum("bi,oi->boi", self.__res_act_fun(x), self._w_b),
             dim=2,
         )