training bar

Author: Vishal-sys-code

requirements.txt

@@ -5,4 +5,5 @@ norse
 pyyaml
 transformers
 bindsnet
-pytest
+pytest
+tqdm

results/all_runs/dt_CartPole-v1/training.log

@@ -79,3 +79,15 @@
 2025-11-05 13:28:17,372 [INFO] Epoch 1, Batch 12
 2025-11-05 13:28:33,096 [INFO] Epoch 1, Batch 13
 2025-11-05 13:28:49,192 [INFO] Epoch 1, Batch 14
+2025-11-05 13:42:57,570 [INFO] Checking for dataset...
+2025-11-05 13:42:57,585 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-05 13:42:57,585 [INFO] Starting training...
+2025-11-05 13:42:57,704 [INFO] Dataset size: 1000 clips
+2025-11-05 13:42:57,709 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-05 13:43:09,026 [INFO] Starting training loop...
+2025-11-05 14:33:37,438 [INFO] Checking for dataset...
+2025-11-05 14:33:37,450 [INFO] Dataset found at D:\Github\neuromorphic_decision_transformer\data\CartPole-v1\dataset.npz.
+2025-11-05 14:33:37,450 [INFO] Starting training...
+2025-11-05 14:33:37,577 [INFO] Dataset size: 1000 clips
+2025-11-05 14:33:37,582 [INFO] DataLoader created with num_workers=0 and pin_memory=False.
+2025-11-05 14:33:48,431 [INFO] Starting training loop...

snn-dt/scripts/train.py

@@ -16,6 +16,7 @@
 import torch
 import yaml
 from torch.utils.data import Dataset
+from tqdm import tqdm
 import warnings
 warnings.filterwarnings('ignore')
 
@@ -151,11 +152,14 @@ def train(cfg, logger):
     for epoch in range(cfg.training.epochs):
         start_time = time.time()
         epoch_losses = []
-        for i, batch in enumerate(train_loader):
-            logger.info(f"Epoch {epoch+1}, Batch {i+1}")
-            if i >= cfg.training.batches_per_epoch:
-                break
+        
+        batch_iter = tqdm(
+            enumerate(train_loader), 
+            total=len(train_loader),
+            desc=f"Epoch {epoch+1}/{cfg.training.epochs}"
+        )
 
+        for i, batch in batch_iter:
             model.train()
 
             for k, v in batch.items():
@@ -177,6 +181,7 @@ def train(cfg, logger):
             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
             optimizer.step()
             epoch_losses.append(loss.item())
+            batch_iter.set_postfix(loss=f"{np.mean(epoch_losses):.4f}")
 
         # Evaluation, Checkpointing, and Logging
         if (epoch + 1) % cfg.training.eval_every == 0: