File size: 2,352 Bytes
031d0d7 | 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 | import argparse
import os
import sys
import torch
sys.path.insert(0, os.path.dirname(__file__))
from modeling import load_model
SAMPLES = {
"cmapss_FD001": torch.full((1, 30, 14), 0.5),
"cmapss_FD002": torch.full((1, 30, 14), 0.5),
"cmapss_FD003": torch.full((1, 30, 14), 0.5),
"cmapss_FD004": torch.full((1, 30, 14), 0.5),
"cwru": torch.linspace(-1, 1, 2048).reshape(1, 2048, 1),
}
VERIFICATION = {
"cmapss_FD001": {"value": 0.390998, "tol": 0.001},
"cmapss_FD002": {"value": 0.038257, "tol": 0.001},
"cmapss_FD003": {"value": 0.788864, "tol": 0.001},
"cmapss_FD004": {"value": 0.414077, "tol": 0.001},
"cwru": {"value": 0, "tol": 0.001},
}
EXPECTED_SHAPE = {
"cmapss_FD001": (1,),
"cmapss_FD002": (1,),
"cmapss_FD003": (1,),
"cmapss_FD004": (1,),
"cwru": (1, 10),
}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--variant", default="cmapss_FD001", choices=list(SAMPLES))
ap.add_argument("--weights", default=os.environ.get("FELA_PDM_WEIGHTS", "."))
args = ap.parse_args()
model = load_model(args.weights, variant=args.variant)
x = SAMPLES[args.variant]
with torch.no_grad():
out = model(x)
exp = EXPECTED_SHAPE[args.variant]
if tuple(out.shape) != exp:
print(f"Fail: output shape {tuple(out.shape)} != expected {exp}")
sys.exit(1)
print(f"Shape OK: {tuple(out.shape)}")
g = VERIFICATION[args.variant]
if g["value"] is None:
if args.variant.startswith("cmapss"):
captured = float(out.reshape(-1)[0])
else:
captured = int(out.argmax(-1).item())
print(f"Captured output: {captured}")
print(
"Verification value is a placeholder. Paste this captured value into VERIFICATION and re-run to enable the check. Shape check passed."
)
return
if args.variant.startswith("cmapss"):
got = float(out.reshape(-1)[0])
if abs(got - g["value"]) > g["tol"]:
print(f"Fail: RUL {got} vs verification {g['value']} (tol {g['tol']})")
sys.exit(1)
else:
got = int(out.argmax(-1).item())
if got != g["value"]:
print(f"Fail: class {got} vs verification {g['value']}")
sys.exit(1)
print("Verification check OK")
if __name__ == "__main__":
main()
|