This repository contains the original PyTorch implementation of the paper 'Self-Supervised Representation Learning for Wafer Bin Map Defect Pattern Classification'.
- Anaconda (>= 4.8.3)
- OpenCV (tested on 4.3.0.66)
- PyTorch (tested on 1.6.0)
- PyTorch Lightning (tested on 0.8.5)
- Albumentations (tested on 0.4.6)
- Tensorboard (tested on 2.2.2, optional)
conda update -n base conda # use 4.8.3 or higher
conda create -n wbm python=3.6
conda activate wbm
conda install anaconda
conda install opencv -c conda-forge
conda install pytorch=1.6.0 cudatoolkit=10.2 -c pytorch
pip install pytorch_lightning
pip install albumentations
- Download from the following link: WM-811K
- Place the
LSWMD.pkl
file under./data/wm811k/
. - Run the following script from the working directory:
python process_wm811k.py
python run_wapirl.py \
--input_size 96 \
--augmentation crop \
--backbone_type resnet \
--backbone_config 18 \
--decouple_input \
--epochs 100 \
--batch_size 256 \
--num_workers 4 \
--gpus 0 \
--optimizer sgd \
--learning_rate 1e-2 \
--weight_decay 1e-3 \
--momentum 0.9 \
--scheduler cosine \
--warmup_steps 0 \
--checkpoint_root ./checkpoints \
--write_summary \
--save_every 10 \
--projector_type linear \
--projector_size 128 \
--temperature 0.07
- Run
python run_wapirl.py --help
for more information on arguments. - If running on a Windows machine, set
num_workers
to 0. (multiprocessing does not function well.)
python run_wapirl.py @experiments/pretrain_wapirl.txt
python run_classification.py \
--input_size 96 \
--augmentation crop \
--backbone_type resnet \
--backbone_config 18 \
--decouple_input \
--epochs 100 \
--batch_size 256 \
--num_workers 4 \
--gpus 0 \
--optimizer sgd \
--learning_rate 1e-2 \
--weight_decay 1e-3 \
--momentum 0.9 \
--scheduler cosine \
--warmup_steps 0 \
--checkpoint_root ./checkpoints \
--write_summary \
--pretrained_model_file /path/to/file \
--pretrained_model_type wapirl \
--label_proportion 1.00 \
--label_smoothing 0.1 \
--dropout 0.5
- IMPORTANT: Provide the correct path for the
pretrained_model_file
argument. - Run
python run_classification.py --help
for more information on arguments. - If running on a Windows machine, set
num_workers
to 0. (multiprocessing does not function well.)
python run_classification.py @experiments/finetune_wapirl.txt