Self-Supervised Learning with Kernel Dependence Maximization
This is the code for SSL-HSIC, a self-supervised learning loss proposed
in the paper Self-Supervised Learning with Kernel Dependence Maximization
(https://arxiv.org/abs/2106.08320).
Using this implementation should achieve a top-1 accuracy on Imagenet around
74.8% using 128 Cloud TPU v2/3.
Installation
To set up a Python3 virtual environment with the required dependencies, run:
python3 -m venv ssl_hsic_env
source ssl_hsic_env/bin/activate
pip install --upgrade pip
pip install -r ssl_hsic/requirements.txt
Usage
Pre-training
For pre-training on ImageNet with SSL-HSIC loss:
mkdir /tmp/ssl_hsic
python3 -m ssl_hsic.experiment \
--config=ssl_hsic/config.py:default \
--jaxline_mode=train
This is going to pre-train for 1000 epochs. Change config to config.py:test
for testing purpose. See
jaxline documentation for more
information on jaxline_mode.
If save_dir is provided in config.py, the last checkpoint is saved and can
be used for evaluation.
Linear Evaluation
For linear evaluation with the saved checkpoint:
mkdir /tmp/ssl_hsic
python3 -m ssl_hsic.eval_experiment \
--config=ssl_hsic/eval_config.py:default \
--jaxline_mode=train
This is going to train a linear layer for 90 epochs. Change config to
eval_config.py:test for testing.
Citing this work
If you use this code in your work, please consider referencing our work:
@inproceedings{
li2021selfsupervised,
title={Self-Supervised Learning with Kernel Dependence Maximization},
author={Yazhe Li and Roman Pogodin and Danica J. Sutherland and Arthur Gretton},
booktitle={Thirty-Fifth Conference on Neural Information Processing Systems},
year={2021},
url={https://openreview.net/forum?id=0HW7A5YZjq7}
}
Disclaimer
This is not an official Google product.
