PyTorch Implementation

1import numpy as np
2
3
4def amnl_training_step(seq:        np.ndarray,
5                         rul_tgt:    np.ndarray,
6                         health_tgt: np.ndarray,
7                         params:     dict,
8                         lr:         float = 1e-3,
9                         max_rul:    float = 125.0) -> dict:
10    """One AMNL gradient-descent step in pure NumPy.
11
12    Mirrors paper_ieee_tii/grace/training/trainer.py::_train_epoch but with
13    a 1-Linear backbone instead of CNN+BiLSTM+Attention, so the algebra is
14    visible. Returns logging dict.
15    """
16    # 1. Forward
17    z            = seq.mean(axis=1) @ params["W_back"]            # (B, 256) shared
18    rul_pred     = np.maximum(0.0, z @ params["W_rul"])            # (B, 1) clamped ≥ 0
19    health_logits = z @ params["W_hs"]                              # (B, 3)
20
21    # 2. AMNL sample weights
22    weights = 1.0 + np.clip(1.0 - rul_tgt[:, 0] / max_rul, 0, 1)   # (B,)
23
24    # 3. Weighted MSE on RUL branch
25    residual = (rul_pred - rul_tgt)[:, 0]                           # (B,)
26    rul_loss = float((weights * residual ** 2).mean())
27
28    # 4. Cross-entropy on health branch (stable log-softmax)
29    z_max  = health_logits.max(-1, keepdims=True)
30    log_p  = (health_logits - z_max) - np.log(
31        np.exp(health_logits - z_max).sum(-1, keepdims=True)
32    )
33    health_loss = float(-log_p[np.arange(len(health_tgt)), health_tgt].mean())
34
35    # 5. Combine: FixedWeightLoss(0.5, 0.5)
36    total_loss = 0.5 * rul_loss + 0.5 * health_loss
37
38    # 6. Backward (analytic - the tedious part)
39    d_rul     = (2.0 / len(rul_tgt)) * weights * residual           # ∂L_rul/∂rul_pred
40    grad_rul  = z.T @ (d_rul * (rul_pred[:, 0] > 0))[:, None]        # (256, 1)
41
42    p          = np.exp(log_p)
43    onehot     = np.zeros_like(p); onehot[np.arange(len(health_tgt)), health_tgt] = 1.0
44    d_logits   = (p - onehot) / len(health_tgt)                     # (B, 3)
45    grad_hs    = z.T @ d_logits                                     # (256, 3)
46    grad_back  = (seq.mean(axis=1)).T @ (
47        0.5 * d_rul[:, None] * params["W_rul"].T +
48        0.5 * d_logits @ params["W_hs"].T
49    )                                                                # (input_dim, 256)
50
51    # 7. Adam-style step (simplified to plain SGD here for clarity)
52    params["W_back"] -= lr * grad_back
53    params["W_rul"]  -= lr * 0.5 * grad_rul
54    params["W_hs"]   -= lr * 0.5 * grad_hs
55
56    return {
57        "rul_loss":   rul_loss,
58        "health_loss": health_loss,
59        "total_loss":  total_loss,
60        "weights_min":  float(weights.min()),
61        "weights_max":  float(weights.max()),
62    }
63
64
65# ---------- Smoke test ----------
66np.random.seed(0)
67B, T, F_in = 4, 30, 14
68params = {
69    "W_back": np.random.randn(F_in, 256).astype(np.float32) * 0.05,
70    "W_rul":  np.random.randn(256, 1).astype(np.float32)    * 0.05,
71    "W_hs":   np.random.randn(256, 3).astype(np.float32)    * 0.05,
72}
73seq        = np.random.randn(B, T, F_in).astype(np.float32)
74rul_tgt    = np.array([[5.0], [40.0], [90.0], [125.0]], dtype=np.float32)
75health_tgt = np.array([2, 1, 0, 0])
76
77stats = amnl_training_step(seq, rul_tgt, health_tgt, params)
78print(f"rul_loss    : {stats['rul_loss']:.4f}")
79print(f"health_loss : {stats['health_loss']:.4f}")
80print(f"total_loss  : {stats['total_loss']:.4f}")
81print(f"weights     : [{stats['weights_min']:.2f}, {stats['weights_max']:.2f}]")

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from torch.utils.data import DataLoader
5
6# Paper-canonical pieces (paper_ieee_tii/grace/...)
7from grace.core.weighted_mse import moderate_weighted_mse_loss   # AMNL RUL loss
8from grace.core.baselines    import FixedWeightLoss              # 0.5/0.5 combiner
9from grace.models.dual_task_model import DualTaskModel            # §11.4 model
10
11
12def amnl_training_step(model:        DualTaskModel,
13                         optimizer:    torch.optim.Optimizer,
14                         mtl_loss:     FixedWeightLoss,
15                         rul_criterion: callable,
16                         hs_criterion:  callable,
17                         seq:          torch.Tensor,
18                         rul_tgt:      torch.Tensor,
19                         hs_tgt:       torch.Tensor,
20                         grad_clip:    float = 1.0) -> dict:
21    """One AMNL training step - paper-canonical recipe.
22
23    Mirrors paper_ieee_tii/grace/training/trainer.py::_train_epoch lines 249-283.
24    """
25    model.train()
26    rul_tgt = rul_tgt.view(-1, 1)                                  # match head output shape
27
28    optimizer.zero_grad()
29    rul_pred, hs_logits = model(seq)
30
31    rul_loss = rul_criterion(rul_pred, rul_tgt)
32    hs_loss  = hs_criterion(hs_logits, hs_tgt)
33    loss     = mtl_loss(rul_loss, hs_loss)                         # 0.5 · rul + 0.5 · hs
34
35    loss.backward()
36    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
37    optimizer.step()
38
39    return {
40        "loss":       loss.item(),
41        "rul_loss":   rul_loss.item(),
42        "hs_loss":    hs_loss.item(),
43        "grad_norm":  grad_norm.item(),
44    }
45
46
47# ---------- Smoke test ----------
48torch.manual_seed(0)
49model        = DualTaskModel(c_in=14, lstm_hidden=256, num_heads=8,
50                              shared_dim=32, num_classes=3)
51optimizer    = torch.optim.Adam(model.parameters(), lr=1e-3)
52mtl_loss     = FixedWeightLoss(rul_weight=0.5, health_weight=0.5)
53rul_criterion = moderate_weighted_mse_loss                          # paper-canonical
54hs_criterion  = F.cross_entropy
55
56seq     = torch.randn(4, 30, 14, requires_grad=False)
57rul_tgt = torch.tensor([5.0, 40.0, 90.0, 125.0])
58hs_tgt  = torch.tensor([2, 1, 0, 0])
59
60stats = amnl_training_step(model, optimizer, mtl_loss,
61                             rul_criterion, hs_criterion,
62                             seq, rul_tgt, hs_tgt)
63
64for k, v in stats.items():
65    print(f"{k:>12s} : {v:>10.4f}")