Official PyTorch implementation of the ICLR 2024 paper Non-negative Contrastive Learning by Yifei Wang*, Qi Zhang*, Yaoyu Guo, and Yisen Wang.
Links: Wandb Logs | Checkpoints
Updates:
- 2024.04.14. Update codes for calculating class consistency, sparsity, and image retrieval precision in main_eval.py.
- 2024.04.01. Add logging of feature stats & wandb pretraining logs. With these stats, we can directly observe the differences between CL and NCL training.
- 2024.03.19. Code is released. 🎉
With non-negative constraints on contrastive features, NCL can significantly enhance feature interpretability, sparsity, and disentanglement over standard contrastive learning, while improving (at least maintaining) its performance on classical tasks.
For implementation, the difference between Contrastive Learning (CL) and Non-negative Contrastive learning (NCL) is minimal. Apart from argparser nuances, NCL only adds one line upon existing methods:
z = torch.nn.functional.relu(z)
which applies a non-negative transformation (eg ReLU) to the output features z
to enforce feature non-negativity. This is all we need to convert a standard CL method (e.g., SimCLR) to a non-negative version, and deliver all the magic.
The codebase is built upon a previous version of solo-learn
(the version on Sep 27, 2022). To avoid unexpected errors, first create a Python3.8
environment, and then install the reposoity as below.
# clone the repository
git clone https://github.com/PKU-ML/non_neg
# create environment
conda create -n non_neg python=3.8
conda activate non_neg
# install dependences
cd non_neg
pip3 install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist --extra-index-url https://download.pytorch.org/whl/cu113
Pretrain with the default configuration files using the following command.
# SimCLR
python3 main_pretrain.py \
--config-path scripts/pretrain/cifar \
--config-name simclr.yaml
# NCL
python3 main_pretrain.py \
--config-path scripts/pretrain/cifar \
--config-name ncl.yaml
The default setting is for CIFAR-100. For CIFAR-10, override data.dataset=cifar10
. Meanwhile, change the experiment name accordingly to avoid collision, e.g., name=simclr-resnet18-cifar10-ncl
. Other experiments follow the same setting.
# SimCLR
python3 main_pretrain.py \
--config-path scripts/pretrain/imagenet-100 \
--config-name simclr.yaml
# NCL
python3 main_pretrain.py \
--config-path scripts/pretrain/imagenet-100 \
--config-name ncl.yaml
By default, we use non_neg=rep_relu
for CIFAR-10 and CIFAR-100 and non_neg=relu
for ImageNet-100.
After that, for linear evaluation, run the following command:
python3 main_linear.py \
--config-path scripts/linear/{dataset} \
--config-name simclr.yaml \
pretrained_feature_extractor=path/to/pretrained/feature/extractor
Here dataset={cifar,imagenet100}
. We use the argument pretrained_feature_extractor
to configure the path of the pretrained checkpoints.
And for fine-tuning evaluation, run the following command:
python3 main_linear.py \
--config-path scripts/finetuning/{dataset} \
--config-name simclr.yaml
And for offline linear probing with selected dimensions, run the following command:
python3 main_linear.py \
--config-path scripts/selected \
--config-name simclr.yaml \
selected_dims=512
where the argument selected_dims
configures the dimensions of selected features.
For offline evaluating the class consistency, the sparsity and the mAP for image retrieval, run the following command:
python3 main_eval.py \
--config-path scripts/eval/{dataset} \
--config-name simclr.yaml \
resume_from_checkpoint=path/to/pretrained/feature/extractor
The following table provides the pre-trained checkpoints for CL and NCL.
CIFAR-10 | CIFAR-100 | ImageNet-100 | |
---|---|---|---|
checkpoints | download | download | download |
If you find the work useful, please cite the accompanying paper:
@inproceedings{
wang2024nonnegative,
title={Non-negative Contrastive Learning},
author={Yifei Wang and Qi Zhang and Yaoyu Guo and Yisen Wang},
booktitle={ICLR},
year={2024},
}
Our codes borrow the implementations of SimCLR in the solo-learn repository: https://github.com/vturrisi/solo-learn