cnn-classifier-toolkit 0.1.0__py3-none-any.whl

cnn_classifier/models/simple_cnn.py

@@ -0,0 +1,55 @@
+"""Minimal stacked CNN baseline."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from pydantic import BaseModel, ConfigDict, Field
+from torch import nn
+
+from cnn_classifier.models.base import (
+    ModelBuildContext,
+    register_model,
+)
+
+
+class SimpleCNNParams(BaseModel):
+    model_config = ConfigDict(extra="forbid", allow_inf_nan=False)
+
+    channels: tuple[int, ...] = Field(default=(32, 64, 128), min_length=1)
+    dropout: float = Field(default=0.5, ge=0.0, lt=1.0)
+    norm: Literal["batch", "none"] = "batch"
+
+
+class SimpleCNN(nn.Module):
+    def __init__(self, params: SimpleCNNParams, context: ModelBuildContext) -> None:
+        super().__init__()
+        layers: list[nn.Module] = []
+        in_channels = context.in_channels
+        for out_channels in params.channels:
+            layers.append(nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=params.norm == "none"))
+            if params.norm == "batch":
+                layers.append(nn.BatchNorm2d(out_channels))
+            layers.extend([nn.ReLU(inplace=True), nn.MaxPool2d(2)])
+            in_channels = out_channels
+        self.features = nn.Sequential(*layers)
+        self.head = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+            nn.Dropout(params.dropout),
+            nn.Linear(in_channels, context.num_classes),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        result: torch.Tensor = self.head(self.features(x))
+        return result
+
+
+@register_model("simple_cnn")
+class SimpleCNNBuilder:
+    Params = SimpleCNNParams
+
+    @staticmethod
+    def build(params: SimpleCNNParams, context: ModelBuildContext) -> nn.Module:
+        return SimpleCNN(params, context)

cnn_classifier/models/summary.py

@@ -0,0 +1,39 @@
+"""Model parameter and shape summaries."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+
+
+def model_summary(model: nn.Module, input_shape: tuple[int, int, int]) -> dict[str, Any]:
+    """Return parameter counts, module details, and output shape."""
+    was_training = model.training
+    parameter = next(model.parameters(), None)
+    device = parameter.device if parameter is not None else torch.device("cpu")
+    model.eval()
+    try:
+        with torch.no_grad():
+            output = model(torch.zeros(1, *input_shape, device=device))
+    finally:
+        model.train(was_training)
+    modules = [
+        {
+            "name": name or "model",
+            "type": type(module).__name__,
+            "params": sum(parameter.numel() for parameter in module.parameters()),
+        }
+        for name, module in model.named_modules()
+    ]
+    return {
+        "input_shape": list(input_shape),
+        "output_shape": list(output.shape),
+        "total_params": sum(parameter.numel() for parameter in model.parameters()),
+        "trainable_params": sum(
+            parameter.numel() for parameter in model.parameters() if parameter.requires_grad
+        ),
+        "modules": modules,
+        "macs": None,
+    }

cnn_classifier/models/torchvision.py

@@ -0,0 +1,46 @@
+"""Torchvision model builder with transfer-learning controls."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torchvision
+from pydantic import BaseModel, ConfigDict, Field
+from torch import nn
+
+from cnn_classifier.core.errors import ConfigError
+from cnn_classifier.models.base import ModelBuildContext, register_model
+from cnn_classifier.models.transfer import (
+    apply_stem_adaptation,
+    apply_transfer,
+    replace_classifier,
+)
+
+
+class TorchvisionParams(BaseModel):
+    model_config = ConfigDict(extra="forbid", allow_inf_nan=False)
+
+    model: str = "resnet18"
+    stem_adaptation: Literal["none", "keep_7x7_stride1", "adapt_3x3"] = "none"
+    dropout: float = Field(default=0.0, ge=0.0, lt=1.0)
+
+
+@register_model("torchvision")
+class TorchvisionBuilder:
+    Params = TorchvisionParams
+
+    @staticmethod
+    def build(params: TorchvisionParams, context: ModelBuildContext) -> nn.Module:
+        try:
+            model = torchvision.models.get_model(
+                params.model,
+                weights=context.transfer.weights,
+            )
+        except (RuntimeError, ValueError) as exc:
+            raise ConfigError(f"cannot build torchvision model {params.model!r}: {exc}") from exc
+        if not isinstance(model, nn.Module):
+            raise ConfigError(f"torchvision model {params.model!r} did not return an nn.Module")
+        replace_classifier(model, context.num_classes, dropout=params.dropout)
+        apply_stem_adaptation(model, params.stem_adaptation, context.in_channels)
+        apply_transfer(model, context.transfer, stem_adaptation=params.stem_adaptation)
+        return model

cnn_classifier/models/transfer.py

@@ -0,0 +1,193 @@
+"""Transfer-learning helpers for stem adaptation, freezing, and param groups."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+from typing import Any, Literal
+
+import torch
+from torch import nn
+
+from cnn_classifier.config.schema import TransferConfig
+from cnn_classifier.core.errors import ConfigError
+
+StemMode = Literal["none", "keep_7x7_stride1", "adapt_3x3"]
+
+
+def _head_name(model: nn.Module) -> str:
+    if hasattr(model, "fc") and isinstance(model.fc, (nn.Linear, nn.Sequential)):
+        return "fc"
+    if hasattr(model, "classifier") and isinstance(model.classifier, (nn.Linear, nn.Sequential)):
+        return "classifier"
+    raise ConfigError("cannot identify the classifier head for this model")
+
+
+def _last_linear(module: nn.Module) -> nn.Linear:
+    if isinstance(module, nn.Linear):
+        return module
+    if isinstance(module, nn.Sequential):
+        for child in reversed(module):
+            if isinstance(child, nn.Linear):
+                return child
+    raise ConfigError("classifier does not contain a Linear layer")
+
+
+def replace_classifier(model: nn.Module, num_classes: int, dropout: float = 0.0) -> nn.Module:
+    """Replace the final classifier with a task-specific head."""
+    if hasattr(model, "fc") and isinstance(model.fc, (nn.Linear, nn.Sequential)):
+        in_features = _last_linear(model.fc).in_features
+        head = nn.Linear(in_features, num_classes)
+        model.fc = nn.Sequential(nn.Dropout(dropout), head) if dropout > 0 else head
+        return model
+    if hasattr(model, "classifier") and isinstance(model.classifier, (nn.Linear, nn.Sequential)):
+        classifier = model.classifier
+        in_features = _last_linear(classifier).in_features
+        new_layer = nn.Linear(in_features, num_classes)
+        if isinstance(classifier, nn.Linear):
+            model.classifier = new_layer
+        else:
+            for index in range(len(classifier) - 1, -1, -1):
+                if isinstance(classifier[index], nn.Linear):
+                    classifier[index] = new_layer
+                    break
+        return model
+    raise ConfigError("unsupported torchvision classifier head")
+
+
+def apply_stem_adaptation(
+    model: nn.Module,
+    mode: StemMode,
+    in_channels: int = 3,
+) -> nn.Module:
+    """Adapt a torchvision stem for native 32x32 inputs."""
+    if mode == "none":
+        return model
+    if not hasattr(model, "conv1") or not isinstance(model.conv1, nn.Conv2d):
+        raise ConfigError(f"model {type(model).__name__} has no supported conv1 stem")
+    conv = model.conv1
+    if conv.in_channels != in_channels:
+        raise ConfigError(
+            f"stem expects {conv.in_channels} input channels but data has {in_channels}"
+        )
+    if mode == "keep_7x7_stride1":
+        conv.stride = (1, 1)
+        model.maxpool = nn.Identity()
+        return model
+    if mode == "adapt_3x3":
+        new_conv = nn.Conv2d(
+            in_channels,
+            conv.out_channels,
+            3,
+            1,
+            1,
+            bias=conv.bias is not None,
+            device=conv.weight.device,
+            dtype=conv.weight.dtype,
+        )
+        kernel = conv.weight.data
+        height, width = kernel.shape[-2:]
+        top = (height - 3) // 2
+        left = (width - 3) // 2
+        with torch.no_grad():
+            new_conv.weight.copy_(kernel[:, :, top : top + 3, left : left + 3])
+            if conv.bias is not None and new_conv.bias is not None:
+                new_conv.bias.copy_(conv.bias.data)
+        model.conv1 = new_conv
+        model.maxpool = nn.Identity()
+        return model
+    raise ConfigError(f"unknown stem adaptation {mode!r}")
+
+
+def _head_parameters(model: nn.Module) -> set[str]:
+    head = _head_name(model)
+    return {
+        name for name, _ in model.named_parameters() if name == head or name.startswith(f"{head}.")
+    }
+
+
+def apply_transfer(
+    model: nn.Module,
+    transfer: TransferConfig,
+    *,
+    stem_adaptation: str,
+) -> nn.Module:
+    """Apply freezing, BatchNorm policy, and validation for a transfer run."""
+    if transfer.freeze in {"backbone", "all_but_head"}:
+        if stem_adaptation == "adapt_3x3":
+            raise ConfigError(
+                "frozen adapt_3x3 is invalid: the uninitialized stem would block ImageNet features"
+            )
+        head_names = _head_parameters(model)
+        for name, parameter in model.named_parameters():
+            parameter.requires_grad = name in head_names
+        if transfer.train_bn:
+            for module in model.modules():
+                if isinstance(module, nn.modules.batchnorm._BatchNorm):
+                    for parameter in module.parameters():
+                        parameter.requires_grad = True
+    return model
+
+
+def build_param_groups(
+    model: nn.Module,
+    transfer: TransferConfig,
+    *,
+    base_lr: float,
+    base_weight_decay: float,
+) -> list[dict[str, Any]]:
+    """Build disjoint optimizer groups, validating every regex pattern."""
+    trainable = {
+        name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad
+    }
+    if not trainable:
+        raise ConfigError("model has no trainable parameters")
+    remaining = dict(trainable)
+    groups: list[dict[str, Any]] = []
+    for group_config in transfer.param_groups:
+        pattern = re.compile(group_config.pattern)
+        matched = [
+            (name, parameter) for name, parameter in remaining.items() if pattern.search(name)
+        ]
+        if not matched:
+            raise ConfigError(
+                f"param group pattern {group_config.pattern!r} did not match "
+                "any trainable parameter"
+            )
+        groups.append(
+            {
+                "params": [parameter for _, parameter in matched],
+                "lr": base_lr * group_config.lr_multiplier,
+                "weight_decay": base_weight_decay * group_config.weight_decay_multiplier,
+            }
+        )
+        for name, _ in matched:
+            del remaining[name]
+    if remaining:
+        groups.insert(
+            0,
+            {
+                "params": list(remaining.values()),
+                "lr": base_lr,
+                "weight_decay": base_weight_decay,
+            },
+        )
+    if not groups:
+        groups = [
+            {
+                "params": list(trainable.values()),
+                "lr": base_lr,
+                "weight_decay": base_weight_decay,
+            }
+        ]
+    return groups
+
+
+def unfreeze_modules(model: nn.Module, names: Sequence[str]) -> None:
+    """Unfreeze every parameter under the named submodules."""
+    modules = dict(model.named_modules())
+    for name in names:
+        if name not in modules:
+            raise ConfigError(f"cannot unfreeze unknown module {name!r}")
+        for parameter in modules[name].parameters():
+            parameter.requires_grad = True

cnn_classifier/runs/__init__.py

@@ -0,0 +1,6 @@
+"""Run directory and artifact helpers."""
+
+from cnn_classifier.runs.artifacts import MetricsWriter
+from cnn_classifier.runs.context import RunContext
+
+__all__ = ["MetricsWriter", "RunContext"]

cnn_classifier/runs/artifacts.py

@@ -0,0 +1,50 @@
+"""CSV and JSONL metrics writers."""
+
+from __future__ import annotations
+
+import csv
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from cnn_classifier.runs.context import RunContext
+
+
+@dataclass
+class MetricsWriter:
+    run: RunContext
+    _header: list[str] | None = field(default=None, init=False)
+
+    def _metrics_path(self) -> Path:
+        return self.run.directory / "metrics.jsonl"
+
+    def _history_path(self) -> Path:
+        return self.run.directory / "history.csv"
+
+    def append(self, metrics: Mapping[str, float], step: int) -> None:
+        payload = {"step": int(step), **{key: float(value) for key, value in metrics.items()}}
+        with self._metrics_path().open("a", encoding="utf-8") as handle:
+            handle.write(json.dumps(payload, sort_keys=True) + "\n")
+        if self._header is None:
+            self._header = self._existing_header()
+            if self._header is None:
+                self._header = ["step", *sorted(metrics)]
+                with self._history_path().open("w", encoding="utf-8", newline="") as handle:
+                    writer = csv.DictWriter(handle, fieldnames=self._header)
+                    writer.writeheader()
+        with self._history_path().open("a", encoding="utf-8", newline="") as handle:
+            writer = csv.DictWriter(handle, fieldnames=self._header, extrasaction="ignore")
+            writer.writerow(payload)
+
+    def _existing_header(self) -> list[str] | None:
+        path = self._history_path()
+        if not path.is_file() or path.stat().st_size == 0:
+            return None
+        with path.open(encoding="utf-8", newline="") as handle:
+            reader = csv.reader(handle)
+            try:
+                header = next(reader)
+            except StopIteration:
+                return None
+        return header or None

cnn_classifier/runs/context.py

@@ -0,0 +1,130 @@
+"""Run directories, metadata, and atomic artifact writes."""
+
+from __future__ import annotations
+
+import json
+import os
+import platform
+import re
+import subprocess
+import tempfile
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import torch
+import torchvision
+import yaml
+
+from cnn_classifier.config.schema import RunConfig
+from cnn_classifier.core.device import device_report
+from cnn_classifier.core.hashing import stable_hash
+
+
+def _slug(value: str) -> str:
+    cleaned = re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-")
+    return cleaned or "run"
+
+
+def _git_metadata() -> dict[str, Any]:
+    try:
+        commit = subprocess.run(
+            ["git", "rev-parse", "--short", "HEAD"],
+            capture_output=True,
+            text=True,
+            check=True,
+        ).stdout.strip()
+        branch = subprocess.run(
+            ["git", "branch", "--show-current"],
+            capture_output=True,
+            text=True,
+            check=True,
+        ).stdout.strip()
+        dirty = bool(
+            subprocess.run(
+                ["git", "status", "--porcelain"],
+                capture_output=True,
+                text=True,
+                check=True,
+            ).stdout.strip()
+        )
+    except (OSError, subprocess.CalledProcessError):
+        return {"commit": None, "branch": None, "dirty": None}
+    return {"commit": commit, "branch": branch, "dirty": dirty}
+
+
+@dataclass
+class RunContext:
+    """A self-describing directory for one experiment run."""
+
+    directory: Path
+    config: RunConfig
+    config_hash: str
+
+    @classmethod
+    def create(cls, config: RunConfig, root: Path = Path("runs")) -> RunContext:
+        config_hash = stable_hash(config.model_dump(mode="json"))
+        timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
+        base = root / f"{timestamp}_{_slug(config.experiment.name)}_{config_hash[:8]}"
+        directory = base
+        suffix = 1
+        while directory.exists():
+            directory = base.with_name(f"{base.name}_{suffix}")
+            suffix += 1
+        for child in ("data", "checkpoints", "predictions", "figures"):
+            (directory / child).mkdir(parents=True, exist_ok=True)
+        run = cls(directory=directory, config=config, config_hash=config_hash)
+        run.write_text(
+            "config.resolved.yaml",
+            yaml.safe_dump(config.model_dump(mode="json"), sort_keys=False),
+        )
+        run.write_json(
+            "environment.json",
+            {
+                "python": platform.python_version(),
+                "platform": platform.platform(),
+                "torch": torch.__version__,
+                "torchvision": torchvision.__version__,
+                "device": device_report(),
+            },
+        )
+        run.write_json("git.json", _git_metadata())
+        run.set_status("running")
+        return run
+
+    def _atomic_write(self, relative_path: str, data: bytes) -> Path:
+        path = self.directory / relative_path
+        path.parent.mkdir(parents=True, exist_ok=True)
+        descriptor, temporary = tempfile.mkstemp(
+            prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
+        )
+        temporary_path = Path(temporary)
+        try:
+            with os.fdopen(descriptor, "wb") as handle:
+                handle.write(data)
+            os.replace(temporary_path, path)
+        except BaseException:
+            temporary_path.unlink(missing_ok=True)
+            raise
+        return path
+
+    def write_json(self, relative_path: str, payload: Any) -> Path:
+        text = json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + "\n"
+        return self._atomic_write(relative_path, text.encode("utf-8"))
+
+    def write_text(self, relative_path: str, text: str) -> Path:
+        return self._atomic_write(relative_path, text.encode("utf-8"))
+
+    def set_status(self, status: str, error: str | None = None) -> None:
+        self.write_json(
+            "status.json",
+            {
+                "status": status,
+                "updated_at": datetime.now(UTC).isoformat(),
+                "error": error,
+            },
+        )
+
+    def close(self, status: str = "completed") -> None:
+        self.set_status(status)

cnn_classifier/tracking/__init__.py

@@ -0,0 +1,34 @@
+"""Tracker construction."""
+
+from __future__ import annotations
+
+import contextlib
+
+from cnn_classifier.config.schema import TrackingConfig
+from cnn_classifier.runs.context import RunContext
+from cnn_classifier.tracking.base import Tracker
+from cnn_classifier.tracking.composite import CompositeTracker
+from cnn_classifier.tracking.console import ConsoleTracker
+from cnn_classifier.tracking.csv import CSVTracker
+from cnn_classifier.tracking.tensorboard import TensorBoardTracker
+
+__all__ = [
+    "CSVTracker",
+    "CompositeTracker",
+    "ConsoleTracker",
+    "TensorBoardTracker",
+    "build_trackers",
+]
+
+
+def build_trackers(config: TrackingConfig, run: RunContext) -> CompositeTracker:
+    """Build the enabled trackers for one run."""
+    trackers: list[Tracker] = []
+    if config.console:
+        trackers.append(ConsoleTracker())
+    if config.csv:
+        trackers.append(CSVTracker(run))
+    if config.tensorboard:
+        with contextlib.suppress(RuntimeError):
+            trackers.append(TensorBoardTracker(str(run.directory / "tensorboard")))
+    return CompositeTracker(trackers)

cnn_classifier/tracking/base.py

@@ -0,0 +1,16 @@
+"""Tracker protocol."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Protocol
+
+
+class Tracker(Protocol):
+    def log_metrics(self, metrics: Mapping[str, float], step: int) -> None:
+        """Record one step of metrics."""
+        ...
+
+    def close(self) -> None:
+        """Release resources."""
+        ...

cnn_classifier/tracking/composite.py

@@ -0,0 +1,20 @@
+"""Composite tracker."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+
+from cnn_classifier.tracking.base import Tracker
+
+
+class CompositeTracker:
+    def __init__(self, trackers: Sequence[Tracker]) -> None:
+        self.trackers = list(trackers)
+
+    def log_metrics(self, metrics: Mapping[str, float], step: int) -> None:
+        for tracker in self.trackers:
+            tracker.log_metrics(metrics, step)
+
+    def close(self) -> None:
+        for tracker in self.trackers:
+            tracker.close()

cnn_classifier/tracking/console.py

@@ -0,0 +1,14 @@
+"""Console tracker."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+
+class ConsoleTracker:
+    def log_metrics(self, metrics: Mapping[str, float], step: int) -> None:
+        rendered = " | ".join(f"{key}={value:.4f}" for key, value in metrics.items())
+        print(f"step {step}: {rendered}")
+
+    def close(self) -> None:
+        return None

cnn_classifier/tracking/csv.py

@@ -0,0 +1,19 @@
+"""CSV/JSONL tracker."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+from cnn_classifier.runs.artifacts import MetricsWriter
+from cnn_classifier.runs.context import RunContext
+
+
+class CSVTracker:
+    def __init__(self, run: RunContext) -> None:
+        self.writer = MetricsWriter(run)
+
+    def log_metrics(self, metrics: Mapping[str, float], step: int) -> None:
+        self.writer.append(metrics, step)
+
+    def close(self) -> None:
+        return None

cnn_classifier/tracking/tensorboard.py

@@ -0,0 +1,22 @@
+"""Optional TensorBoard tracker."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+
+class TensorBoardTracker:
+    def __init__(self, log_dir: str) -> None:
+        try:
+            from torch.utils.tensorboard import SummaryWriter
+        except ImportError as exc:
+            raise RuntimeError("TensorBoard tracking requires the 'tracking' extra") from exc
+        self.writer: Any = SummaryWriter(log_dir)
+
+    def log_metrics(self, metrics: Mapping[str, float], step: int) -> None:
+        for key, value in metrics.items():
+            self.writer.add_scalar(key, value, step)
+
+    def close(self) -> None:
+        self.writer.close()