File size: 3,818 Bytes
2adb41a | 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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """
ml/scripts/train_fault_classifier.py
Trains the Bi-LSTM Geometry Fault Classifier (tc.v1 SOTA).
"""
import os
import sys
import argparse
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from ml.data.synthetic_geometry import SyntheticGeometryDataset
from ml.models.geometry.fault_classifier import BiLSTMAttention
def train(
num_samples: int = 5000,
val_samples: int = 1000,
epochs: int = 20,
batch_size: int = 64,
lr: float = 1e-3,
save_path: str = "artifacts/checkpoints/geometry/bilstm_fault_typing.pt",
):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 1. Data
train_ds = SyntheticGeometryDataset(num_samples=num_samples)
val_ds = SyntheticGeometryDataset(num_samples=val_samples)
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=batch_size)
# 2. Model
model = BiLSTMAttention().to(device)
# 3. SOTA Loss & Optimizer
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
# 4. Training Loop
print(f"Starting training for {epochs} epochs on {len(train_ds)} samples...")
for epoch in range(epochs):
model.train()
train_loss = 0.0
for X, y in train_loader:
X, y = X.to(device), y.to(device)
optimizer.zero_grad()
logits, _ = model(X)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
train_loss += loss.item()
scheduler.step()
# Validation
model.eval()
correct = 0
total = 0
with torch.no_grad():
for X, y in val_loader:
X, y = X.to(device), y.to(device)
logits, _ = model(X)
preds = torch.argmax(logits, dim=1)
correct += (preds == y).sum().item()
total += y.size(0)
val_acc = correct / max(1, total)
print(f"Epoch {epoch+1:02d}/{epochs:02d} | Loss: {train_loss/len(train_loader):.4f} | Val Acc: {val_acc * 100:.2f}%")
# 5. Save Checkpoint
out_dir = os.path.dirname(save_path)
os.makedirs(out_dir, exist_ok=True)
torch.save(model.state_dict(), save_path)
print(f"[OK] Model saved to {save_path}")
# Mirror to ml/models/geometry/weights for unified resolution
mirror_path = "ml/models/geometry/weights/fault_classifier.pt"
os.makedirs(os.path.dirname(mirror_path), exist_ok=True)
torch.save({"model_state_dict": model.state_dict(), "temperature": 1.5}, mirror_path)
print(f"[OK] Mirrored weights to {mirror_path}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Train Bi-LSTM Geometry Classifier.")
parser.add_argument("--epochs", type=int, default=20, help="Training epochs")
parser.add_argument("--batch_size", "--batch-size", "--batch", type=int, default=64, dest="batch_size", help="Batch size")
parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
parser.add_argument("--save_path", "--save-path", default="artifacts/checkpoints/geometry/bilstm_fault_typing.pt", dest="save_path")
parser.add_argument("--num_samples", "--num-samples", type=int, default=5000, dest="num_samples")
args = parser.parse_args()
train(
epochs=args.epochs,
batch_size=args.batch_size,
lr=args.lr,
save_path=args.save_path,
num_samples=args.num_samples,
)
|