| import argparse, glob, os, time |
| import numpy as np, torch, torch.nn as nn, torch.nn.functional as F |
| import scipy.io as sio |
| from modeling import FELAPDM, PDMConfig |
|
|
| cwru_classes = [ |
| "Normal", |
| "IR007", |
| "B007", |
| "OR007", |
| "IR014", |
| "B014", |
| "OR014", |
| "IR021", |
| "B021", |
| "OR021", |
| ] |
| cwru_idx = {c: i for i, c in enumerate(cwru_classes)} |
| sensor_cols = [6, 7, 8, 11, 12, 13, 15, 16, 17, 18, 19, 21, 24, 25] |
| rul_cap = 125 |
|
|
|
|
| def cwru_label(f): |
| b = os.path.basename(f).split("_")[0] |
| if b in cwru_idx: |
| return cwru_idx[b] |
| raise ValueError(f) |
|
|
|
|
| def load_cwru(d, win=2048, stride=1024, loads=("0", "1", "2", "3"), seed=0): |
| rng = np.random.default_rng(seed) |
| xs, ys = ([], []) |
| for f in sorted(glob.glob(os.path.join(d, "*.mat"))): |
| load = os.path.basename(f).split("_")[1].replace(".mat", "") |
| if load not in loads: |
| continue |
| m = sio.loadmat(f) |
| de = [k for k in m if k.endswith("DE_time")] |
| if not de: |
| continue |
| sig = m[de[0]].ravel().astype(np.float32) |
| sig = (sig - sig.mean()) / (sig.std() + 1e-08) |
| lab = cwru_label(f) |
| for s in range(0, len(sig) - win + 1, stride): |
| xs.append(sig[s : s + win]) |
| ys.append(lab) |
| x = np.stack(xs).astype(np.float32)[..., None] |
| y = np.array(ys, np.int64) |
| idx = rng.permutation(len(x)) |
| return (x[idx], y[idx]) |
|
|
|
|
| def cwru_split(x, y, test_frac=0.25, seed=0): |
| rng = np.random.default_rng(seed) |
| n = len(x) |
| perm = rng.permutation(n) |
| nte = int(n * test_frac) |
| te, tr = (perm[:nte], perm[nte:]) |
| return (x[tr], y[tr], x[te], y[te]) |
|
|
|
|
| def read_cmapss(p): |
| return np.loadtxt(p) |
|
|
|
|
| def load_cmapss(d, subset="FD001", win=30, cap=rul_cap): |
| train = read_cmapss(os.path.join(d, f"train_{subset}.txt")) |
| test = read_cmapss(os.path.join(d, f"test_{subset}.txt")) |
| truth = read_cmapss(os.path.join(d, f"RUL_{subset}.txt")).ravel() |
| s = sensor_cols |
| tr_s = train[:, s] |
| smin, smax = (tr_s.min(0), tr_s.max(0)) |
| rng = smax - smin |
| rng[rng == 0] = 1.0 |
|
|
| def norm(a): |
| return (a[:, s] - smin) / rng |
|
|
| xtr, ytr = ([], []) |
| for u in np.unique(train[:, 0]): |
| eng = train[train[:, 0] == u] |
| feats = norm(eng).astype(np.float32) |
| L = len(feats) |
| mc = eng[:, 1].max() |
| for e in range(win, L + 1): |
| xtr.append(feats[e - win : e]) |
| ytr.append(min(mc - eng[e - 1, 1], cap)) |
| if L < win: |
| pad = np.zeros((win - L, len(s)), np.float32) |
| xtr.append(np.concatenate([pad, feats], 0)) |
| ytr.append(min(mc - eng[-1, 1], cap)) |
| xtr = np.stack(xtr).astype(np.float32) |
| ytr = np.array(ytr, np.float32) |
| xte, yte = ([], []) |
| for i, u in enumerate(np.unique(test[:, 0])): |
| eng = test[test[:, 0] == u] |
| feats = norm(eng).astype(np.float32) |
| L = len(feats) |
| if L >= win: |
| w = feats[L - win :] |
| else: |
| pad = np.zeros((win - L, len(s)), np.float32) |
| w = np.concatenate([pad, feats], 0) |
| xte.append(w) |
| yte.append(min(truth[i], cap)) |
| return (xtr, ytr, np.stack(xte).astype(np.float32), np.array(yte, np.float32)) |
|
|
|
|
| def cmapss_score(yt, yp): |
| d = yp - yt |
| s = np.where(d < 0, np.exp(-d / 13.0) - 1.0, np.exp(d / 10.0) - 1.0) |
| return (float(np.sqrt(np.mean(d**2))), float(np.sum(s))) |
|
|
|
|
| def batches(x, y, bs, seed=0): |
| idx = np.arange(len(x)) |
| np.random.default_rng(seed).shuffle(idx) |
| for i in range(0, len(x), bs): |
| j = idx[i : i + bs] |
| yield (x[j], y[j]) |
|
|
|
|
| def cfg_cwru(**kw): |
| c = PDMConfig( |
| in_channels=1, |
| patch=4, |
| n_embd=64, |
| n_layer=4, |
| n_head=4, |
| fno_modes=64, |
| ffn_hidden=128, |
| n_classes=10, |
| seq_len=2048, |
| ) |
| for k, v in kw.items(): |
| setattr(c, k, v) |
| return c |
|
|
|
|
| def cfg_cmapss(**kw): |
| c = PDMConfig( |
| in_channels=14, |
| patch=1, |
| n_embd=64, |
| n_layer=4, |
| n_head=4, |
| fno_modes=32, |
| ffn_hidden=128, |
| rul_head=True, |
| seq_len=64, |
| ) |
| for k, v in kw.items(): |
| setattr(c, k, v) |
| return c |
|
|
|
|
| def train_cwru(a): |
| dev = torch.device(a.device) |
| x, y = load_cwru(a.data, win=2048, stride=a.stride, seed=a.seed) |
| xtr, ytr, xte, yte = cwru_split(x, y, test_frac=0.25, seed=a.seed) |
| assert len(x) == 5886 |
| if a.smoke: |
| print( |
| f"Cwru total {len(x)} train {len(xtr)} test {len(xte)} classes {len(set(y.tolist()))}" |
| ) |
| return |
| cfg = cfg_cwru(dropout=a.dropout) |
| m = FELAPDM(cfg).to(dev) |
| opt = torch.optim.AdamW(m.parameters(), lr=a.lr, weight_decay=0.0001) |
| sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, a.epochs) |
| xte_t = torch.from_numpy(xte).to(dev) |
| yte_t = torch.from_numpy(yte).to(dev) |
| best = 0.0 |
| for ep in range(a.epochs): |
| m.train() |
| for xb, yb in batches(xtr, ytr, a.bs, seed=a.seed + ep): |
| xb = torch.from_numpy(xb).to(dev) |
| yb = torch.from_numpy(yb).to(dev) |
| loss = F.cross_entropy(m(xb, task="cls"), yb) |
| opt.zero_grad() |
| loss.backward() |
| nn.utils.clip_grad_norm_(m.parameters(), 1.0) |
| opt.step() |
| sch.step() |
| m.eval() |
| with torch.no_grad(): |
| pr = torch.cat( |
| [ |
| m(xte_t[i : i + 256], task="cls").argmax(-1) |
| for i in range(0, len(xte_t), 256) |
| ] |
| ) |
| acc = (pr == yte_t).float().mean().item() |
| if acc > best: |
| best = acc |
| torch.save( |
| {"cfg": cfg.__dict__, "model": m.state_dict(), "classes": cwru_classes}, |
| a.out, |
| ) |
| print(f"Ep {ep:02d} acc {acc * 100:.2f} best {best * 100:.2f}") |
|
|
|
|
| def train_cmapss(a): |
| dev = torch.device(a.device) |
| xtr, ytr, xte, yte = load_cmapss(a.data, subset=a.subset, win=a.win) |
| if a.subset == "FD001": |
| assert len(xtr) == 17731 |
| if a.smoke: |
| print(f"Cmapss {a.subset} train {xtr.shape} test {xte.shape}") |
| return |
| cfg = cfg_cmapss(dropout=a.dropout, seq_len=a.win) |
| m = FELAPDM(cfg).to(dev) |
| opt = torch.optim.AdamW(m.parameters(), lr=a.lr, weight_decay=0.0001) |
| sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, a.epochs) |
| cap = 125.0 |
| xte_t = torch.from_numpy(xte).to(dev) |
| best = 1000000000.0 |
| for ep in range(a.epochs): |
| m.train() |
| for xb, yb in batches(xtr, ytr / cap, a.bs, seed=a.seed + ep): |
| xb = torch.from_numpy(xb).to(dev) |
| yb = torch.from_numpy(yb).to(dev) |
| loss = F.mse_loss(m(xb, task="rul"), yb) |
| opt.zero_grad() |
| loss.backward() |
| nn.utils.clip_grad_norm_(m.parameters(), 1.0) |
| opt.step() |
| sch.step() |
| m.eval() |
| with torch.no_grad(): |
| pr = np.clip( |
| torch.cat( |
| [ |
| m(xte_t[i : i + 256], task="rul") |
| for i in range(0, len(xte_t), 256) |
| ] |
| ) |
| .cpu() |
| .numpy() |
| * cap, |
| 0, |
| cap, |
| ) |
| rmse, score = cmapss_score(yte, pr) |
| if rmse < best: |
| best = rmse |
| torch.save( |
| {"cfg": cfg.__dict__, "model": m.state_dict(), "rul_cap": cap}, a.out |
| ) |
| print(f"Ep {ep:02d} RMSE {rmse:.2f} score {score:.0f} best {best:.2f}") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--task", choices=["cwru", "cmapss"], required=True) |
| ap.add_argument("--data", default=None) |
| ap.add_argument("--subset", default="FD001") |
| ap.add_argument("--out", default=None) |
| ap.add_argument("--epochs", type=int, default=40) |
| ap.add_argument("--bs", type=int, default=64) |
| ap.add_argument("--lr", type=float, default=0.002) |
| ap.add_argument("--win", type=int, default=30) |
| ap.add_argument("--stride", type=int, default=1024) |
| ap.add_argument("--dropout", type=float, default=0.0) |
| ap.add_argument("--device", default="cpu") |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--smoke", action="store_true") |
| a = ap.parse_args() |
| if a.task == "cwru": |
| a.data = a.data or "../data/cwru" |
| a.out = a.out or "./pdm_cwru_fno.pt" |
| train_cwru(a) |
| else: |
| a.data = a.data or "../data/cmapss" |
| a.out = a.out or f"./pdm_cmapss_{a.subset}_fno.pt" |
| train_cmapss(a) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
