feat: baseline model

.gitignore

@@ -4,4 +4,7 @@ __pycache__/
 wandb/
 checkpoints/
 *.egg-info/
-dist/
+dist/
+mnist/
+data/
+runs/

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "python.analysis.typeCheckingMode": "standard"
+}

mnist_classifier/configs/config.yaml

@@ -1,15 +0,0 @@
-training:
-  batch_size: 64
-  max_epochs: 10
-  learning_rate: 0.001
-  early_stopping_patience: 5
-
-model:
-  conv1_channels: 32
-  conv2_channels: 64
-  fc1_size: 128
-  dropout_rate: 0.25
-
-wandb:
-  project: "mnist-classifier"
-  entity: "bardenha"

mnist_classifier/data/datamodule.py

@@ -1,27 +0,0 @@
-from typing import Dict, Any
-
-import pytorch_lightning as pl
-from datasets import load_dataset
-from torch.utils.data import DataLoader
-
-class MNISTDataModule(pl.LightningDataModule):
-    def __init__(self, config: Dict[str, Any]):
-        super().__init__()
-        self.config = config
-
-    def setup(self, stage=None):
-        self.dataset = load_dataset('mnist')
-        self.dataset = self.dataset.with_transform(self.config.transform_dataset)
-    
-    def train_dataloader(self):
-        return DataLoader(
-            self.dataset['train'], 
-            batch_size=self.config.batch_size, 
-            shuffle=True
-        )
-    
-    def val_dataloader(self):
-        return DataLoader(
-            self.dataset['test'],  # Using test set as validation
-            batch_size=self.config.batch_size
-        )

mnist_classifier/dataset.py

@@ -0,0 +1,21 @@
+import torch
+from torchvision import datasets, transforms
+
+class MNISTDataModule:
+    def __init__(self, batch_size=64, val_batch_size=1000):
+        self.batch_size = batch_size
+        self.val_batch_size = val_batch_size
+        
+    def get_dataloaders(self):
+        """Create training and test dataloaders."""
+        transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))
+        ])
+        train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
+        test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)
+
+        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
+        test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=self.val_batch_size, shuffle=False)
+
+        return train_loader, test_loader

mnist_classifier/model.py

@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MNISTModel(nn.Module):
+    def __init__(self):
+        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.fc1 = nn.Linear(9216, 128)
+        self.fc2 = nn.Linear(128, 10)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = F.relu(x)
+        x = F.max_pool2d(x, 2)
+        x = self.dropout1(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout2(x)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)

mnist_classifier/models/mnist_model.py

@@ -1,124 +0,0 @@
-from typing import Dict, Any
-
-import pytorch_lightning as pl
-import torch
-import torch.nn as nn
-import torchmetrics
-import wandb
-
-# Simple CNN architecture for MNIST
-class MNISTNet(nn.Module):
-    def __init__(self, config: Dict[str, Any]):
-        super().__init__()
-        self.conv1 = nn.Conv2d(1, config['model']['conv1_channels'], kernel_size=3)
-        self.conv2 = nn.Conv2d(config['model']['conv1_channels'], 
-                              config['model']['conv2_channels'], kernel_size=3)
-        self.pool = nn.MaxPool2d(2)
-        self.dropout = nn.Dropout(config['model']['dropout_rate'])
-        self.fc1 = nn.Linear(config['model']['conv2_channels'] * 5 * 5, 
-                            config['model']['fc1_size'])
-        self.fc2 = nn.Linear(config['model']['fc1_size'], 10)
-
-    def forward(self, x):
-        x = torch.relu(self.conv1(x))
-        x = self.pool(torch.relu(self.conv2(x)))
-        x = self.dropout(x)
-        x = x.view(-1, 64 * 5 * 5)
-        x = torch.relu(self.fc1(x))
-        x = self.fc2(x)
-        return x
-    
-    
-
-class MNISTModule(pl.LightningModule):
-    def __init__(self, config: Dict[str, Any]):
-        super().__init__()
-        self.config = config
-        self.model = MNISTNet(config)
-
-        # Initialize metrics
-        self.train_accuracy = torchmetrics.Accuracy(task='multiclass', num_classes=10)
-        self.val_accuracy = torchmetrics.Accuracy(task='multiclass', num_classes=10)
-        self.train_f1 = torchmetrics.F1Score(task='multiclass', num_classes=10)
-        self.val_f1 = torchmetrics.F1Score(task='multiclass', num_classes=10)
-        self.confusion_matrix = torchmetrics.ConfusionMatrix(task='multiclass', num_classes=10)
-    
-    def forward(self, x):
-        return self.model(x)
-    
-    def training_step(self, batch, batch_idx):
-        x, y = batch['pixel_values'], batch['label']
-        logits = self(x)
-        loss = nn.CrossEntropyLoss()(logits, y)
-        
-        # Calculate and log metrics
-        preds = torch.argmax(logits, dim=1)
-        self.train_accuracy(preds, y)
-        self.train_f1(preds, y)
-        
-        # Log metrics
-        self.log('train_loss', loss, prog_bar=True)
-        self.log('train_accuracy', self.train_accuracy, prog_bar=True)
-        self.log('train_f1', self.train_f1, prog_bar=True)
-        
-        return loss
-    
-    def validation_step(self, batch, batch_idx):
-        x, y = batch['pixel_values'], batch['label']
-        logits = self(x)
-        loss = nn.CrossEntropyLoss()(logits, y)
-        
-        # Calculate metrics
-        preds = torch.argmax(logits, dim=1)
-        self.val_accuracy(preds, y)
-        self.val_f1(preds, y)
-        self.confusion_matrix(preds, y)
-        
-        # Log metrics
-        self.log('val_loss', loss, prog_bar=True)
-        self.log('val_accuracy', self.val_accuracy, prog_bar=True)
-        self.log('val_f1', self.val_f1, prog_bar=True)
-        
-        # Log sample predictions periodically
-        if batch_idx == 0:  # First batch of each epoch
-            self._log_sample_predictions(x, y, preds)
-    
-    def _log_sample_predictions(self, images, labels, predictions):
-        # Log a grid of sample predictions
-        if self.logger:
-            n_samples = min(16, len(images))
-            self.logger.experiment.log({
-                "sample_predictions": [
-                    wandb.Image(
-                        images[i],
-                        caption=f"True: {labels[i].item()} Pred: {predictions[i].item()}"
-                    )
-                    for i in range(n_samples)
-                ]
-            })
-    
-    def on_validation_epoch_end(self):
-        # Log confusion matrix at the end of each validation epoch
-        conf_mat = self.confusion_matrix.compute()
-        self.logger.experiment.log({
-            "confusion_matrix": wandb.plot.confusion_matrix(
-                probs=None,
-                y_true=conf_mat.flatten(),
-                preds=None,
-                class_names=range(10)
-            )
-        })
-        self.confusion_matrix.reset()
-    
-    def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
-        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optimizer, mode='min', factor=0.1, patience=3, verbose=True
-        )
-        return {
-            "optimizer": optimizer,
-            "lr_scheduler": {
-                "scheduler": scheduler,
-                "monitor": "val_loss"
-            }
-        }

mnist_classifier/train.py

@@ -1,49 +1,79 @@
-import pytorch_lightning as pl
-from pytorch_lightning.loggers import WandbLogger
-from pathlib import Path
-
-from mnist_classifier.models.mnist_model import MNISTModule
-from mnist_classifier.data.datamodule import MNISTDataModule
-from mnist_classifier.utils.metrics import load_config
-
-def main():
-    config = load_config(Path("mnist_classifier/configs/config.yaml"))
-    
-    # Initialize wandb logger
-    wandb_logger = WandbLogger(
-        project=config['wandb']['project'],
-        entity=config['wandb']['entity']
-    )
-    
-    # Initialize trainer
-    trainer = pl.Trainer(
-        max_epochs=config['training']['max_epochs'],
-        accelerator='gpu',
-        devices=[0],
-        logger=wandb_logger,
-        callbacks=[
-            pl.callbacks.ModelCheckpoint(
-                dirpath='checkpoints',
-                filename='mnist-{epoch:02d}-{val_loss:.2f}',
-                save_top_k=3,
-                monitor='val_loss',
-                mode='min'
-            ),
-            pl.callbacks.EarlyStopping(
-                monitor='val_loss',
-                patience=config['training']['early_stopping_patience'],
-                mode='min'
-            ),
-            pl.callbacks.LearningRateMonitor(logging_interval='epoch')
-        ]
-    )
-    
-    # Initialize data module and model
-    data_module = MNISTDataModule(config)
-    model = MNISTModule(config)
-    
-    # Train
-    trainer.fit(model, data_module)
+import torch
+import torch.nn as nn
+import torch.optim as optim
+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
+
+def train():
+    # Set device
+    device = torch.device('cuda')
+    print(f"Using device: {device}")
+
+    # Initialize tensorboard
+    writer = SummaryWriter('runs/mnist_experiment')
+
+    # Setup data
+    data_module = MNISTDataModule(batch_size=64, 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())
+    criterion = nn.CrossEntropyLoss()
+
+    # Training loop
+    num_epochs = 10
+    for epoch in range(num_epochs):
+        model.train()
+        running_loss = 0.0
+        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
+
+        # Validation phase
+        model.eval()
+        test_loss = 0
+        correct = 0
+        total = 0
+        with torch.no_grad():
+            for batch in test_loader:
+                images = batch[0].to(device)
+                labels = batch[1].to(device)
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                
+                test_loss += loss.item()
+                _, predicted = outputs.max(1)
+                total += labels.size(0)
+                correct += predicted.eq(labels).sum().item()
+
+        accuracy = 100. * correct / total
+        writer.add_scalar('test accuracy', accuracy, epoch)
+        print(f'Epoch {epoch+1}: Test Accuracy: {accuracy:.2f}%')
 
 if __name__ == "__main__":
-    main()
+    train()

mnist_classifier/utils/metrics.py

@@ -1,6 +0,0 @@
-import yaml
-from pathlib import Path
-
-def load_config(config_path: str):
-    with open(config_path, 'r') as f:
-        return yaml.safe_load(f)

pyproject.toml

@@ -8,14 +8,14 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.10"
-torch = "^2.4.0"
-torchvision = "^0.15.0"
-pytorch-lightning = "^2.0.0"
-wandb = "^0.15.0"
-torchmetrics = "^1.0.0"
-datasets = "^2.0.0"
-huggingface-hub = "^0.16.0"
 pyyaml = "^6.0"
+torch = "^2.5.1"
+torchvision = "^0.20.1"
+transformers = "^4.46.3"
+datasets = "^3.1.0"
+tensorboard = "^2.18.0"
+tqdm = "^4.67.0"
+types-tqdm = "^4.67.0.20241119"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.0.0"

tests/test_model.py

@@ -1,11 +0,0 @@
-import pytest
-import torch
-from mnist_classifier.models.mnist_model import MNISTNet
-from mnist_classifier.utils.metrics import load_config
-
-def test_mnist_net_forward():
-    config = load_config('mnist_classifier/configs/config.yaml')
-    model = MNISTNet(config)
-    x = torch.randn(1, 1, 28, 28)
-    output = model(x)
-    assert output.shape == (1, 10)