init

Author: n-shevko

bindsnet/models/models.py

@@ -169,7 +169,7 @@ def __init__(
         )
 
         # Connections
-        w = 0.3 * torch.rand(self.n_inpt, self.n_neurons)
+        w = 0.3 * torch.rand(self.n_inpt, self.n_neurons, dtype=torch.float16)
         input_exc_conn = Connection(
             source=input_layer,
             target=exc_layer,
@@ -181,13 +181,13 @@ def __init__(
             wmax=wmax,
             norm=norm,
         )
-        w = self.exc * torch.diag(torch.ones(self.n_neurons))
+        w = self.exc * torch.diag(torch.ones(self.n_neurons, dtype=torch.float16))
         exc_inh_conn = Connection(
             source=exc_layer, target=inh_layer, w=w, wmin=0, wmax=self.exc
         )
         w = -self.inh * (
-            torch.ones(self.n_neurons, self.n_neurons)
-            - torch.diag(torch.ones(self.n_neurons))
+            torch.ones(self.n_neurons, self.n_neurons, dtype=torch.float16)
+            - torch.diag(torch.ones(self.n_neurons, dtype=torch.float16))
         )
         inh_exc_conn = Connection(
             source=inh_layer, target=exc_layer, w=w, wmin=-self.inh, wmax=0

bindsnet/network/topology.py

@@ -322,7 +322,7 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
         """
         # Compute multiplication of spike activations by weights and add bias.
         if self.b is None:
-            post = s.view(s.size(0), -1).float() @ self.w
+            post = s.view(s.size(0), -1).to(dtype=torch.float16) @ self.w
         else:
             post = s.view(s.size(0), -1).float() @ self.w + self.b
         return post.view(s.size(0), *self.target.shape)

examples/mnist/eth_mnist.py

@@ -42,7 +42,7 @@
 parser.add_argument("--test", dest="train", action="store_false")
 parser.add_argument("--plot", dest="plot", action="store_true")
 parser.add_argument("--gpu", dest="gpu", action="store_true")
-parser.set_defaults(plot=True, gpu=True)
+parser.set_defaults(plot=False, gpu=True)
 
 args = parser.parse_args()