Experimental Code: Quantization-Robust Edge Predictive Maintenance

What's in this repository

src/
  cwru_loader.py              # CWRU bearing dataset loader
  model_1dcnn.py              # 1D-CNN architecture and training loop
  quantize_and_calibrate.py   # TFLite quantization + calibration metrics
  multi_seed_ablation.py      # Multi-seed orchestration
  selective_inference.py      # Risk-coverage / AURC analysis
  make_paper_figures.py       # Figure rendering (figs 1, 2, 3)
  make_tables.py              # Generates Tables I and II as plain text from JSON
  cpd_viability_check.py      # (Optional) FP32-INT8 disagreement diagnostic; not used in final paper

data/                          # CWRU .mat files (not included; download instructions below)
figures/                       # Output figures (PNG and PDF)
tflite_models_per_cell/        # Per-cell .tflite files saved during multi-seed run
multi_seed_results.json        # Aggregated results from multi_seed_ablation
selective_inference_results.json  # Aggregated results from selective_inference

Files to keep in the experimental scripts directory

Core pipeline (required):

File	Purpose	Status
cwru_loader.py	Loads CWRU .mat files, segments, splits	Required
model_1dcnn.py	1D-CNN model + training	Required
quantize_and_calibrate.py	TFLite conversion + ECE/Brier/NLL	Required
multi_seed_ablation.py	Multi-seed orchestration	Required
selective_inference.py	Risk-coverage + AURC	Required
make_paper_figures.py	Renders figures 1, 2, 3	Required
make_tables.py	Renders Tables I and II	Required

Optional/diagnostic (can be archived but not deleted):

File	Purpose	Note
cpd_viability_check.py	Cross-Precision Disagreement viability test	Used during exploration; documents the negative result that CPD-based abstention is not viable on this dataset. Keep for reproducibility audit; not needed for the final paper.

Files NOT in the repository (large artifacts):

• data/ — CWRU .mat files (40 files, ~85 MB). Download separately (see below).
• tflite_models_per_cell/ — Generated by the multi-seed run; ~30 MB total.
• *.keras model checkpoints — Optional intermediate artifacts; only generated if --save-models was passed.

Data acquisition

Download the CWRU 12 kHz drive-end fault data from:

https://engineering.case.edu/bearingdatacenter/download-data-file

Required files (40 total):

• Normal Baseline: 097, 098, 099, 100
• Inner Race fault @ 0.007": 105, 106, 107, 108
• Inner Race fault @ 0.014": 169, 170, 171, 172
• Inner Race fault @ 0.021": 209, 210, 211, 212
• Ball fault @ 0.007": 118, 119, 120, 121
• Ball fault @ 0.014": 185, 186, 187, 188
• Ball fault @ 0.021": 222, 223, 224, 225
• Outer Race @ 6 o'clock @ 0.007": 130, 131, 132, 133
• Outer Race @ 6 o'clock @ 0.014": 197, 198, 199, 200
• Outer Race @ 6 o'clock @ 0.021": 234, 235, 236, 237

Place all .mat files in ./data/. The loader expects either 097.mat or X097.mat style naming; subdirectories are walked automatically.

Software environment

• Python 3.10 or 3.11
• TensorFlow 2.15+ (tested on 2.21 on Windows, CPU)
• NumPy, scipy, matplotlib

Install:

pip install tensorflow scipy numpy matplotlib

GPU is optional; the full multi-seed pipeline runs in ~5 minutes on CPU.

Reproducing the paper results

Run these commands in order. Total time on CPU: ~10 minutes.

Step 1: Verify the data loader works

python src/cwru_loader.py --data-dir ./data --split-mode cross_load

Expected output: 40 files used, 0 missing, train/val/calib/test split sizes (approx 6998 / 1753 / 1540 / 1541).

Step 2: Run the multi-seed capacity ablation

python src/multi_seed_ablation.py \
    --data-dir ./data \
    --save-tflite-per-cell tflite_models_per_cell

This trains 9 models (3 capacities × 3 seeds), quantizes each to FP32/FP16/INT8, measures all metrics, and saves:

• multi_seed_results.json — aggregated metrics
• tflite_models_per_cell/{cap}_seed{seed}/model_{prec}.tflite — for selective inference

Step 3: Run the selective inference analysis

python src/selective_inference.py \
    --data-dir ./data \
    --tflite-base-dir tflite_models_per_cell

Reads the per-cell tflite files, computes risk-coverage curves and AURC, saves:

• selective_inference_results.json — aggregated AURC and coverage-at-risk
• figures/fig5_risk_coverage.{pdf,png} — the paper's selective inference figure

Step 4: Render figures 1, 2, 3

python src/make_paper_figures.py \
    --results-json multi_seed_results.json \
    --out-dir figures

Produces figures/fig1_ece_ratio_box.{pdf,png}, fig2_per_seed_ece.{pdf,png}, fig3_pareto.{pdf,png}. (Note: fig 4 was deliberately excluded from the final paper; if needed, the --tflite-base-dir flag enables it.)

Step 5: Render Tables I and II as plain text

python src/make_tables.py

Outputs tables/table_I_metrics.txt and tables/table_II_coverage.txt, which are the verbatim numerical tables that appear in the paper. Use these for copy-paste into the IEEE template.

Reproducibility notes

• Random seeds are set across Python, NumPy, and TensorFlow at the start of each training run. On CPU, training is bitwise deterministic for a fixed seed; on GPU, results may vary by 0.1–0.5% due to non-deterministic CuDNN kernels.
• The CWRU split is also seeded; the cross-load partition (which load is held out) is deterministic across seeds, but train/val division within the training loads varies with seed.
• TFLite INT8 quantization uses 200 random samples from the training set as the representative dataset for activation range estimation. The seed for this sample is fixed within convert_to_tflite_int8, so quantization is deterministic for a fixed model.

Known issues / limitations

• Windows users: TensorFlow >= 2.11 does not support GPU on native Windows. Use CPU (the pipeline is fast enough) or WSL2 if GPU is required.
• The tf.lite.Interpreter deprecation warning is harmless; the API works through TF 2.20+. Migration to ai_edge_litert is left as future cleanup.
• The smoke tests in each module use synthetic data and are NOT validated against the published paper's results. Run the full pipeline (Steps 1–5) for verification.