Updates

models/mlp.py

@@ -1,3 +1,5 @@
+import time
+
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -8,16 +10,16 @@ class MLP(nn.Module):
     def __init__(self, hidden_sizes=(128, 64)):
         super().__init__()
         layers = []
-        in_size = 784
-        for h in hidden_sizes:
-            layers.append(nn.Linear(in_size, h))
+        input_size = 784
+        for hidden_size in hidden_sizes:
+            layers.append(nn.Linear(input_size, hidden_size))
             layers.append(nn.ReLU())
-            in_size = h
-        layers.append(nn.Linear(in_size, 10))
+            input_size = hidden_size
+        layers.append(nn.Linear(input_size, 10))
         self.net = nn.Sequential(*layers)
 
-    def forward(self, x):
-        return self.net(x)
+    def forward(self, pixels):
+        return self.net(pixels)
 
 
 class MLPClassifier:
@@ -26,53 +28,56 @@ class MLPClassifier:
         self.epochs = epochs
 
     def fit(self, X, y):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self._device = device
+        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self._model = MLP(hidden_sizes=self.hidden_sizes).to(self._device)
 
-        X_tr = torch.tensor(X, dtype=torch.float32)
-        y_tr = torch.tensor(y, dtype=torch.long)
+        images = torch.tensor(X, dtype=torch.float32)
+        labels = torch.tensor(y, dtype=torch.long)
+        train_images, train_labels, val_images, val_labels = self._split(images, labels)
 
-        # Hold out 10% of the training data to track progress each epoch
-        n_val = len(X_tr) // 10
-        X_val, X_tr = X_tr[:n_val], X_tr[n_val:]
-        y_val, y_tr = y_tr[:n_val], y_tr[n_val:]
-
-        loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64, shuffle=True)
-
-        model = MLP(hidden_sizes=self.hidden_sizes).to(device)
-        optimizer = optim.Adam(model.parameters(), lr=1e-3)
+        batches = DataLoader(TensorDataset(train_images, train_labels), batch_size=64, shuffle=True)
+        optimizer = optim.Adam(self._model.parameters(), lr=1e-3)
         loss_fn = nn.CrossEntropyLoss()
 
         print(f"\nTraining MLP (hidden_sizes={self.hidden_sizes}, epochs={self.epochs})")
         for epoch in range(1, self.epochs + 1):
-            model.train()
-            total_loss = 0
-            for xb, yb in loader:
-                xb, yb = xb.to(device), yb.to(device)
-                optimizer.zero_grad()
-                loss = loss_fn(model(xb), yb)
-                loss.backward()
-                optimizer.step()
-                total_loss += loss.item()
-
-            model.eval()
-            with torch.no_grad():
-                val_pred = model(X_val.to(device)).argmax(dim=1).cpu()
-                val_accuracy = (val_pred == y_val).float().mean().item()
-
-            print(f"  epoch {epoch:2d}/{self.epochs}  loss={total_loss / len(loader):.3f}  val_accuracy={val_accuracy:.3f}")
+            t0 = time.time()
+            avg_loss = self._train_one_epoch(batches, optimizer, loss_fn)
+            val_accuracy = self._accuracy(val_images, val_labels)
+            elapsed = time.time() - t0
+            print(f"  epoch {epoch:2d}/{self.epochs}  loss={avg_loss:.3f}  val_accuracy={val_accuracy:.3f}  {elapsed:.1f}s")
         print()
-
-        self._model = model
         return self
 
-    def predict_proba(self, X):
-        X_te = torch.tensor(X, dtype=torch.float32)
+    def _split(self, images, labels):
+        n_val = len(images) // 10
+        return images[n_val:], labels[n_val:], images[:n_val], labels[:n_val]
+
+    def _train_one_epoch(self, batches, optimizer, loss_fn):
+        self._model.train()
+        total_loss = 0
+        for image_batch, label_batch in batches:
+            image_batch = image_batch.to(self._device)
+            label_batch = label_batch.to(self._device)
+            optimizer.zero_grad()
+            loss = loss_fn(self._model(image_batch), label_batch)
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+        return total_loss / len(batches)
+
+    def _accuracy(self, images, labels):
         self._model.eval()
         with torch.no_grad():
-            logits = self._model(X_te.to(self._device))
-            probabilities = torch.softmax(logits, dim=1).cpu().numpy()
-        return probabilities
+            predictions = self._model(images.to(self._device)).argmax(dim=1).cpu()
+            return (predictions == labels).float().mean().item()
+
+    def predict_proba(self, X):
+        images = torch.tensor(X, dtype=torch.float32)
+        self._model.eval()
+        with torch.no_grad():
+            logits = self._model(images.to(self._device))
+            return torch.softmax(logits, dim=1).cpu().numpy()
 
     def predict(self, X):
         return self.predict_proba(X).argmax(axis=1)