feat: update model training and logging, add mypy cache to .gitignore

.gitignore

@@ -1,6 +1,7 @@
 __pycache__/
 *.pyc
 .pytest_cache/
+.mypy_cache/
 wandb/
 checkpoints/
 *.egg-info/

mnist_classifier/model.py

@@ -7,8 +7,8 @@ class MNISTModel(nn.Module):
         super().__init__()
         self.conv1 = nn.Conv2d(1, 32, 3, 1)
         self.conv2 = nn.Conv2d(32, 64, 3, 1)
-        self.dropout1 = nn.Dropout2d(0.25)
-        self.dropout2 = nn.Dropout2d(0.5)
+        self.dropout1 = nn.Dropout(0.25)
+        self.dropout2 = nn.Dropout(0.5)
         self.fc1 = nn.Linear(9216, 128)
         self.fc2 = nn.Linear(128, 10)
 

mnist_classifier/train.py

@@ -5,6 +5,8 @@ from torch.utils.data import DataLoader
 from torch.utils.tensorboard.writer import SummaryWriter
 from mnist_classifier.dataset import MNISTDataModule
 from mnist_classifier.model import MNISTModel
+from datetime import datetime
+import os
 
 def train():
     # Set device
@@ -12,7 +14,8 @@ def train():
     print(f"Using device: {device}")
 
     # Initialize tensorboard
-    writer = SummaryWriter('runs/mnist_experiment')
+    log_dir = 'runs/mnist_experiment_' + datetime.now().strftime('%Y%m%d-%H%M%S')
+    writer = SummaryWriter(log_dir)
 
     # Setup data
     data_module = MNISTDataModule(batch_size=64, val_batch_size=1000)
@@ -24,7 +27,11 @@ def train():
     criterion = nn.CrossEntropyLoss()
 
     # Training loop
-    num_epochs = 10
+    learning_rate = 0.001
+    batch_size = 64
+    epochs = 10
+
+    num_epochs = epochs
     for epoch in range(num_epochs):
         model.train()
         running_loss = 0.0
@@ -75,5 +82,14 @@ def train():
         writer.add_scalar('test accuracy', accuracy, epoch)
         print(f'Epoch {epoch+1}: Test Accuracy: {accuracy:.2f}%')
 
+    writer.close()
+
+    # Ensure the directory exists
+    os.makedirs("./models", exist_ok=True)
+
+    # Format the filename with the config parameters
+    filename = f"./models/mnist_model_lr{learning_rate}_bs{batch_size}_ep{epochs}.pth"
+    torch.save(model.state_dict(), filename)
+
 if __name__ == "__main__":
     train()

models/mnist_model_lr0.001_bs64_ep10.pth

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

torchvision.pyi

@@ -0,0 +1,9 @@
+from typing import Any
+
+class datasets:
+    MNIST: Any
+
+class transforms:
+    Compose: Any
+    ToTensor: Any
+    Normalize: Any