File size: 1,248 Bytes
69928d6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | """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``).
"""
@abstractmethod
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
|