Full Training Script Walkthrough

1# Pseudo-pipeline. Real Python uses PyTorch (next block); this version
2# wires the abstract pieces together so the SHAPES and CONTROL FLOW are
3# visible without DL framework specifics.
4
5import numpy as np
6
7
8def train_one_dataset(dataset_name: str = "FD002",
9                       epochs:        int = 200,
10                       lr:            float = 1e-3,
11                       weight_decay:  float = 1e-4,
12                       grad_clip:     float = 1.0,
13                       ema_decay:     float = 0.999) -> dict:
14    """End-to-end AMNL training for one C-MAPSS subset.
15
16    Returns history dict with per-epoch loss / RMSE / NASA score.
17    """
18    # 1-2. Data: build dataset and loader (paper §7)
19    train_ds = build_dataset(dataset_name, split="train")
20    test_ds  = build_dataset(dataset_name, split="test")
21    train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
22    test_loader  = DataLoader(test_ds,  batch_size=32, shuffle=False)
23
24    # 3. Model: §11.4 DualTaskModel
25    model = DualTaskModel(c_in=14, lstm_hidden=256, num_heads=8,
26                           shared_dim=32, num_classes=3)
27
28    # 4-5. Loss stack: §14.1 + §15.1
29    rul_criterion = moderate_weighted_mse_loss
30    hs_criterion  = cross_entropy
31    mtl_loss      = FixedWeightLoss(rul_weight=0.5, health_weight=0.5)
32
33    # 6. Optimiser + scheduler: §15.2
34    optimizer = AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
35    scheduler = ReduceLROnPlateau(optimizer, factor=0.5, patience=30, min_lr=5e-6)
36
37    # 7. EMA + grad clip: §15.3
38    ema = ExponentialMovingAverage(model, decay=ema_decay)
39
40    # ---------- Per-epoch loop ----------
41    history = {"rmse": [], "nasa": [], "rul_loss": [], "hs_loss": [], "lr": []}
42    best_rmse = float("inf")
43    for epoch in range(epochs):
44        # 8. Warmup (§15.2)
45        for pg in optimizer.param_groups:
46            pg["lr"] = lr * (0.1 + 0.9 * min(epoch, 10) / 10)
47
48        # 9. Training step (§14.4)
49        train_stats = train_epoch(model, train_loader,
50                                    rul_criterion, hs_criterion, mtl_loss,
51                                    optimizer, ema,
52                                    grad_clip=grad_clip)
53
54        # 10. Eval with EMA shadow (§15.3 + §13)
55        ema.apply_shadow(model)
56        eval_stats = evaluate(model, test_loader)
57        ema.restore(model)
58
59        # 11. Scheduler step (§15.2)
60        if epoch >= 10:
61            scheduler.step(eval_stats["rmse"])
62
63        # 12. Track best
64        if eval_stats["rmse"] < best_rmse:
65            best_rmse = eval_stats["rmse"]
66
67        history["rmse"]      .append(eval_stats["rmse"])
68        history["nasa"]      .append(eval_stats["nasa_score"])
69        history["rul_loss"]  .append(train_stats["rul_loss"])
70        history["hs_loss"]   .append(train_stats["hs_loss"])
71        history["lr"]        .append(optimizer.param_groups[0]["lr"])
72
73    return {"history": history, "best_rmse": best_rmse}
74
75
76# ---------- Smoke test ----------
77np.random.seed(0)
78print("training pipeline (pseudo)")
79print("  - dataset       : FD002")
80print("  - epochs        : 200")
81print("  - paper RMSE    : 13.36 (FD002 + FD004 average)")
82print("  - paper NASA    : 1302  (FD002 + FD004 average)")
83# In real code: out = train_one_dataset()

1import torch
2import torch.nn.functional as F
3import torch.optim as optim
4from torch.utils.data import DataLoader
5
6# Paper-canonical pieces (paper_ieee_tii/grace/...)
7from grace.data.cmapss_dataset      import CMAPSSFullDataset       # §7
8from grace.models.dual_task_model    import DualTaskModel            # §11.4
9from grace.core.weighted_mse         import moderate_weighted_mse_loss # §14
10from grace.core.baselines            import FixedWeightLoss           # §15.1
11from grace.training.callbacks        import ExponentialMovingAverage   # §15.3
12from grace.training.evaluator        import evaluate_model             # §13
13
14
15def train_amnl(dataset_name: str = "FD002",
16                epochs:        int = 200,
17                lr:            float = 1e-3) -> dict:
18    """Reproduce the paper's AMNL training run on one C-MAPSS subset.
19
20    Source: paper_ieee_tii/experiments/train_amnl_v7.py (top-level driver).
21    """
22    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23
24    # 1-2. Data
25    train_ds = CMAPSSFullDataset(dataset_name, split="train")
26    test_ds  = CMAPSSFullDataset(dataset_name, split="test")
27    train_loader = DataLoader(train_ds, batch_size=32, shuffle=True,  num_workers=4, pin_memory=True)
28    test_loader  = DataLoader(test_ds,  batch_size=32, shuffle=False, num_workers=2, pin_memory=True)
29
30    # 3. Model
31    model = DualTaskModel(c_in=14, lstm_hidden=256, num_heads=8,
32                           shared_dim=32, num_classes=3).to(device)
33
34    # 4-5. Loss stack
35    rul_criterion = moderate_weighted_mse_loss
36    hs_criterion  = F.cross_entropy
37    mtl_loss      = FixedWeightLoss(rul_weight=0.5, health_weight=0.5)
38
39    # 6. Optim + sched
40    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4,
41                              betas=(0.9, 0.999), eps=1e-8)
42    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min",
43                                                        factor=0.5, patience=30,
44                                                        min_lr=5e-6)
45
46    # 7. EMA
47    ema = ExponentialMovingAverage(model, decay=0.999)
48
49    # ---------- Training loop ----------
50    history = {"rmse": [], "nasa": [], "lr": []}
51    best_rmse = float("inf")
52    for epoch in range(epochs):
53        # 8. Warmup
54        if epoch < 10:
55            for pg in optimizer.param_groups:
56                pg["lr"] = lr * (0.1 + 0.9 * epoch / 10)
57
58        # 9. Train
59        model.train()
60        for seq, rul_tgt, hs_tgt, _ in train_loader:
61            seq, rul_tgt, hs_tgt = seq.to(device), rul_tgt.to(device).view(-1, 1), hs_tgt.to(device)
62            optimizer.zero_grad()
63            rul_pred, hs_logits = model(seq)
64            rul_loss = rul_criterion(rul_pred, rul_tgt)
65            hs_loss  = hs_criterion(hs_logits, hs_tgt)
66            loss     = mtl_loss(rul_loss, hs_loss)
67            loss.backward()
68            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
69            optimizer.step()
70            ema.update(model)
71
72        # 10. Evaluate with EMA
73        ema.apply_shadow(model)
74        eval_stats = evaluate_model(model, test_loader, device)
75        ema.restore(model)
76
77        # 11. Scheduler step (after warmup)
78        if epoch >= 10:
79            scheduler.step(eval_stats["rmse"])
80
81        # 12. Track best
82        if eval_stats["rmse"] < best_rmse:
83            best_rmse = eval_stats["rmse"]
84
85        history["rmse"].append(eval_stats["rmse"])
86        history["nasa"].append(eval_stats["nasa_score"])
87        history["lr"]  .append(optimizer.param_groups[0]["lr"])
88
89        if epoch % 20 == 0:
90            print(f"epoch={epoch:>3d}  rmse={eval_stats['rmse']:.3f}  nasa={eval_stats['nasa_score']:.1f}  lr={optimizer.param_groups[0]['lr']:.2e}")
91
92    return {"history": history, "best_rmse": best_rmse}
93
94
95if __name__ == "__main__":
96    out = train_amnl("FD002", epochs=200)
97    print(f"FD002 best RMSE : {out['best_rmse']:.3f}")