https://github.com/artemantonov/gearbox-speed-estimation-via-vibration-analysis
Gearbox speed estimation using vibration data transformed via FFT and a lightweight PyTorch CNN
https://github.com/artemantonov/gearbox-speed-estimation-via-vibration-analysis
cnn early-stopping fft-analysis great-expectations mlflow pandas python pytorch seaborn time-series vibrational-analysis
Last synced: 27 days ago
JSON representation
Gearbox speed estimation using vibration data transformed via FFT and a lightweight PyTorch CNN
- Host: GitHub
- URL: https://github.com/artemantonov/gearbox-speed-estimation-via-vibration-analysis
- Owner: ArtemAntonov
- License: mit
- Created: 2025-04-26T19:48:07.000Z (about 1 month ago)
- Default Branch: main
- Last Pushed: 2025-05-04T18:38:24.000Z (29 days ago)
- Last Synced: 2025-05-04T19:23:56.481Z (29 days ago)
- Topics: cnn, early-stopping, fft-analysis, great-expectations, mlflow, pandas, python, pytorch, seaborn, time-series, vibrational-analysis
- Language: Jupyter Notebook
- Homepage:
- Size: 15.4 MB
- Stars: 0
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
[](https://opensource.org/licenses/MIT)
[](#)
[](#)
# Gearbox Speed Estimation via Vibration Analysis
This repository presents a machine-learning model for estimating shaft rotational speed from two-axis accelerometer vibration measurements. Raw time-domain vibration signals are transformed into the frequency domain via Fast Fourier Transform(FFT) and subsequently processed by a PyTorch convolutional neural network(CNN). The final model achieves a mean absolute error (MAE) of 1.17rps on hold-out data. By eliminating the need for dedicated tachometers, this tachometer-less approach offers cost-effective, non-intrusive speed monitoring for legacy industrial machinery.
### Objective
Develop a neural network that predicts rotational speed from raw vibration signals.
### Data
[Mechanical Gear Vibration Dataset](https://www.kaggle.com/datasets/hieudaotrung/gear-vibration/data) was used for this project. It contains 36 time series(runs), representing shaft vibration readings along two axis for six different gear faults under two loads and at three speeds.
columns description
- sensor1 - displacement along x-axis(mm)
- sensor2 - displacement along y-axis(mm)
- time_x - sampling timestamp
- speedSet - speed(revolutions per second)
- load_value - load(Nm)
- gear_fault_desc - gear fault type
### Key Features
- **EDA**: Great Expectations, autocorrelation analysis, MSTL decomposition, FFT visualizations
- **Data Preparation**: FFT extraction, sliding-window segmentation, custom oversampling
- **Model Development**: PyTorch CNN, MLflow experiment tracking, early stopping, learning rate scheduling
### Project Flow
1. [Exploratory Data Analysis](https://github.com/ArtemAntonov/Gearbox-Speed-Estimation-via-Vibration-Analysis/#1-exploratory-data-analysis)
2. [Data Preparation](https://github.com/ArtemAntonov/Gearbox-Speed-Estimation-via-Vibration-Analysis/#2-data-preparation)
3. [Model training](https://github.com/ArtemAntonov/Gearbox-Speed-Estimation-via-Vibration-Analysis/#3-model-training)
4. [Conclusion](https://github.com/ArtemAntonov/Gearbox-Speed-Estimation-via-Vibration-Analysis/#4-conclusion)
### 1. Exploratory Data Analysis
The dataset has data as per description, without missed values. Each gearbox run has 25000 samples and was performed in different time, resulting in non-overlapping readings time that can be used as index in time series. The dataset provides axial and radial shaft vibration readings.
The readings exhibit a visible trend, possibly due to the gearbox not being fully warmed up before measurements. The data shows multiple seasonal patterns with non linear behavior, which together with high noise levels suggest multiple resonting components in the system. Each gear fault produces a unique vibration signature.
Frequency-domain plots provide the clearest differentiation between runs.
### 2. Data Preparation
The dataset is relatively small compared to complexity of the system. Thus, the primary data preparation challenge was maximizing the number of usable samples.
As EDA showed, the most informative representation was the 0–50 Hz band of the FFT, which became the basis for feature extraction. The minimal number of sensor readings needed to form a recognizable spectrum was determined to be 12500. Then, a moving window method with a step of 100 samples was applied to each run to create multiple spectra.
Since the number of spectra was still insufficient for model training, synthetic augmentation was applied. Synthetic spectra were generated by shifting values in the FFT array and filling gaps with noise, significantly increasing the available training data.
### 3. Model training
The data was split into three sets:
- Train(synthetic data)
- Test(real data, excluding validation samples)
- Validation(real unseen data)
The Mean Squared Error (MSE) metric was chosen for training, being highly sensitive to large errors. Training was performed in three phases:
1. Architecture search
2. Layer configuration tuning
3. Weight decay selection and final training
All experiments and models were tracked with MLflow.
Champion model architecture:
```
nn.Conv2d(1, 33, kernel_size=(2, 15), padding=(0, 7)),
nn.ReLU(),
nn.Conv2d(33, 1, kernel_size=(1, 1)),
nn.Linear(115, 40),
nn.ReLU(),
nn.Linear(40, 1)
```
The best model achieved MSE 4.8862 and MAE 1.1665 on validation set. As expected, performance on validation set is worse than on test. 8 and 25 rps have predictions nearly as good as for test dataset. 40 rps samples have bigger errors, which can be explained by not exactly well selected run as source and thus, result suffers from unrecognised features of this spectrum.
### 4. Conclusion
In closing, this project validates that a concise PyTorch CNN, fed with thoughtfully preprocessed FFT spectra, can accurately infer gearbox rotational speed without any physical tachometer, achieving an MAE of just 1.17 rps(~3.65% of the operating range) on small amount of available data. All experiments are fully tracked in MLflow, ensuring reproducibility and transparency of model development. Along the way, we uncovered that training hyperparameters, especially batch size, have a pronounced effect on performance, underscoring the importance of systematic tuning and robust cross‑validation.