Initial commit
This commit is contained in:
@@ -0,0 +1,265 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
try:
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision.io import read_image
|
||||
from torchvision.transforms import v2
|
||||
except ModuleNotFoundError as exc: # pragma: no cover - optional runtime dependency
|
||||
raise SystemExit(
|
||||
"A9 requires torch and torchvision in the experiment .venv. "
|
||||
"Install them before running run_a9_deep_baselines.py."
|
||||
) from exc
|
||||
|
||||
from evaluation.metrics import compute_binary_metrics, optimize_threshold
|
||||
|
||||
|
||||
IMAGE_SIZE = 128
|
||||
RANDOM_SEED = 42
|
||||
|
||||
|
||||
def seed_everything(seed: int = RANDOM_SEED) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def choose_device() -> torch.device:
|
||||
if torch.backends.mps.is_available():
|
||||
return torch.device("mps")
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda")
|
||||
return torch.device("cpu")
|
||||
|
||||
|
||||
class CropDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
frame: pd.DataFrame,
|
||||
image_column: str,
|
||||
feature_columns: list[str] | None = None,
|
||||
train_mode: bool = False,
|
||||
) -> None:
|
||||
self.frame = frame.reset_index(drop=True)
|
||||
self.image_column = image_column
|
||||
self.feature_columns = feature_columns or []
|
||||
self.train_mode = train_mode
|
||||
self.transform = v2.Compose(
|
||||
[
|
||||
v2.Resize((IMAGE_SIZE, IMAGE_SIZE), antialias=True),
|
||||
v2.ToDtype(torch.float32, scale=True),
|
||||
v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.frame)
|
||||
|
||||
def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
|
||||
row = self.frame.iloc[index]
|
||||
image = read_image(str(row[self.image_column]))
|
||||
image = self.transform(image)
|
||||
sample = {
|
||||
"image": image,
|
||||
"target": torch.tensor(int(row["target_binary"]), dtype=torch.float32),
|
||||
}
|
||||
if self.feature_columns:
|
||||
tabular_values = row[self.feature_columns].to_numpy(dtype=np.float32, copy=True)
|
||||
tabular_values = np.nan_to_num(tabular_values, nan=0.0, posinf=0.0, neginf=0.0)
|
||||
sample["tabular"] = torch.tensor(tabular_values, dtype=torch.float32)
|
||||
return sample
|
||||
|
||||
|
||||
class SmallCNN(nn.Module):
|
||||
def __init__(self, tabular_dim: int = 0) -> None:
|
||||
super().__init__()
|
||||
self.features = nn.Sequential(
|
||||
nn.Conv2d(3, 32, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(2),
|
||||
nn.Conv2d(32, 64, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.MaxPool2d(2),
|
||||
nn.Conv2d(64, 128, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.AdaptiveAvgPool2d((1, 1)),
|
||||
)
|
||||
self.image_head = nn.Sequential(
|
||||
nn.Flatten(),
|
||||
nn.Linear(128, 64),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout(0.2),
|
||||
)
|
||||
self.tabular_dim = tabular_dim
|
||||
fusion_dim = 64 + tabular_dim
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(fusion_dim, 64),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout(0.2),
|
||||
nn.Linear(64, 1),
|
||||
)
|
||||
|
||||
def forward(self, image: torch.Tensor, tabular: torch.Tensor | None = None) -> torch.Tensor:
|
||||
image_embedding = self.image_head(self.features(image))
|
||||
if self.tabular_dim > 0:
|
||||
assert tabular is not None
|
||||
logits = self.classifier(torch.cat([image_embedding, tabular], dim=1))
|
||||
else:
|
||||
logits = self.classifier(image_embedding)
|
||||
return logits.squeeze(1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingResult:
|
||||
model_path: str
|
||||
history: list[dict[str, float]]
|
||||
threshold: float
|
||||
test_predictions: pd.DataFrame
|
||||
test_metrics: dict[str, float]
|
||||
|
||||
|
||||
def _make_loader(frame: pd.DataFrame, image_column: str, feature_columns: list[str] | None, shuffle: bool) -> DataLoader:
|
||||
dataset = CropDataset(frame, image_column=image_column, feature_columns=feature_columns, train_mode=shuffle)
|
||||
return DataLoader(
|
||||
dataset,
|
||||
batch_size=128,
|
||||
shuffle=shuffle,
|
||||
num_workers=0,
|
||||
pin_memory=False,
|
||||
)
|
||||
|
||||
|
||||
def _run_epoch(
|
||||
model: nn.Module,
|
||||
loader: DataLoader,
|
||||
device: torch.device,
|
||||
optimizer: torch.optim.Optimizer | None,
|
||||
) -> tuple[float, np.ndarray, np.ndarray]:
|
||||
if optimizer is None:
|
||||
model.eval()
|
||||
else:
|
||||
model.train()
|
||||
criterion = nn.BCEWithLogitsLoss()
|
||||
losses: list[float] = []
|
||||
all_targets: list[np.ndarray] = []
|
||||
all_scores: list[np.ndarray] = []
|
||||
|
||||
for batch in loader:
|
||||
image = batch["image"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
tabular = batch.get("tabular")
|
||||
if tabular is not None:
|
||||
tabular = tabular.to(device)
|
||||
|
||||
if optimizer is not None:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
logits = model(image, tabular)
|
||||
loss = criterion(logits, target)
|
||||
if optimizer is not None:
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
losses.append(float(loss.detach().cpu()))
|
||||
all_targets.append(target.detach().cpu().numpy())
|
||||
all_scores.append(torch.sigmoid(logits).detach().cpu().numpy())
|
||||
|
||||
return float(np.mean(losses)), np.concatenate(all_targets), np.concatenate(all_scores)
|
||||
|
||||
|
||||
def train_deep_baseline(
|
||||
train_df: pd.DataFrame,
|
||||
val_df: pd.DataFrame,
|
||||
test_df: pd.DataFrame,
|
||||
image_column: str,
|
||||
feature_columns: list[str] | None,
|
||||
checkpoint_path: Path,
|
||||
max_epochs: int = 12,
|
||||
patience: int = 3,
|
||||
) -> TrainingResult:
|
||||
seed_everything()
|
||||
device = choose_device()
|
||||
model = SmallCNN(tabular_dim=len(feature_columns or [])).to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
||||
|
||||
train_loader = _make_loader(train_df, image_column=image_column, feature_columns=feature_columns, shuffle=True)
|
||||
val_loader = _make_loader(val_df, image_column=image_column, feature_columns=feature_columns, shuffle=False)
|
||||
test_loader = _make_loader(test_df, image_column=image_column, feature_columns=feature_columns, shuffle=False)
|
||||
|
||||
best_state: dict[str, torch.Tensor] | None = None
|
||||
best_val_pr_auc = -np.inf
|
||||
best_threshold = 0.5
|
||||
epochs_without_improvement = 0
|
||||
history: list[dict[str, float]] = []
|
||||
|
||||
for epoch_idx in range(1, max_epochs + 1):
|
||||
train_loss, train_true, train_score = _run_epoch(model, train_loader, device, optimizer)
|
||||
val_loss, val_true, val_score = _run_epoch(model, val_loader, device, optimizer=None)
|
||||
threshold = optimize_threshold(val_true, val_score)
|
||||
train_metrics = compute_binary_metrics(train_true, train_score, 0.5)
|
||||
val_metrics = compute_binary_metrics(val_true, val_score, threshold)
|
||||
history.append(
|
||||
{
|
||||
"epoch": epoch_idx,
|
||||
"train_loss": train_loss,
|
||||
"val_loss": val_loss,
|
||||
"train_pr_auc": train_metrics.pr_auc,
|
||||
"val_pr_auc": val_metrics.pr_auc,
|
||||
"val_recall_blast": val_metrics.recall_blast,
|
||||
}
|
||||
)
|
||||
if val_metrics.pr_auc > best_val_pr_auc:
|
||||
best_val_pr_auc = val_metrics.pr_auc
|
||||
best_threshold = threshold
|
||||
best_state = {key: value.detach().cpu().clone() for key, value in model.state_dict().items()}
|
||||
epochs_without_improvement = 0
|
||||
else:
|
||||
epochs_without_improvement += 1
|
||||
if epochs_without_improvement >= patience:
|
||||
break
|
||||
|
||||
if best_state is None:
|
||||
raise ValueError("Deep model did not produce a valid checkpoint.")
|
||||
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(best_state, checkpoint_path)
|
||||
model.load_state_dict(best_state)
|
||||
|
||||
_, test_true, test_score = _run_epoch(model, test_loader, device, optimizer=None)
|
||||
test_metrics = compute_binary_metrics(test_true, test_score, best_threshold)
|
||||
test_predictions = test_df.copy()
|
||||
test_predictions["y_true"] = test_true
|
||||
test_predictions["y_score"] = test_score
|
||||
test_predictions["y_pred"] = (test_score >= best_threshold).astype(int)
|
||||
|
||||
return TrainingResult(
|
||||
model_path=str(checkpoint_path),
|
||||
history=history,
|
||||
threshold=best_threshold,
|
||||
test_predictions=test_predictions,
|
||||
test_metrics={
|
||||
"pr_auc": test_metrics.pr_auc,
|
||||
"roc_auc": test_metrics.roc_auc,
|
||||
"balanced_accuracy": test_metrics.balanced_accuracy,
|
||||
"macro_f1": test_metrics.macro_f1,
|
||||
"recall_blast": test_metrics.recall_blast,
|
||||
"specificity": test_metrics.specificity,
|
||||
"matthews_corrcoef": test_metrics.matthews_corrcoef,
|
||||
"brier_score": test_metrics.brier_score,
|
||||
"expected_calibration_error": test_metrics.expected_calibration_error,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def save_training_history(history_payload: dict[str, list[dict[str, float]]], path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(history_payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
Reference in New Issue
Block a user