Spectral Tensor Train Parameterization of Deep Learning Layers
This repository is the official implementation of our AISTATS 2021 paper titled "Spectral Tensor Train Parameterization of Deep Learning Layers" by Anton Obukhov, Maxim Rakhuba, Alexander Liniger, Zhiwu Huang, Stamatios Georgoulis, Dengxin Dai, and Luc Van Gool [arXiv] [PMLR].
It demonstrates how to perform low-rank neural network reparameterization and its stable training in a compressed form. The code provides all experiments (GAN and Image Classification) from the paper (see configs/aistats21 directory) with the following types of reparameterizations: SNGAN, SRGAN, SVDP, or STTP.
Installation
All experiments can be reproduced on a single 11Gb GPU.
Clone the repository, then create a new virtual environment, and install python dependencies into it:
python3 -m venv venv_sttp
source venv_sttp/bin/activate
pip3 install --upgrade pip
pip3 install -r requirements.txt
In case of problems with generic requirements, fall back to requirements_reproducibility.txt.
Logging
The code performs logging to the console, tensorboard file in the experiment log directory, and also Weights and Biases (wandb). Upon the first run, please enter your wandb credentials, which can be obtained by registering a free account with the service.
Creating Environment Config
The training script allows specifying multiple yml config files, which will be concatenated during execution. This is done to separate experiment configs from environment configs. To start running experiments, create your own config file with a few environment settings, similar to configs/env_lsf.yml. Generally, you only need to update paths; see other fields explained in the config reference.
Training
Choose a preconfigured experiment from any of the configs/aistats21 directories, or compose your own config using the config reference, and run the following command:
CUDA_VISIBLE_DEVICES=0 python -m src.train --cfg configs/env_yours.yml --cfg configs/experiment.yml
Poster
