Merge pull request #557 from JeremCab/master

Update examples/mnist/batch_eth_mnist.py

examples/mnist/batch_eth_mnist.py

@@ -32,7 +32,7 @@
 parser.add_argument("--n_test", type=int, default=10000)
 parser.add_argument("--n_train", type=int, default=60000)
 parser.add_argument("--n_workers", type=int, default=-1)
-parser.add_argument("--update_steps", type=int, default=256)
+parser.add_argument("--n_updates", type=int, default=10)
 parser.add_argument("--exc", type=float, default=22.5)
 parser.add_argument("--inh", type=float, default=120)
 parser.add_argument("--theta_plus", type=float, default=0.05)
@@ -44,7 +44,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()
 
@@ -55,7 +55,7 @@
 n_test = args.n_test
 n_train = args.n_train
 n_workers = args.n_workers
-update_steps = args.update_steps
+n_updates = args.n_updates
 exc = args.exc
 inh = args.inh
 theta_plus = args.theta_plus
@@ -67,6 +67,7 @@
 plot = args.plot
 gpu = args.gpu
 
+update_steps = int(n_train / batch_size / n_updates)
 update_interval = update_steps * batch_size
 
 device = "cpu"
@@ -162,14 +163,14 @@
 spike_record = torch.zeros((update_interval, int(time / dt), n_neurons), device=device)
 
 # Train the network.
-print("\nBegin training.\n")
+print("\nBegin training...")
 start = t()
 
 for epoch in range(n_epochs):
     labels = []
 
     if epoch % progress_interval == 0:
-        print("\n Progress: %d / %d (%.4f seconds)" % (epoch, n_epochs, t() - start))
+        print("\nProgress: %d / %d (%.4f seconds)" % (epoch, n_epochs, t() - start))
         start = t()
 
     # Create a dataloader to iterate and batch data
@@ -183,13 +184,10 @@
 
     pbar_training = tqdm(total=n_train)
     for step, batch in enumerate(train_dataloader):
-        if step > n_train:
+        if step * batch_size > n_train:
             break
-        # Get next input sample.
-        inputs = {"X": batch["encoded_image"]}
-        if gpu:
-            inputs = {k: v.cuda() for k, v in inputs.items()}
 
+        # Assign labels to excitatory neurons.
         if step % update_steps == 0 and step > 0:
             # Convert the array of labels into a tensor
             label_tensor = torch.tensor(labels, device=device)
@@ -245,6 +243,12 @@
 
             labels = []
 
+        # Get next input sample.
+        inputs = {"X": batch["encoded_image"]}
+        if gpu:
+            inputs = {k: v.cuda() for k, v in inputs.items()}
+
+        # Remember labels.
         labels.extend(batch["label"].tolist())
 
         # Run the network on the input.
@@ -293,9 +297,10 @@
 
         network.reset_state_variables()  # Reset state variables.
         pbar_training.update(batch_size)
+    pbar_training.close()
 
 print("Progress: %d / %d (%.4f seconds)" % (epoch + 1, n_epochs, t() - start))
-print("Training complete.\n")
+print("\nTraining complete.\n")
 
 # Load MNIST data.
 test_dataset = MNIST(
@@ -322,13 +327,15 @@
 accuracy = {"all": 0, "proportion": 0}
 
 # Train the network.
-print("\nBegin testing\n")
+print("\nBegin testing...\n")
 network.train(mode=False)
 start = t()
 
 pbar = tqdm(total=n_test)
-for step, batch in enumerate(test_dataset):
-    if step > n_test:
+pbar.set_description_str("Test progress: ")
+
+for step, batch in enumerate(test_dataloader):
+    if step * batch_size > n_test:
         break
     # Get next input sample.
     inputs = {"X": batch["encoded_image"]}
@@ -362,11 +369,11 @@
     )
 
     network.reset_state_variables()  # Reset state variables.
-    pbar.set_description_str("Test progress: ")
-    pbar.update()
+    pbar.update(batch_size)
+pbar.close()
 
 print("\nAll activity accuracy: %.2f" % (accuracy["all"] / n_test))
 print("Proportion weighting accuracy: %.2f \n" % (accuracy["proportion"] / n_test))
 
 print("Progress: %d / %d (%.4f seconds)" % (epoch + 1, n_epochs, t() - start))
-print("Testing complete.\n")
+print("\nTesting complete.\n")