345 lines
12 KiB
Python
345 lines
12 KiB
Python
"""
|
|
Training pipeline for active learning head-only fine-tuning.
|
|
|
|
Loads a pre-trained MobileNetV3 backbone (frozen), replaces the classification
|
|
head, and fine-tunes on versioned crop datasets using AdamW + early stopping.
|
|
All training is FP32 for Tesla P40 compatibility (CC 5.2, no Tensor Cores).
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Dict, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Trainer:
|
|
"""Head-only fine-tuning pipeline for MobileNetV3 classifier."""
|
|
|
|
def __init__(self, config):
|
|
self._cfg = config
|
|
self._models_path = Path(config.get("storage.models_path", "/models"))
|
|
self._candidate_path = self._models_path / "candidate"
|
|
self._candidate_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
al = config.get_section("active_learning")
|
|
tr = al.get("training", {})
|
|
self._epochs: int = int(tr.get("epochs", 20))
|
|
self._batch_size: int = int(tr.get("batch_size", 32))
|
|
self._lr: float = float(tr.get("learning_rate", 1e-3))
|
|
self._weight_decay: float = float(tr.get("weight_decay", 1e-2))
|
|
self._patience: int = int(tr.get("early_stopping_patience", 5))
|
|
self._lr_factor: float = float(tr.get("lr_factor", 0.5))
|
|
self._lr_patience: int = int(tr.get("lr_patience", 3))
|
|
|
|
aug = al.get("augmentation", {})
|
|
self._aug_flip: bool = bool(aug.get("horizontal_flip", True))
|
|
self._aug_jitter: bool = bool(aug.get("color_jitter", True))
|
|
self._aug_affine: bool = bool(aug.get("affine", True))
|
|
self._aug_degrees: int = int(aug.get("affine_degrees", 10))
|
|
self._aug_scale: float = float(aug.get("affine_scale", 0.1))
|
|
|
|
seed = int(config.get("active_learning.seed", 42))
|
|
self._seed = seed
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public API
|
|
# ------------------------------------------------------------------
|
|
|
|
def train(self, dataset_path: str, version: str) -> Optional[str]:
|
|
"""
|
|
Fine-tune the classification head on the versioned dataset.
|
|
|
|
Returns the path to the best checkpoint or None on failure.
|
|
"""
|
|
try:
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader
|
|
|
|
self._seed_everything(self._seed)
|
|
self._assert_cuda_compatibility(torch)
|
|
|
|
train_loader, val_loader = self._build_dataloaders(dataset_path, torch)
|
|
model = self._build_model(nn)
|
|
best_path = self._run_training(model, train_loader, val_loader, version, torch, nn)
|
|
return best_path
|
|
except ImportError as exc:
|
|
logger.error("PyTorch not available: %s", exc)
|
|
return None
|
|
except Exception as exc:
|
|
logger.error("Training failed: %s", exc, exc_info=True)
|
|
return None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Model construction
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_model(self, nn):
|
|
"""Load MobileNetV3-Small backbone, freeze it, attach trainable head."""
|
|
import torchvision.models as models
|
|
|
|
model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1)
|
|
|
|
# Freeze backbone
|
|
for param in model.features.parameters():
|
|
param.requires_grad = False
|
|
|
|
# Replace classification head
|
|
model.classifier = nn.Sequential(
|
|
nn.Linear(576, 256),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.3),
|
|
nn.Linear(256, 2),
|
|
)
|
|
return model
|
|
|
|
# ------------------------------------------------------------------
|
|
# Data loading
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_dataloaders(self, dataset_path: str, torch) -> Tuple:
|
|
from torch.utils.data import DataLoader
|
|
import torchvision.transforms as T
|
|
|
|
train_transform = self._build_train_transform(T)
|
|
val_transform = self._build_val_transform(T)
|
|
|
|
train_dataset = _CropDataset(
|
|
Path(dataset_path) / "crops", split="train",
|
|
transform=train_transform, seed=self._seed,
|
|
)
|
|
val_dataset = _CropDataset(
|
|
Path(dataset_path) / "crops", split="val",
|
|
transform=val_transform, seed=self._seed,
|
|
)
|
|
|
|
g = torch.Generator()
|
|
g.manual_seed(self._seed)
|
|
|
|
train_loader = DataLoader(
|
|
train_dataset,
|
|
batch_size=self._batch_size,
|
|
shuffle=True,
|
|
num_workers=2,
|
|
pin_memory=True,
|
|
generator=g,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_dataset,
|
|
batch_size=self._batch_size,
|
|
shuffle=False,
|
|
num_workers=2,
|
|
pin_memory=True,
|
|
)
|
|
logger.info(
|
|
"Dataset: %d train / %d val samples",
|
|
len(train_dataset), len(val_dataset),
|
|
)
|
|
return train_loader, val_loader
|
|
|
|
def _build_train_transform(self, T):
|
|
ops = [T.Resize((224, 224))]
|
|
if self._aug_flip:
|
|
ops.append(T.RandomHorizontalFlip())
|
|
if self._aug_jitter:
|
|
ops.append(T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1))
|
|
if self._aug_affine:
|
|
ops.append(T.RandomAffine(
|
|
degrees=self._aug_degrees,
|
|
scale=(1.0 - self._aug_scale, 1.0 + self._aug_scale),
|
|
))
|
|
ops += [T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
|
|
return T.Compose(ops)
|
|
|
|
def _build_val_transform(self, T):
|
|
return T.Compose([
|
|
T.Resize((224, 224)),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
# ------------------------------------------------------------------
|
|
# Training loop
|
|
# ------------------------------------------------------------------
|
|
|
|
def _run_training(self, model, train_loader, val_loader, version: str, torch, nn) -> Optional[str]:
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = model.to(device)
|
|
logger.info("Training on device: %s", device)
|
|
|
|
class_weights = self._compute_class_weights(train_loader, torch, device)
|
|
criterion = nn.CrossEntropyLoss(weight=class_weights)
|
|
|
|
optimizer = torch.optim.AdamW(
|
|
filter(lambda p: p.requires_grad, model.parameters()),
|
|
lr=self._lr,
|
|
weight_decay=self._weight_decay,
|
|
betas=(0.9, 0.999),
|
|
)
|
|
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
|
optimizer, mode="max", factor=self._lr_factor, patience=self._lr_patience,
|
|
)
|
|
|
|
best_f1 = -1.0
|
|
best_path: Optional[str] = None
|
|
no_improve = 0
|
|
|
|
for epoch in range(1, self._epochs + 1):
|
|
train_loss = self._train_epoch(model, train_loader, criterion, optimizer, device, torch)
|
|
val_f1 = self._val_epoch(model, val_loader, device, torch)
|
|
scheduler.step(val_f1)
|
|
|
|
ckpt_path = self._candidate_path / f"{version}_epoch_{epoch:02d}.pt"
|
|
torch.save(model.state_dict(), ckpt_path)
|
|
logger.info(
|
|
"Epoch %d/%d train_loss=%.4f val_f1=%.4f",
|
|
epoch, self._epochs, train_loss, val_f1,
|
|
)
|
|
|
|
if val_f1 > best_f1:
|
|
best_f1 = val_f1
|
|
best_path = str(self._candidate_path / f"{version}_best.pt")
|
|
torch.save(model.state_dict(), best_path)
|
|
no_improve = 0
|
|
else:
|
|
no_improve += 1
|
|
if no_improve >= self._patience:
|
|
logger.info("Early stopping at epoch %d (patience=%d)", epoch, self._patience)
|
|
break
|
|
|
|
logger.info("Training complete. Best val_f1=%.4f checkpoint=%s", best_f1, best_path)
|
|
return best_path
|
|
|
|
@staticmethod
|
|
def _train_epoch(model, loader, criterion, optimizer, device, torch) -> float:
|
|
model.train()
|
|
total_loss = 0.0
|
|
for images, labels in loader:
|
|
images, labels = images.to(device), labels.to(device)
|
|
optimizer.zero_grad()
|
|
outputs = model(images)
|
|
loss = criterion(outputs, labels)
|
|
loss.backward()
|
|
optimizer.step()
|
|
total_loss += loss.item() * len(images)
|
|
return total_loss / max(len(loader.dataset), 1)
|
|
|
|
@staticmethod
|
|
def _val_epoch(model, loader, device, torch) -> float:
|
|
from sklearn.metrics import f1_score
|
|
|
|
model.eval()
|
|
all_preds, all_labels = [], []
|
|
with torch.no_grad():
|
|
for images, labels in loader:
|
|
images = images.to(device)
|
|
outputs = model(images)
|
|
preds = outputs.argmax(dim=1).cpu().numpy()
|
|
all_preds.extend(preds)
|
|
all_labels.extend(labels.numpy())
|
|
|
|
return float(f1_score(all_labels, all_preds, average="macro", zero_division=0))
|
|
|
|
@staticmethod
|
|
def _compute_class_weights(train_loader, torch, device):
|
|
"""Compute inverse-frequency class weights to handle imbalance."""
|
|
counts = np.zeros(2, dtype=np.int64)
|
|
for _, labels in train_loader:
|
|
for lbl in labels.numpy():
|
|
counts[int(lbl)] += 1
|
|
total = counts.sum()
|
|
if total == 0 or counts.min() == 0:
|
|
return None
|
|
weights = total / (2.0 * counts.astype(np.float32))
|
|
return torch.tensor(weights, dtype=torch.float32, device=device)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Misc
|
|
# ------------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _assert_cuda_compatibility(torch):
|
|
if torch.cuda.is_available():
|
|
cuda_ver = torch.version.cuda or ""
|
|
torch_ver = torch.__version__
|
|
logger.info("CUDA version: %s PyTorch: %s", cuda_ver, torch_ver)
|
|
major = int(cuda_ver.split(".")[0]) if cuda_ver else 0
|
|
if major > 11:
|
|
logger.warning(
|
|
"CUDA %s detected; system requires CUDA ≤ 11.8 for Tesla P40", cuda_ver
|
|
)
|
|
|
|
@staticmethod
|
|
def _seed_everything(seed: int):
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
try:
|
|
import torch
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Dataset helper
|
|
# ------------------------------------------------------------------
|
|
|
|
class _CropDataset:
|
|
"""Reads face crop images from class_0/ and class_1/ subdirectories."""
|
|
|
|
_EXTENSIONS = {".jpg", ".jpeg", ".png"}
|
|
|
|
def __init__(self, crops_root: Path, split: str, transform, seed: int):
|
|
all_samples = self._discover(crops_root)
|
|
train_samples, val_samples = _split_samples(all_samples, 0.80, seed)
|
|
self._samples = train_samples if split == "train" else val_samples
|
|
self._transform = transform
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._samples)
|
|
|
|
def __getitem__(self, idx: int):
|
|
from PIL import Image
|
|
path, label = self._samples[idx]
|
|
image = Image.open(path).convert("RGB")
|
|
if self._transform:
|
|
image = self._transform(image)
|
|
return image, label
|
|
|
|
def _discover(self, crops_root: Path):
|
|
samples = []
|
|
for label, subdir in ((0, "class_0"), (1, "class_1")):
|
|
d = crops_root / subdir
|
|
if d.is_dir():
|
|
for p in sorted(d.iterdir()):
|
|
if p.suffix.lower() in self._EXTENSIONS:
|
|
samples.append((p, label))
|
|
return samples
|
|
|
|
|
|
def _split_samples(samples, train_frac: float, seed: int):
|
|
"""Stratified split mirroring LabelIngestor logic."""
|
|
rng = random.Random(seed)
|
|
by_label: Dict[int, list] = {}
|
|
for path, label in samples:
|
|
by_label.setdefault(label, []).append((path, label))
|
|
|
|
train, val = [], []
|
|
for group in by_label.values():
|
|
shuffled = list(group)
|
|
rng.shuffle(shuffled)
|
|
split = max(1, int(len(shuffled) * train_frac))
|
|
train.extend(shuffled[:split])
|
|
val.extend(shuffled[split:])
|
|
|
|
return train, val
|