Improving Model Robustness to Weight Noise via Consistency Regularization - PyTorch Implementation

Official PyTorch implementation of the paper "Improving Model Robustness to Weight Noise via Consistency Regularization" published in Machine Learning: Science and Technology.

Features

• Implementation of "Improving Model Robustness to Weight Noise via Consistency Regularization"
• Support for multi-GPU training.

Requirements

• Python 3.8+
• PyTorch 2.0+
• Torchvision
• Tensorboard
• Scikit-learn
• Matplotlib

Dataset Preparation

1. Download CIFAR10 dataset from CIFAR-10 and CIFAR-100 datasets
2. Extract dataset to data/cifar-10 directory

Quick Start

Training Example

bash

python main_QAT_classifier.py config.yaml

Project Structure

├── model/                 # Model definitions
├── data/                  # Dataset storage
├── pretrained_models/     # Pretrained models
├── quantization_and_noise/ # Quantization and weight noise tools
├── train_utils/           # Training/testing code
├── util/                  # General utilities
├── config.yaml            # Configuration file
├── main_QAT_classifier.py # Quantization-aware training classifier script
└── README.md

License

This project is licensed under the MIT License - see the LICENSE file for details.

Citation

If you use this code in your research, please cite our paper:

bibtex

@article{Hou_2024,
doi = {10.1088/2632-2153/ad734a},
url = {https://dx.doi.org/10.1088/2632-2153/ad734a},
year = {2024},
month = {sep},
publisher = {IOP Publishing},
volume = {5},
number = {3},
pages = {035065},
author = {Hou, Yaoqi and Zhang, Qingtian and Wang, Namin and Wu, Huaqiang},
title = {Improving model robustness to weight noise via consistency regularization},
journal = {Machine Learning: Science and Technology},
}

Contact

For questions and suggestions, please contact:

• Houyaoqi - houyaoqi@mail.tsinghua.edu.cn

References

This implementation references GitHub - zhutmost/lsq-net: Unofficial implementation of LSQ-Net, a neural network quantization framework.