feat: enhance training loop with tqdm progress bar and configurable parameters

mnist_classifier/train.py

@@ -9,6 +9,7 @@ from datetime import datetime
 import os
 import random
 import numpy as np
+from tqdm import tqdm
 
 def set_seed(seed):
     torch.manual_seed(seed)
@@ -20,6 +21,11 @@ def set_seed(seed):
     torch.backends.cudnn.benchmark = False
 
 def train():
+    # Training loop
+    learning_rate = 0.001
+    batch_size = 128
+    epochs = 10
+
     # Set seed for reproducibility
     set_seed(42)
     
@@ -32,18 +38,14 @@ def train():
     writer = SummaryWriter(log_dir)
 
     # Setup data
-    data_module = MNISTDataModule(batch_size=64, val_batch_size=1000)
+    data_module = MNISTDataModule(batch_size=batch_size, val_batch_size=1000)
     train_loader, test_loader = data_module.get_dataloaders()
 
     # Initialize model, optimizer, and loss function
     model = MNISTModel().to(device)
-    optimizer = optim.Adam(model.parameters())
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
     criterion = nn.CrossEntropyLoss()
 
-    # Training loop
-    learning_rate = 0.001
-    batch_size = 64
-    epochs = 10
 
     num_epochs = epochs
     for epoch in range(num_epochs):
@@ -52,28 +54,44 @@ def train():
         correct = 0
         total = 0
 
-        for batch_idx, batch in enumerate(train_loader):
-            images, labels = batch[0].to(device), batch[1].to(device)
-
-            optimizer.zero_grad()
-            outputs = model(images)
-            loss = criterion(outputs, labels)
-            loss.backward()
-            optimizer.step()
-
-            running_loss += loss.item()
-            _, predicted = outputs.max(1)
-            total += labels.size(0)
-            correct += predicted.eq(labels).sum().item()
-
-            if batch_idx % 100 == 99:
-                writer.add_scalar('training loss',
-                                running_loss / 100,
-                                epoch * len(train_loader) + batch_idx)
-                writer.add_scalar('training accuracy',
-                                100. * correct / total,
-                                epoch * len(train_loader) + batch_idx)
-                running_loss = 0.0
+        with tqdm(total=len(train_loader), desc=f"Epoch {epoch+1}/{num_epochs}", unit="batch") as pbar:
+            for batch_idx, batch in enumerate(train_loader):
+                images, labels = batch[0].to(device), batch[1].to(device)
+
+                if batch_idx == 0:
+                    print(f"images shape: {images.shape}")
+                    print(f"labels shape: {labels.shape}")
+                    
+                    # print number of images in batch
+                    print(f"Number of images in batch: {len(images)}")
+
+                optimizer.zero_grad()
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+
+                running_loss += loss.item()
+                _, predicted = outputs.max(1)
+                total += labels.size(0)
+                correct += predicted.eq(labels).sum().item()
+
+                # Update tqdm progress bar
+                pbar.set_postfix({
+                    'loss': running_loss / (batch_idx + 1),
+                    'accuracy': 100. * correct / total,
+                    'step': batch_idx + 1
+                })
+                pbar.update(1)
+
+                if batch_idx % 100 == 99:
+                    writer.add_scalar('training loss',
+                                      running_loss / 100,
+                                      epoch * len(train_loader) + batch_idx)
+                    writer.add_scalar('training accuracy',
+                                      100. * correct / total,
+                                      epoch * len(train_loader) + batch_idx)
+                    running_loss = 0.0
 
         # Validation phase
         model.eval()

models/mnist_model_lr0.001_bs128_ep10.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f9d6050aca93a46463f77e1a9dd4566da96e07905b9b872b519fa964f6984fc
+size 4803156