{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Memory comparison: ffocp_eq vs cvxpylayer\n",
"\n",
"This notebook compares the peak RSS (Resident Set Size) memory usage between `ffocp_eq` and `cvxpylayer` across problem sizes.\n",
"\n",
"- batch_size = 8\n",
"- ydim ∈ {100, 200, …, 1000}\n",
"- Only **one iteration** per ydim (first batch forward + backward)\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import sys\n",
"import time\n",
"import gc\n",
"\n",
"import numpy as np\n",
"import torch\n",
"import pandas as pd\n",
"import matplotlib.pyplot as plt\n",
"\n",
"# Add synthetic_task directory directly to path (no __init__.py in that folder)\n",
"sys.path.insert(0, os.path.join(os.path.abspath('.'), 'synthetic_task'))\n",
"\n",
"from models import OptModel\n",
"from data import genData\n",
"from utils_synthetic import PeakRSS, bytes_to_gb"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Using device: cpu\n",
"ydim values to benchmark: [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000]\n"
]
}
],
"source": [
"# Configuration\n",
"BATCH_SIZE = 8\n",
"SEED = 3\n",
"INPUT_DIM = 640\n",
"YDIM_LIST = list(range(100, 1001, 100))\n",
"NUM_SAMPLES = 2000\n",
"LEARNING_RATE = 0.001\n",
"\n",
"# Hyperparameters for ffocp_eq\n",
"ALPHA = 100\n",
"DUAL_CUTOFF = 1e-3\n",
"SLACK_TOL = 1e-8\n",
"BACKWARD_EPS = 1e-3\n",
"\n",
"# Device\n",
"device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')\n",
"print(f\"Using device: {device}\")\n",
"print(f\"ydim values to benchmark: {YDIM_LIST}\")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"def set_seed(seed):\n",
" \"\"\"Set random seeds for reproducibility\"\"\"\n",
" np.random.seed(seed)\n",
" torch.manual_seed(seed)\n",
" torch.cuda.manual_seed_all(seed)\n",
" torch.backends.cudnn.deterministic = True\n",
"\n",
"def cleanup():\n",
" \"\"\"Clean up memory\"\"\"\n",
" gc.collect()\n",
" if torch.cuda.is_available():\n",
" torch.cuda.empty_cache()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"def benchmark_method(method_name, ydim):\n",
" \"\"\"\n",
" Benchmark a single method for initialization and first iteration peak RSS.\n",
" \n",
" Args:\n",
" method_name: 'ffocp_eq' or 'cvxpylayer'\n",
" ydim: dimension of the optimization variable\n",
" \n",
" Returns:\n",
" dict: Contains init_peak_rss, forward_peak_rss, backward_peak_rss (all in GB)\n",
" \"\"\"\n",
" cleanup()\n",
" set_seed(SEED)\n",
" \n",
" results = {'ydim': ydim, 'method': method_name}\n",
" \n",
" # Generate data\n",
" train_loader, test_loader = genData(device, INPUT_DIM, ydim, NUM_SAMPLES, BATCH_SIZE)\n",
" \n",
" # Measure initialization peak RSS\n",
" with PeakRSS(include_children=True, interval=0.0001) as m_init:\n",
" model = OptModel(\n",
" INPUT_DIM, ydim, \n",
" layer_type=method_name, \n",
" constraint_learnable=False, \n",
" batch_size=BATCH_SIZE, \n",
" device=device, \n",
" alpha=ALPHA, \n",
" dual_cutoff=DUAL_CUTOFF, \n",
" slack_tol=SLACK_TOL, \n",
" backward_eps=BACKWARD_EPS, \n",
" is_QP=True\n",
" ).to(device)\n",
" \n",
" results['init_peak_rss_gb'] = bytes_to_gb(m_init.peak)\n",
" \n",
" # Setup optimizer and loss\n",
" optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=0)\n",
" loss_fn = torch.nn.MSELoss()\n",
" \n",
" # Get first batch\n",
" x, y = next(iter(train_loader))\n",
" \n",
" model.train()\n",
" \n",
" # Measure forward pass peak RSS\n",
" with PeakRSS(include_children=True, interval=0.0001) as m_fwd:\n",
" start_fwd = time.time()\n",
" z, y_pred = model(x)\n",
" forward_time = time.time() - start_fwd\n",
" \n",
" results['forward_peak_rss_gb'] = bytes_to_gb(m_fwd.peak)\n",
" results['forward_time'] = forward_time\n",
" \n",
" # Compute loss\n",
" ts_loss = loss_fn(y_pred, y)\n",
" df_loss = torch.mean(y * z)\n",
" loss = df_loss\n",
" \n",
" # Measure backward pass peak RSS\n",
" with PeakRSS(include_children=True, interval=0.0001) as m_bwd:\n",
" start_bwd = time.time()\n",
" loss.backward()\n",
" backward_time = time.time() - start_bwd\n",
" \n",
" results['backward_peak_rss_gb'] = bytes_to_gb(m_bwd.peak)\n",
" results['backward_time'] = backward_time\n",
" results['total_iter_time'] = forward_time + backward_time\n",
" \n",
" # Cleanup\n",
" del model, optimizer, train_loader, test_loader\n",
" cleanup()\n",
" \n",
" return results"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Run Benchmarks for All ydim Values"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"============================================================\n",
"Benchmarking ydim=100\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 0.3490 GB, Forward RSS: 0.3977 GB, Backward RSS: 0.4037 GB\n",
" Running cvxpylayer...\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"/nethome/zzhao628/anaconda3/envs/bilevel/lib/python3.10/site-packages/cvxpy/reductions/solvers/solving_chain.py:250: UserWarning: Your problem has too many parameters for efficient DPP compilation. We suggest setting 'ignore_dpp = True'.\n",
" warnings.warn(\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
" Init RSS: 0.5497 GB, Forward RSS: 0.5671 GB, Backward RSS: 0.5808 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=200\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 0.5896 GB, Forward RSS: 0.5896 GB, Backward RSS: 0.5945 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 1.2633 GB, Forward RSS: 1.2686 GB, Backward RSS: 1.3053 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=300\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 1.3261 GB, Forward RSS: 1.3312 GB, Backward RSS: 1.3318 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 2.6546 GB, Forward RSS: 2.6542 GB, Backward RSS: 2.7367 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=400\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 2.7676 GB, Forward RSS: 2.7676 GB, Backward RSS: 2.7688 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 4.3196 GB, Forward RSS: 4.3207 GB, Backward RSS: 4.5112 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=500\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 4.5102 GB, Forward RSS: 4.5102 GB, Backward RSS: 4.5129 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 7.0341 GB, Forward RSS: 7.2490 GB, Backward RSS: 7.3287 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=600\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 7.1074 GB, Forward RSS: 7.2074 GB, Backward RSS: 7.2079 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 12.0712 GB, Forward RSS: 7.4286 GB, Backward RSS: 7.5332 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=700\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 7.2419 GB, Forward RSS: 7.2593 GB, Backward RSS: 7.2898 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 23.8247 GB, Forward RSS: 8.7073 GB, Backward RSS: 8.8440 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=800\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 8.3415 GB, Forward RSS: 8.4895 GB, Backward RSS: 8.4902 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 31.4607 GB, Forward RSS: 8.7852 GB, Backward RSS: 8.8766 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=900\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 8.3628 GB, Forward RSS: 8.5696 GB, Backward RSS: 8.5752 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 55.3312 GB, Forward RSS: 55.7627 GB, Backward RSS: 55.9967 GB\n",
"\n",
"============================================================\n",
"Benchmarking ydim=1000\n",
"============================================================\n",
" Running ffocp_eq...\n",
"FFOLayer forward eps = 1e-12, backward eps = 0.001\n",
"max_workers: 8\n",
" Init RSS: 55.5935 GB, Forward RSS: 55.5935 GB, Backward RSS: 55.5935 GB\n",
" Running cvxpylayer...\n",
" Init RSS: 88.7127 GB, Forward RSS: 88.7127 GB, Backward RSS: 88.9277 GB\n",
"\n",
"============================================================\n",
"All benchmarks completed!\n",
"============================================================\n"
]
}
],
"source": [
"# Run benchmarks for all ydim values\n",
"all_results = []\n",
"\n",
"for ydim in YDIM_LIST:\n",
" print(f\"\\n{'='*60}\")\n",
" print(f\"Benchmarking ydim={ydim}\")\n",
" print(f\"{'='*60}\")\n",
" \n",
" # Benchmark ffocp_eq\n",
" print(f\" Running ffocp_eq...\")\n",
" ffocp_result = benchmark_method('ffocp_eq', ydim)\n",
" all_results.append(ffocp_result)\n",
" print(f\" Init RSS: {ffocp_result['init_peak_rss_gb']:.4f} GB, \"\n",
" f\"Forward RSS: {ffocp_result['forward_peak_rss_gb']:.4f} GB, \"\n",
" f\"Backward RSS: {ffocp_result['backward_peak_rss_gb']:.4f} GB\")\n",
" \n",
" # Benchmark cvxpylayer\n",
" print(f\" Running cvxpylayer...\")\n",
" cvxpy_result = benchmark_method('cvxpylayer', ydim)\n",
" all_results.append(cvxpy_result)\n",
" print(f\" Init RSS: {cvxpy_result['init_peak_rss_gb']:.4f} GB, \"\n",
" f\"Forward RSS: {cvxpy_result['forward_peak_rss_gb']:.4f} GB, \"\n",
" f\"Backward RSS: {cvxpy_result['backward_peak_rss_gb']:.4f} GB\")\n",
"\n",
"print(f\"\\n{'='*60}\")\n",
"print(\"All benchmarks completed!\")\n",
"print(f\"{'='*60}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Results Table"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Full Results:\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" ydim | \n",
" method | \n",
" init_peak_rss_gb | \n",
" forward_peak_rss_gb | \n",
" forward_time | \n",
" backward_peak_rss_gb | \n",
" backward_time | \n",
" total_iter_time | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" 100 | \n",
" ffocp_eq | \n",
" 0.562920 | \n",
" 0.652424 | \n",
" 0.337005 | \n",
" 0.675358 | \n",
" 0.204346 | \n",
" 0.541352 | \n",
"
\n",
" \n",
" | 1 | \n",
" 100 | \n",
" cvxpylayer | \n",
" 0.704193 | \n",
" 0.778027 | \n",
" 0.166445 | \n",
" 0.781296 | \n",
" 0.028233 | \n",
" 0.194678 | \n",
"
\n",
" \n",
" | 2 | \n",
" 200 | \n",
" ffocp_eq | \n",
" 0.564537 | \n",
" 0.672375 | \n",
" 0.209352 | \n",
" 0.690907 | \n",
" 0.184464 | \n",
" 0.393817 | \n",
"
\n",
" \n",
" | 3 | \n",
" 200 | \n",
" cvxpylayer | \n",
" 0.966274 | \n",
" 0.737301 | \n",
" 0.263529 | \n",
" 0.763065 | \n",
" 0.049828 | \n",
" 0.313357 | \n",
"
\n",
" \n",
" | 4 | \n",
" 300 | \n",
" ffocp_eq | \n",
" 0.564751 | \n",
" 0.677315 | \n",
" 0.255209 | \n",
" 0.732018 | \n",
" 0.236742 | \n",
" 0.491951 | \n",
"
\n",
" \n",
" | 5 | \n",
" 300 | \n",
" cvxpylayer | \n",
" 1.952305 | \n",
" 0.906479 | \n",
" 0.686642 | \n",
" 0.950520 | \n",
" 0.103317 | \n",
" 0.789960 | \n",
"
\n",
" \n",
" | 6 | \n",
" 400 | \n",
" ffocp_eq | \n",
" 0.566673 | \n",
" 0.728924 | \n",
" 0.361483 | \n",
" 0.787350 | \n",
" 0.331312 | \n",
" 0.692795 | \n",
"
\n",
" \n",
" | 7 | \n",
" 400 | \n",
" cvxpylayer | \n",
" 4.026390 | \n",
" 1.370781 | \n",
" 1.072643 | \n",
" 1.432758 | \n",
" 0.182669 | \n",
" 1.255312 | \n",
"
\n",
" \n",
" | 8 | \n",
" 500 | \n",
" ffocp_eq | \n",
" 0.569153 | \n",
" 0.814888 | \n",
" 0.530983 | \n",
" 0.860821 | \n",
" 0.509320 | \n",
" 1.040303 | \n",
"
\n",
" \n",
" | 9 | \n",
" 500 | \n",
" cvxpylayer | \n",
" 7.714897 | \n",
" 2.415863 | \n",
" 1.835934 | \n",
" 2.506645 | \n",
" 0.265122 | \n",
" 2.101056 | \n",
"
\n",
" \n",
" | 10 | \n",
" 600 | \n",
" ffocp_eq | \n",
" 0.563145 | \n",
" 0.834091 | \n",
" 0.363253 | \n",
" 0.902866 | \n",
" 0.349034 | \n",
" 0.712287 | \n",
"
\n",
" \n",
" | 11 | \n",
" 600 | \n",
" cvxpylayer | \n",
" 337.181324 | \n",
" 337.578888 | \n",
" 2.881727 | \n",
" 337.705185 | \n",
" 0.426281 | \n",
" 3.308007 | \n",
"
\n",
" \n",
" | 12 | \n",
" 700 | \n",
" ffocp_eq | \n",
" 0.568317 | \n",
" 0.877663 | \n",
" 0.450494 | \n",
" 1.015194 | \n",
" 0.438758 | \n",
" 0.889252 | \n",
"
\n",
" \n",
" | 13 | \n",
" 700 | \n",
" cvxpylayer | \n",
" 383.725735 | \n",
" 384.244263 | \n",
" 3.595613 | \n",
" 384.360676 | \n",
" 0.609306 | \n",
" 4.204918 | \n",
"
\n",
" \n",
" | 14 | \n",
" 800 | \n",
" ffocp_eq | \n",
" 0.571453 | \n",
" 0.927097 | \n",
" 0.561200 | \n",
" 1.126850 | \n",
" 0.571702 | \n",
" 1.132902 | \n",
"
\n",
" \n",
" | 15 | \n",
" 800 | \n",
" cvxpylayer | \n",
" 406.781963 | \n",
" 407.450722 | \n",
" 5.034194 | \n",
" 407.637691 | \n",
" 0.810200 | \n",
" 5.844394 | \n",
"
\n",
" \n",
" | 16 | \n",
" 900 | \n",
" ffocp_eq | \n",
" 0.574322 | \n",
" 0.971909 | \n",
" 0.762432 | \n",
" 1.202812 | \n",
" 0.772691 | \n",
" 1.535123 | \n",
"
\n",
" \n",
" | 17 | \n",
" 900 | \n",
" cvxpylayer | \n",
" 436.398518 | \n",
" 437.217331 | \n",
" 6.556314 | \n",
" 437.362904 | \n",
" 1.157879 | \n",
" 7.714193 | \n",
"
\n",
" \n",
" | 18 | \n",
" 1000 | \n",
" ffocp_eq | \n",
" 0.570454 | \n",
" 1.040497 | \n",
" 1.023682 | \n",
" 1.312859 | \n",
" 1.247311 | \n",
" 2.270993 | \n",
"
\n",
" \n",
" | 19 | \n",
" 1000 | \n",
" cvxpylayer | \n",
" 487.272373 | \n",
" 488.176178 | \n",
" 8.727450 | \n",
" 488.280502 | \n",
" 1.387577 | \n",
" 10.115027 | \n",
"
\n",
" \n",
"
\n",
"