| """Abstract base model that all SoftASP models must implement.""" | |
| from __future__ import annotations | |
| from abc import ABC, abstractmethod | |
| from typing import Any | |
| import torch | |
| import torch.nn as nn | |
| class BaseSoftASPModel(nn.Module, ABC): | |
| """Base class for all SoftASP models. | |
| Every model receives a batch dict and returns an output dict containing | |
| at minimum ``loss`` and ``logits`` (or ``predictions``). | |
| """ | |
| def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: | |
| """Forward pass. | |
| Args: | |
| batch: Dict from DataLoader with keys like ``images``, ``labels``, | |
| and optional masks. | |
| Returns: | |
| Dict with at least ``loss`` and ``logits`` keys. | |
| """ | |
| def count_parameters(self) -> dict[str, int]: | |
| """Return total and trainable parameter counts.""" | |
| total = sum(p.numel() for p in self.parameters()) | |
| trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) | |
| return {"total_params": total, "trainable_params": trainable} | |
| def get_backbone(self) -> nn.Module | None: | |
| """Optional component controlled by the trainer's freeze/ramp schedule.""" | |
| return None | |