train_with_emnist.py

"""
train_with_emnist.py
====================
Fine-tune the CRNN model with EMNIST character data.

FIXES vs old version:
  - Phase 1: CNN FROZEN — only RNN+FC trained (prevents catastrophic forgetting)
  - Phase 2: Full model at 10x lower LR for final polish
  - log_softmax applied before CTCLoss (was missing — caused garbage loss)
  - Loads from best_model.pth (synthetic, 0.12% CER baseline)
  - Saves best_model_emnist.pth only when val improves
"""

import os
import sys
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, ConcatDataset

sys.path.append('.')
from crnn_model import get_crnn_model
from dataset import CivilRegistryDataset, collate_fn

print("=" * 55)
print("Fine-tuning CRNN with EMNIST dataset")
print("=" * 55)

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Device: {DEVICE}")

emnist_dataset = CivilRegistryDataset(
    data_dir='data/train',
    annotations_file='data/emnist_train_annotations.json',
    img_height=64, img_width=512, augment=True
)
# FIXED: mix synthetic data in so the model never forgets multi-word sequences
synth_dataset = CivilRegistryDataset(
    data_dir='data/train',
    annotations_file='data/train_annotations.json',
    img_height=64, img_width=512, augment=True
)
train_dataset = emnist_dataset  # keep reference for char_to_idx / num_chars
mixed_train   = ConcatDataset([emnist_dataset, synth_dataset])
val_dataset = CivilRegistryDataset(
    data_dir='data/val',
    annotations_file='data/val_annotations.json',  # FIXED: was emnist_val — must match real task
    img_height=64, img_width=512, augment=False
)
print(f"EMNIST train  : {len(emnist_dataset)}")
print(f"Synthetic train: {len(synth_dataset)}")
print(f"Mixed train   : {len(mixed_train)}")
print(f"Val           : {len(val_dataset)}")

train_loader = DataLoader(mixed_train, batch_size=32, shuffle=True,
                          num_workers=0, collate_fn=collate_fn)
val_loader   = DataLoader(val_dataset,   batch_size=32, shuffle=False,
                          num_workers=0, collate_fn=collate_fn)

# ── Load best synthetic checkpoint ───────────────────────────
BASE = 'checkpoints/best_model.pth'
if not os.path.exists(BASE):
    print(f"ERROR: {BASE} not found. Run: python train.py")
    sys.exit(1)

ckpt   = torch.load(BASE, map_location=DEVICE, weights_only=False)
config = ckpt.get('config', {})

model = get_crnn_model(
    model_type      = config.get('model_type', 'standard'),
    img_height      = config.get('img_height', 64),
    num_chars       = train_dataset.num_chars,
    hidden_size     = config.get('hidden_size', 128),
    num_lstm_layers = config.get('num_lstm_layers', 1),
).to(DEVICE)

missing, _ = model.load_state_dict(ckpt['model_state_dict'], strict=False)
if missing:
    print(f"  Note: {len(missing)} layers re-initialized (expected for fc layer)")
print(f"  Loaded epoch {ckpt.get('epoch')} "
      f"(val_loss={ckpt.get('val_loss', ckpt.get('val_cer', 0)):.4f})")

criterion = torch.nn.CTCLoss(blank=0, reduction='mean', zero_infinity=True)


def run_epoch(loader, training, optimizer=None):
    model.train() if training else model.eval()
    total, n = 0, 0
    ctx = torch.enable_grad() if training else torch.no_grad()
    with ctx:
        for images, targets, target_lengths, _ in loader:
            images        = images.to(DEVICE)
            batch_size    = images.size(0)
            if training:
                optimizer.zero_grad()
            # CRITICAL: log_softmax before CTCLoss
            outputs       = F.log_softmax(model(images), dim=2)
            seq_len       = outputs.size(0)
            input_lengths = torch.full((batch_size,), seq_len, dtype=torch.long)
            loss = criterion(outputs, targets, input_lengths, target_lengths)
            if not torch.isnan(loss) and not torch.isinf(loss):
                if training:
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
                    optimizer.step()
                total += loss.item()
                n     += 1
    return total / max(n, 1)


def run_phase(num, epochs, lr, freeze_cnn, patience):
    print(f"\n{'='*55}")
    print(f"  PHASE {num} — "
          f"{'CNN FROZEN  (RNN+FC only)' if freeze_cnn else 'FULL MODEL  (all layers)'}"
          f"   LR={lr}")
    print(f"{'='*55}")

    # Freeze or unfreeze CNN
    for name, param in model.named_parameters():
        param.requires_grad = not (freeze_cnn and 'cnn' in name)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"  Trainable params : {trainable:,}")

    opt      = optim.Adam(
        filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
    sched    = optim.lr_scheduler.ReduceLROnPlateau(opt, patience=3, factor=0.5)
    best     = float('inf')
    counter  = 0

    for epoch in range(1, epochs + 1):
        tr = run_epoch(train_loader, True,  opt)
        vl = run_epoch(val_loader,   False, None)
        sched.step(vl)

        if vl < best:
            best    = vl
            counter = 0
            torch.save({
                'model_state_dict': model.state_dict(),
                'config':           config,
                'char_to_idx':      train_dataset.char_to_idx,
                'idx_to_char':      train_dataset.idx_to_char,
                'epoch':            epoch,
                'val_loss':         vl,   # FIXED: renamed from val_cer — this is val loss, not CER%
            }, 'checkpoints/best_model_emnist.pth')
            print(f"  Epoch {epoch:02d}/{epochs}  Train={tr:.4f}  Val={vl:.4f}  <- saved")
        else:
            counter += 1
            print(f"  Epoch {epoch:02d}/{epochs}  Train={tr:.4f}  Val={vl:.4f}"
                  f"  (patience {counter}/{patience})")
            if counter >= patience:
                print(f"  Early stopping at epoch {epoch}.")
                break
    return best


# ── Phase 1: Freeze CNN — teach RNN+FC to handle EMNIST chars ─
p1_best = run_phase(1, epochs=30, lr=1e-4, freeze_cnn=True, patience=7)

# ── Phase 2: Unfreeze all — gentle full-model polish ──────────
p2_best = run_phase(2, epochs=20, lr=1e-6, freeze_cnn=False, patience=5)

print(f"\n{'='*55}")
print(f"EMNIST fine-tuning complete!")
print(f"  Phase 1 best val loss : {p1_best:.4f}")
print(f"  Phase 2 best val loss : {p2_best:.4f}")
print(f"  Saved : checkpoints/best_model_emnist.pth")
print(f"\nNext step: python IAM_train.py --prepare --train")