support mypy

.github/workflows/poetry.yml

@@ -3,9 +3,11 @@ name: Poetry
 on:
   workflow_dispatch:
   push:
-    branches: ["main"]
+    branches:
+      - main
   pull_request:
-    branches: ["main"]
+    branches:
+      - main
 
 jobs:
   poetry:
@@ -27,6 +29,8 @@ jobs:
         run: poetry install
       - name: Lint
         run: poetry run ruff check .
+      - name: Type check
+        run: poetry run mypy --install-types --non-interactive . 
       - name: Test
         run: poetry run pytest -v -s --cov=. --cov-report=xml tests
       - name: Upload coverage reports to Codecov

.github/workflows/python.yml

@@ -3,9 +3,11 @@ name: Python
 on:
   workflow_dispatch:
   push:
-    branches: ["main"]
+    branches:
+      - main
   pull_request:
-    branches: ["main"]
+    branches:
+      - main
 
 jobs:
   python:

pyproject.toml

@@ -55,3 +55,6 @@ known-third-party = ["wandb"]
 
 [tool.pytest.ini_options]
 filterwarnings = ["ignore::DeprecationWarning"]
+
+[tool.mypy]
+ignore_missing_imports = true

template/jobs/job.py

@@ -1,6 +1,6 @@
-from omegaconf import OmegaConf
+from omegaconf import DictConfig
 
 
 class Job:
-    def run(self, config: OmegaConf, resume=None) -> None:
+    def run(self, config: DictConfig, resume: str | None = None) -> None:
         raise NotImplementedError

template/jobs/mnist.py

@@ -1,15 +1,15 @@
 import torch
 from mlconfig import instantiate
 from mlconfig import register
-from omegaconf import OmegaConf
+from omegaconf import DictConfig
 
 from ..utils import manual_seed
 from .job import Job
 
 
 @register
 class MNISTTrainingJob(Job):
-    def run(self, config: OmegaConf, resume=None) -> None:
+    def run(self, config: DictConfig, resume: str | None = None) -> None:
         manual_seed()
 
         device = torch.device(config.device if torch.cuda.is_available() else "cpu")

template/trainers/mnist.py

@@ -6,8 +6,8 @@
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
-from torchmetrics import Accuracy
 from torchmetrics import MeanMetric
+from torchmetrics.classification import MulticlassAccuracy
 from tqdm import tqdm
 from tqdm import trange
 
@@ -36,7 +36,7 @@ def __init__(
         self.num_epochs = num_epochs
         self.num_classes = num_classes
 
-        self.best_acc = 0
+        self.best_acc = 0.0
         self.state = {"epoch": 1}
 
     def fit(self) -> None:
@@ -61,11 +61,11 @@ def fit(self) -> None:
 
             self.state["epoch"] = epoch
 
-    def train(self) -> None:
+    def train(self) -> tuple[float, float]:
         self.model.train()
 
-        loss_metric = MeanMetric().to(self.device)
-        acc_metric = Accuracy(task="multiclass", num_classes=self.num_classes).to(self.device)
+        loss_metric = MeanMetric()
+        acc_metric = MulticlassAccuracy(num_classes=self.num_classes)
 
         for x, y in tqdm(self.train_loader):
             x = x.to(self.device)
@@ -78,17 +78,17 @@ def train(self) -> None:
             loss.backward()
             self.optimizer.step()
 
-            loss_metric.update(loss, weight=x.size(0))
-            acc_metric.update(output, y)
+            loss_metric.update(loss.item(), weight=x.size(0))
+            acc_metric.update(output.cpu(), y.cpu())
 
-        return loss_metric.compute().item(), acc_metric.compute().item()
+        return float(loss_metric.compute()), float(acc_metric.compute())
 
     @torch.no_grad()
-    def evaluate(self) -> None:
+    def evaluate(self) -> tuple[float, float]:
         self.model.eval()
 
-        loss_metric = MeanMetric().to(self.device)
-        acc_metric = Accuracy(task="multiclass", num_classes=self.num_classes).to(self.device)
+        loss_metric = MeanMetric()
+        acc_metric = MulticlassAccuracy(num_classes=self.num_classes)
 
         for x, y in tqdm(self.test_loader):
             x = x.to(self.device)
@@ -97,15 +97,15 @@ def evaluate(self) -> None:
             output = self.model(x)
             loss = f.cross_entropy(output, y)
 
-            loss_metric.update(loss, weight=x.size(0))
-            acc_metric.update(output, y)
+            loss_metric.update(loss.item(), weight=x.size(0))
+            acc_metric.update(output.cpu(), y.cpu())
 
-        test_acc = acc_metric.compute().item()
+        test_acc = float(acc_metric.compute())
         if test_acc > self.best_acc:
             self.best_acc = test_acc
             self.save_checkpoint("best.pth")
 
-        return loss_metric.compute().item(), test_acc
+        return float(loss_metric.compute()), test_acc
 
     def save_checkpoint(self, f) -> None:
         self.model.eval()

template/trainers/trainer.py

@@ -1,3 +1,3 @@
 class Trainer:
-    def train(self) -> None:
+    def train(self):
         raise NotImplementedError