Few-shot-via-ensembling-Transformer-with-Mahalanobis-distance

(CODE WILL RELEASE SOON!!)

We provide a Pytorch implement code of paper "Few-Shot Bearing Fault Diagnosis via Ensembling Transformer-based Model with Mahalanobis Distance Metric Learning from Multiscale Features" accepted by IEEE Transactions on Instrumentation and Measurement.

Prerequisites

• Python 3
• Linux
• Pytorch 0.4+
• GPU + CUDA CuDNN

Dataset

In this paper, we ultilize 2 datasets: CWRU and PU.

Note, if you use these datasets, please cite the corresponding papers. (Feel free to contact me if you need PU dataset in .pt file)

Getting Started

• Installation

git clone https://github.com/HungVu307/Few-shot-via-ensembling-Transformer-with-Mahalanobis-distance

• Training for 1 shot

python train_1shot.py --dataset 'CWRU' --training_samples_CWRU 30 --training_samples_PDB 195 --model_name 'Net'

• Testing for 1 shot

python test_1shot.py --dataset 'CWRU' --best_weight 'PATH TO BEST WEIGHT'

• Training for 5 shot

python train_5shot.py --dataset 'CWRU' --training_samples_CWRU 60 --training_samples_PDB 300 --model_name 'Net'

• Testing for 5 shot

python test_5shot.py --dataset 'CWRU' --best_weight 'PATH TO BEST WEIGHT'

• Result

1. CWRU dataset
2. PU dataset

Contact

Please feel free to contact me via email [email protected] or [email protected] if you need anything related to this repo!

Citation

If you feel this code is useful, please give us 1 ⭐ and cite our paper.