ramdrop / stun

Implementation for the paper: STUN: Self-Teaching Uncertainty Estimation for Place Recognition
BSD 3-Clause "New" or "Revised" License
30 stars 3 forks source link

[IROS2022] STUN: Self-Teaching Uncertainty Estimation for Place Recognition

0. Environment Setup ⚙️

1. Download the Pittsburgh dataset and pretrained models 📨

bash ./download.sh

the files will be downloaded and saved in the following folders:

pittsburgh
├── database
├── query
└── structure

logs
├── student_contrast
├── student_quadruplet
├── student_triplet
└── teacher_triplet

2. Evaluate the pretrained models 🔍

   # STUN
   python main.py  --resume=logs/student_triplet/ckpt.pth.tar

   # STUN (Constrast)
   python main.py --resume=logs/student_constrast/ckpt.pth.tar

   # STUN (Quadruplet)
   python main.py --resume=logs/student_quadruplet/ckpt.pth.tar

   # Standard Triplet
   python main.py --phase=test_tea   --resume=logs/teacher_triplet/ckpt.pth.tar

3. Plot results 📈

python vis_results.py
# you can plot results of different models by populate the NETWORK variable.

4. Train and evaluate STUN from scratch 🧭

   # train the teacher net
   python main.py --phase=train_tea --loss=tri

   # train the student net supervised by the pretrained teacher net
   python main.py --phase=train_stu --resume=[teacher_net_xxx/ckpt_best.pth.tar]

After analyzing empirical figures, we found the correlation between recall@N and uncertainty level evolve into a sensible trend after 30 epochs. But ECE (Expected Calibration Error) will diverge if the student network is excessively trained. As a result, we focused our examination on the model's performance from epoch=30 to epoch=35 and chose the one with the lowest ECE.

   # evaluate 
   ./eval_batch.sh

ece.png

If you find our work useful, please consider citing:

@INPROCEEDINGS{stun_cai,
  author={Cai, Kaiwen and Lu, Chris Xiaoxuan and Huang, Xiaowei},
  booktitle={2022 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, 
  title={STUN: Self-Teaching Uncertainty Estimation for Place Recognition}, 
  year={2022},
  volume={},
  number={},
  pages={6614-6621},
  doi={10.1109/IROS47612.2022.9981546}}