# Isoloss

## Installing the environment

All the simulations use 
[Jax](https://github.com/google/jax), 
[Flax](https://github.com/google/flax), 
[TensorFlow](https://github.com/tensorflow/tensorflow), and 
[TensorFlow Datasets](https://www.tensorflow.org/datasets?hl=en).
To install the environment with the supported versions of these packages, run the following commands (requires python>=3.8): 

```
python3 -m venv isoloss
source isoloss/bin/activate
pip install --upgrade pip setuptools
./jax_installer.sh
pip install -r requirements.txt
```

To run the notebooks, you will also need to install JupyterLab and Matplotlib.
```
pip install jupyterlab
pip install matplotlib
```

## Sample run scripts

The following example scripts can be adapted to run the experiments in the paper.

### BYOL
```
python main.py --nolog --dataset cifar100 --cosine_lr_decay --distance_metric cosine --proj mlp --pred mlp --loss simsiam --ema
```

### DirectPred
```
python main.py --nolog --dataset cifar100 --cosine_lr_decay --distance_metric cosine --proj mlp --pred dp --loss directloss --ema
```

### IsoLoss
```
python main.py --nolog --dataset cifar100 --cosine_lr_decay --distance_metric cosine --proj mlp --pred dp --loss directloss --ema --iso
```

The flag `--nolog` disables logging to Weights & Biases.
Once you have a Weights & Biases account configured and added to `main.py`, you can remove this flag to log the results of your experiments.