Changes coming soon!
Code for the paper:
Fast, Scalable, Warm-Start Semidefinite Programming with Spectral Bundling and Sketching
Rico Angell and Andrew McCallum
arXiv:2312.11801
ICML 2024
USBS is an optimization algorithm for solving large semidefinite programs.
This repo was developed with Python 3.12.4. To get started we recommend creating a fresh virutal environment using virtualenv, conda, or mamba. The main package that needs to be installed is JAX (this repo was tested with version 0.4.31.dev20240701). Follow the instructions here to install the version of JAX needed for the desired hardware (CPU/GPU/TPU). Although it is likely unnecessary for most users, we download and build JAX from source using the instructions here.
We build JAX from source in order to run USBS on a single NVIDIA A100 GPU which uses different versions of CUDA and CUDNN than the pre-built versions of JAX available here. We see a massive performance improvement from using a GPU, and thus, were inclined to build JAX from source. In the case that you have access to a GPU and want to get the most out of USBS, we provide some tips for building JAX from source for GPU. Use these tips to augment the official build instructions (read the full build instructions from the official JAX website first!). As a note, we built JAX on a Linux machine running Ubuntu 20.04.6 LTS. You should be able to ignore the following tips if there exists a pre-built version of JAX that suites your available hardware.
- Before building, we add the file paths to both CUDA and CUDNN to the enviornment variable
LD_LIBRARY_PATH
using the following commands:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/cuda/
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/cudnn/
-
We found success building
jaxlib
using theclang
compiler instead ofgcc
. It seems that this has something to due with builing XLA (see this). Assumingclang
is already installed on your system you can use it to buildjaxlib
usingclang
by adding the--use_clang
flag to the end of thepython build/build.py ...
command. -
We found that after
jax
andjaxlib
are built and installed into the virtual environment we needed to install thebitsandbytes
package using the following command:
pip install -U bitsandbytes
After JAX is installed, the remaining packages can be installed using the following command:
pip install numpy scipy scikit-learn numba GitPython IPython mat73
All of the data can be downloaded here. The max-cut data was aggregated from Gset and DIMACS10. The QAP data was aggregated from QAPLIB and TSPLIB.
After downloading the data and moving it to $PROJECT_ROOT
, the following command will extract the data into the expected location:
tar xzvf data.tgz
The following are example commands to execute for each of the three problem types presented in the paper.
PYTHONPATH="." python -u scripts/warm_start_maxcut_usbs.py --data_path=data/maxcut/Gset/G1.mat --max_iters=5000 --max_time=360 --trace_factor=2.0 --rho=0.01 --beta=0.25 --k_curr=10 --k_past=1 --sketch_dim=10 --obj_gap_eps=1e-07 --infeas_gap_eps=1e-07 --max_infeas_eps=1e-07 --subprob_max_iters=100 --subprob_eps=1e-15 --lanczos_max_restarts=10 --warm_start_strategy="none"
PYTHONPATH="." python -u scripts/warm_start_qap_usbs.py --data_path=data/qap/qapdata/chr12a.dat --max_iters=5000 --max_time=360 --trace_factor=2.0 --rho=0.005 --beta=0.25 --k_curr=2 --k_past=0 --obj_gap_eps=1e-07 --infeas_gap_eps=1e-07 --max_infeas_eps=1e-07 --subprob_max_iters=100 --subprob_eps=1e-7 --lanczos_max_restarts=10 --warm_start_strategy="none"
PYTHONPATH="." python -u scripts/warm_start_ecc.py --seed=0 --debug --output_dir=test_out --data_path=data/ecc/merged_pubmed_processed.pkl --max_rounds=100 --max_iters=100000 --trace_factor=2.0 --k_curr=3 --k_past=0 --rho=0.01 --beta=0.25 --sketch_dim=-1 --subprob_max_iters=30 --subprob_eps=1e-7 --lanczos_max_restarts=10 --obj_gap_eps=0.1 --infeas_gap_eps=0.1 --max_infeas_eps=0.1
If you use USBS in your work, please cite the following paper:
@misc{angell2023fast,
title={Fast, Scalable, Warm-Start Semidefinite Programming with Spectral Bundling and Sketching},
author={Rico Angell and Andrew McCallum},
year={2023},
eprint={2312.11801},
archivePrefix={arXiv},
primaryClass={math.OC}
}
If you have any questions, comments, or feedback on our work, please reach out at rangell@cs.umass.edu! (or open a GitHub issue)
USBS is MIT licensed. See the LICENSE file for details.
- Add fixed frequency cond evaluation
- Add the best default params
- Add better timing and diagnostic information
- README
- GPU non-determinism note: google/jax#10674
- Explanation of how to specify problem data
- Explanation of parameters
- Maxcut walkthrough example
- Useful snippets contained in repo:
- While loop primitives
- Thick Restart Lanczos
- Munkres’ algorithm
- Fix citation to official ICML
- Convert problem from ineq (prob) -> eq [prob -> soln] -> ineq (soln)
- SDPLR(+)