Spaces:

Moni38
/

smartvision-ai

Sleeping

smartvision-ai / src /models /base_classifier.py

Priya Monisha

fix: pass ruff lint+format, clean README, remove stray script

fa0ce5b 5 months ago

11.5 kB

# filename: src/models/base_classifier.py
# purpose:  Shared training/evaluation engine for all 4 CNN classifiers
# version:  2.0

import logging
import time
from pathlib import Path
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)

try:
    import torch
    import torch.nn as nn
    from sklearn.metrics import (
        accuracy_score,
        precision_score,
        recall_score,
        f1_score,
    )

    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


def _freeze_backbone_bn(model: "nn.Module") -> None:
    """
    Set BatchNorm layers to eval() if their affine params are frozen.
    Must be called AFTER model.train() each epoch — model.train() resets all
    layers to training mode first; we then selectively revert frozen BN layers
    back to eval().
    Why module.weight.requires_grad:
      When a backbone block is frozen (param.requires_grad=False), its BN's
      affine weight inherits requires_grad=False. This is the reliable proxy
      for "is this BN inside a frozen block?"
    Why not module.parameters():
      running_mean and running_var are buffers, not parameters — they never
      appear in .parameters(). They update whenever module.training==True,
      regardless of requires_grad. Checking requires_grad alone is insufficient.
    Why affine=False BN (weight is None) is always frozen:
      No learnable params at all — either frozen or has no learning signal.
      Either way, updating running stats from current mini-batches corrupts
      pretrained statistics.
    Handles mixed-freeze correctly:
      e.g. MobileNetV2 features[0:13] frozen, features[14:] unfrozen:
      BN in features[0:13] get eval(), BN in features[14:] stays train().
    """
    for module in model.modules():
        if not isinstance(module, (nn.BatchNorm2d, nn.BatchNorm1d)):
            continue
        if module.weight is None:
            module.eval()  # affine=False: always freeze stats
        elif not module.weight.requires_grad:
            module.eval()  # affine frozen: parent block is frozen


def train_one_epoch(
    model,
    loader,
    optimizer,
    criterion,
    device,
    scaler=None,
    freeze_bn: bool = False,
    grad_clip: Optional[float] = None,
) -> dict:
    """Run one training epoch.
    freeze_bn: call _freeze_backbone_bn() after model.train(). Use True whenever
               any backbone layers are frozen. Per-module check handles partial
               unfreeze (e.g., features[14:] only) correctly.
    grad_clip: max norm for gradient clipping. 1.0 for both phases. Do not use
               0.5 — too aggressive, suppresses signal in newly unfrozen layers.
    """
    model.train()
    # CRITICAL: call AFTER model.train() — model.train() resets everything to
    # training mode first; we then selectively revert frozen BN layers to eval().
    if freeze_bn:
        _freeze_backbone_bn(model)

    total_loss = 0.0
    correct = 0
    total = 0

    for images, labels in loader:
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)
        optimizer.zero_grad(set_to_none=True)

        if scaler is not None:
            with torch.cuda.amp.autocast():
                outputs = model(images)
                loss = criterion(outputs, labels)
            scaler.scale(loss).backward()

            if grad_clip is not None:
                # unscale_ MUST precede clip_grad_norm_ in AMP path —
                # gradients are stored scaled; unscale converts to true grads first
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(
                    (p for p in model.parameters() if p.requires_grad),
                    max_norm=grad_clip,
                )

            scaler.step(optimizer)
            scaler.update()
        else:
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()

            if grad_clip is not None:
                torch.nn.utils.clip_grad_norm_(
                    (p for p in model.parameters() if p.requires_grad),
                    max_norm=grad_clip,
                )

            optimizer.step()

        total_loss += loss.item() * images.size(0)
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += images.size(0)

    return {
        "loss": total_loss / total,
        "accuracy": correct / total,
    }


def evaluate(
    model,
    loader,
    criterion,
    device,
    num_classes: int = 25,
) -> dict:
    """Evaluate model. Returns loss, accuracy, precision, recall, F1 (macro)."""
    model.eval()
    total_loss = 0.0
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for images, labels in loader:
            images = images.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            outputs = model(images)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * images.size(0)
            preds = outputs.argmax(dim=1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    n = len(all_labels)
    return {
        "loss": total_loss / n,
        "accuracy": accuracy_score(all_labels, all_preds),
        "precision": precision_score(
            all_labels, all_preds, average="macro", zero_division=0
        ),
        "recall": recall_score(all_labels, all_preds, average="macro", zero_division=0),
        "f1": f1_score(all_labels, all_preds, average="macro", zero_division=0),
    }


def train(
    model,
    train_loader,
    val_loader,
    optimizer,
    scheduler,
    criterion,
    device,
    epochs: int,
    patience: int = 5,
    scaler=None,
    model_name: str = "model",
    save_path: Optional[Path] = None,
    freeze_bn: bool = False,
    grad_clip: Optional[float] = None,
) -> dict:
    """Full training loop with early stopping (monitors val_accuracy).
    Monitors val_accuracy, not val_loss. Rationale: label_smoothing=0.1 inflates
    loss by ~0.31 nats uniformly — loss trends correctly but absolute values are
    misleading for early stopping. val_accuracy is scale-invariant.
    Reloads best checkpoint before returning so caller always gets best model,
    not last epoch.
    """
    history = {
        "train_loss": [],
        "train_acc": [],
        "val_loss": [],
        "val_acc": [],
        "val_f1": [],
    }
    best_val_acc = 0.0
    patience_count = 0

    # Resolve checkpoint path ONCE before the loop — avoid reassigning inside loop
    checkpoint_path = None
    if save_path is not None:
        checkpoint_path = Path(save_path)
        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)

    for epoch in range(1, epochs + 1):
        t0 = time.time()

        train_metrics = train_one_epoch(
            model,
            train_loader,
            optimizer,
            criterion,
            device,
            scaler=scaler,
            freeze_bn=freeze_bn,
            grad_clip=grad_clip,
        )
        val_metrics = evaluate(model, val_loader, criterion, device)

        if scheduler is not None:
            if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                scheduler.step(val_metrics["loss"])
            else:
                scheduler.step()

        history["train_loss"].append(round(train_metrics["loss"], 6))
        history["train_acc"].append(round(train_metrics["accuracy"], 6))
        history["val_loss"].append(round(val_metrics["loss"], 6))
        history["val_acc"].append(round(val_metrics["accuracy"], 6))
        history["val_f1"].append(round(val_metrics["f1"], 6))

        elapsed = time.time() - t0
        logger.info(
            f"[{model_name}] Epoch {epoch:3d}/{epochs} | "
            f"train_acc={train_metrics['accuracy']:.4f} | "
            f"val_acc={val_metrics['accuracy']:.4f} | "
            f"val_f1={val_metrics['f1']:.4f} | {elapsed:.1f}s"
        )

        if val_metrics["accuracy"] > best_val_acc:
            best_val_acc = val_metrics["accuracy"]
            patience_count = 0
            if checkpoint_path is not None:
                # Cast to Python native types — numpy.float64 (from accuracy_score)
                # is rejected by weights_only=True during reload
                torch.save(
                    {
                        "epoch": int(epoch),
                        "model_state_dict": model.state_dict(),
                        "val_acc": float(best_val_acc),
                        "val_loss": float(val_metrics["loss"]),
                        "val_f1": float(val_metrics["f1"]),
                    },
                    checkpoint_path,
                )
                logger.info(
                    f"  Saved best: {checkpoint_path.name} (val_acc={best_val_acc:.4f})"
                )
        else:
            patience_count += 1
            if patience_count >= patience:
                logger.info(
                    f"[{model_name}] Early stopping at epoch {epoch} (patience={patience})"
                )
                break

    # Reload best checkpoint so caller receives best model, not last epoch
    if checkpoint_path is not None and checkpoint_path.exists():
        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=True)
        model.load_state_dict(ckpt["model_state_dict"])
        logger.info(f"Reloaded best checkpoint: val_acc={ckpt['val_acc']:.4f}")

    history["best_val_acc"] = best_val_acc
    return history


# ── Inference helpers ─────────────────────────────────────────────────────────


def save_model(model, path: Path) -> None:
    """Save state_dict only — never full model pickle (Rule 5)."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(model.state_dict(), path)


def load_model(model, path: Path, device: str = "cpu"):
    """Load weights with weights_only=True (Rule 36). Handles both raw state_dict
    and checkpoint dict (keys: model_state_dict, epoch, val_acc, ...) formats."""
    state_dict = torch.load(path, map_location=device, weights_only=True)
    if isinstance(state_dict, dict) and "model_state_dict" in state_dict:
        state_dict = state_dict["model_state_dict"]
    model.load_state_dict(state_dict)
    model.eval()
    return model


def benchmark_inference(
    model, device, image_size: int = 224, n: int = 100, warmup: int = 10
) -> float:
    """Returns mean inference time in milliseconds over n runs after warmup.
    CUDA ops are async — synchronize before/after each timing window so that
    time.perf_counter() measures actual compute, not just kernel scheduling.
    """
    model.eval()
    dummy = torch.randn(1, 3, image_size, image_size).to(device)
    is_cuda = str(device).startswith("cuda")

    with torch.no_grad():
        for _ in range(warmup):
            model(dummy)
        if is_cuda:
            torch.cuda.synchronize()

    times = []
    with torch.no_grad():
        for _ in range(n):
            if is_cuda:
                torch.cuda.synchronize()
            t0 = time.perf_counter()
            model(dummy)
            if is_cuda:
                torch.cuda.synchronize()
            times.append((time.perf_counter() - t0) * 1000)

    return float(np.mean(times))

	# filename: src/models/base_classifier.py
	# purpose: Shared training/evaluation engine for all 4 CNN classifiers
	# version: 2.0

	import logging
	import time
	from pathlib import Path
	from typing import Optional

	import numpy as np

	logger = logging.getLogger(__name__)

	try:
	import torch
	import torch.nn as nn
	from sklearn.metrics import (
	accuracy_score,
	precision_score,
	recall_score,
	f1_score,
	)

	TORCH_AVAILABLE = True
	except ImportError:
	TORCH_AVAILABLE = False


	def _freeze_backbone_bn(model: "nn.Module") -> None:
	"""
	Set BatchNorm layers to eval() if their affine params are frozen.

	Must be called AFTER model.train() each epoch — model.train() resets all
	layers to training mode first; we then selectively revert frozen BN layers
	back to eval().

	Why module.weight.requires_grad:
	When a backbone block is frozen (param.requires_grad=False), its BN's
	affine weight inherits requires_grad=False. This is the reliable proxy
	for "is this BN inside a frozen block?"

	Why not module.parameters():
	running_mean and running_var are buffers, not parameters — they never
	appear in .parameters(). They update whenever module.training==True,
	regardless of requires_grad. Checking requires_grad alone is insufficient.

	Why affine=False BN (weight is None) is always frozen:
	No learnable params at all — either frozen or has no learning signal.
	Either way, updating running stats from current mini-batches corrupts
	pretrained statistics.

	Handles mixed-freeze correctly:
	e.g. MobileNetV2 features[0:13] frozen, features[14:] unfrozen:
	BN in features[0:13] get eval(), BN in features[14:] stays train().
	"""
	for module in model.modules():
	if not isinstance(module, (nn.BatchNorm2d, nn.BatchNorm1d)):
	continue
	if module.weight is None:
	module.eval() # affine=False: always freeze stats
	elif not module.weight.requires_grad:
	module.eval() # affine frozen: parent block is frozen


	def train_one_epoch(
	model,
	loader,
	optimizer,
	criterion,
	device,
	scaler=None,
	freeze_bn: bool = False,
	grad_clip: Optional[float] = None,
	) -> dict:
	"""Run one training epoch.

	freeze_bn: call _freeze_backbone_bn() after model.train(). Use True whenever
	any backbone layers are frozen. Per-module check handles partial
	unfreeze (e.g., features[14:] only) correctly.
	grad_clip: max norm for gradient clipping. 1.0 for both phases. Do not use
	0.5 — too aggressive, suppresses signal in newly unfrozen layers.
	"""
	model.train()
	# CRITICAL: call AFTER model.train() — model.train() resets everything to
	# training mode first; we then selectively revert frozen BN layers to eval().
	if freeze_bn:
	_freeze_backbone_bn(model)

	total_loss = 0.0
	correct = 0
	total = 0

	for images, labels in loader:
	images = images.to(device, non_blocking=True)
	labels = labels.to(device, non_blocking=True)
	optimizer.zero_grad(set_to_none=True)

	if scaler is not None:
	with torch.cuda.amp.autocast():
	outputs = model(images)
	loss = criterion(outputs, labels)
	scaler.scale(loss).backward()

	if grad_clip is not None:
	# unscale_ MUST precede clip_grad_norm_ in AMP path —
	# gradients are stored scaled; unscale converts to true grads first
	scaler.unscale_(optimizer)
	torch.nn.utils.clip_grad_norm_(
	(p for p in model.parameters() if p.requires_grad),
	max_norm=grad_clip,
	)

	scaler.step(optimizer)
	scaler.update()
	else:
	outputs = model(images)
	loss = criterion(outputs, labels)
	loss.backward()

	if grad_clip is not None:
	torch.nn.utils.clip_grad_norm_(
	(p for p in model.parameters() if p.requires_grad),
	max_norm=grad_clip,
	)

	optimizer.step()

	total_loss += loss.item() * images.size(0)
	preds = outputs.argmax(dim=1)
	correct += (preds == labels).sum().item()
	total += images.size(0)

	return {
	"loss": total_loss / total,
	"accuracy": correct / total,
	}


	def evaluate(
	model,
	loader,
	criterion,
	device,
	num_classes: int = 25,
	) -> dict:
	"""Evaluate model. Returns loss, accuracy, precision, recall, F1 (macro)."""
	model.eval()
	total_loss = 0.0
	all_preds = []
	all_labels = []

	with torch.no_grad():
	for images, labels in loader:
	images = images.to(device, non_blocking=True)
	labels = labels.to(device, non_blocking=True)
	outputs = model(images)
	loss = criterion(outputs, labels)
	total_loss += loss.item() * images.size(0)
	preds = outputs.argmax(dim=1)
	all_preds.extend(preds.cpu().numpy())
	all_labels.extend(labels.cpu().numpy())

	n = len(all_labels)
	return {
	"loss": total_loss / n,
	"accuracy": accuracy_score(all_labels, all_preds),
	"precision": precision_score(
	all_labels, all_preds, average="macro", zero_division=0
	),
	"recall": recall_score(all_labels, all_preds, average="macro", zero_division=0),
	"f1": f1_score(all_labels, all_preds, average="macro", zero_division=0),
	}


	def train(
	model,
	train_loader,
	val_loader,
	optimizer,
	scheduler,
	criterion,
	device,
	epochs: int,
	patience: int = 5,
	scaler=None,
	model_name: str = "model",
	save_path: Optional[Path] = None,
	freeze_bn: bool = False,
	grad_clip: Optional[float] = None,
	) -> dict:
	"""Full training loop with early stopping (monitors val_accuracy).

	Monitors val_accuracy, not val_loss. Rationale: label_smoothing=0.1 inflates
	loss by ~0.31 nats uniformly — loss trends correctly but absolute values are
	misleading for early stopping. val_accuracy is scale-invariant.

	Reloads best checkpoint before returning so caller always gets best model,
	not last epoch.
	"""
	history = {
	"train_loss": [],
	"train_acc": [],
	"val_loss": [],
	"val_acc": [],
	"val_f1": [],
	}
	best_val_acc = 0.0
	patience_count = 0

	# Resolve checkpoint path ONCE before the loop — avoid reassigning inside loop
	checkpoint_path = None
	if save_path is not None:
	checkpoint_path = Path(save_path)
	checkpoint_path.parent.mkdir(parents=True, exist_ok=True)

	for epoch in range(1, epochs + 1):
	t0 = time.time()

	train_metrics = train_one_epoch(
	model,
	train_loader,
	optimizer,
	criterion,
	device,
	scaler=scaler,
	freeze_bn=freeze_bn,
	grad_clip=grad_clip,
	)
	val_metrics = evaluate(model, val_loader, criterion, device)

	if scheduler is not None:
	if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
	scheduler.step(val_metrics["loss"])
	else:
	scheduler.step()

	history["train_loss"].append(round(train_metrics["loss"], 6))
	history["train_acc"].append(round(train_metrics["accuracy"], 6))
	history["val_loss"].append(round(val_metrics["loss"], 6))
	history["val_acc"].append(round(val_metrics["accuracy"], 6))
	history["val_f1"].append(round(val_metrics["f1"], 6))

	elapsed = time.time() - t0
	logger.info(
	f"[{model_name}] Epoch {epoch:3d}/{epochs} \| "
	f"train_acc={train_metrics['accuracy']:.4f} \| "
	f"val_acc={val_metrics['accuracy']:.4f} \| "
	f"val_f1={val_metrics['f1']:.4f} \| {elapsed:.1f}s"
	)

	if val_metrics["accuracy"] > best_val_acc:
	best_val_acc = val_metrics["accuracy"]
	patience_count = 0
	if checkpoint_path is not None:
	# Cast to Python native types — numpy.float64 (from accuracy_score)
	# is rejected by weights_only=True during reload
	torch.save(
	{
	"epoch": int(epoch),
	"model_state_dict": model.state_dict(),
	"val_acc": float(best_val_acc),
	"val_loss": float(val_metrics["loss"]),
	"val_f1": float(val_metrics["f1"]),
	},
	checkpoint_path,
	)
	logger.info(
	f" Saved best: {checkpoint_path.name} (val_acc={best_val_acc:.4f})"
	)
	else:
	patience_count += 1
	if patience_count >= patience:
	logger.info(
	f"[{model_name}] Early stopping at epoch {epoch} (patience={patience})"
	)
	break

	# Reload best checkpoint so caller receives best model, not last epoch
	if checkpoint_path is not None and checkpoint_path.exists():
	ckpt = torch.load(checkpoint_path, map_location=device, weights_only=True)
	model.load_state_dict(ckpt["model_state_dict"])
	logger.info(f"Reloaded best checkpoint: val_acc={ckpt['val_acc']:.4f}")

	history["best_val_acc"] = best_val_acc
	return history


	# ── Inference helpers ─────────────────────────────────────────────────────────


	def save_model(model, path: Path) -> None:
	"""Save state_dict only — never full model pickle (Rule 5)."""
	path = Path(path)
	path.parent.mkdir(parents=True, exist_ok=True)
	torch.save(model.state_dict(), path)


	def load_model(model, path: Path, device: str = "cpu"):
	"""Load weights with weights_only=True (Rule 36). Handles both raw state_dict
	and checkpoint dict (keys: model_state_dict, epoch, val_acc, ...) formats."""
	state_dict = torch.load(path, map_location=device, weights_only=True)
	if isinstance(state_dict, dict) and "model_state_dict" in state_dict:
	state_dict = state_dict["model_state_dict"]
	model.load_state_dict(state_dict)
	model.eval()
	return model


	def benchmark_inference(
	model, device, image_size: int = 224, n: int = 100, warmup: int = 10
	) -> float:
	"""Returns mean inference time in milliseconds over n runs after warmup.

	CUDA ops are async — synchronize before/after each timing window so that
	time.perf_counter() measures actual compute, not just kernel scheduling.
	"""
	model.eval()
	dummy = torch.randn(1, 3, image_size, image_size).to(device)
	is_cuda = str(device).startswith("cuda")

	with torch.no_grad():
	for _ in range(warmup):
	model(dummy)
	if is_cuda:
	torch.cuda.synchronize()

	times = []
	with torch.no_grad():
	for _ in range(n):
	if is_cuda:
	torch.cuda.synchronize()
	t0 = time.perf_counter()
	model(dummy)
	if is_cuda:
	torch.cuda.synchronize()
	times.append((time.perf_counter() - t0) * 1000)

	return float(np.mean(times))