demo_aws_multi_gpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "source": [
    "# Privacy Meter Demo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "337837f6-ee78-4c98-ad1f-72c74aef63b8",
   "metadata": {},
   "source": [
    "## Setting up the multi-GPU environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b243137d",
   "metadata": {},
   "source": [
    "## Importing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "70647b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import math\n",
    "import time\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import yaml\n",
    "from torch.utils.data import Subset\n",
    "\n",
    "from audit import get_average_audit_results, audit_models, sample_auditing_dataset\n",
    "from get_signals import get_model_signals\n",
    "from models.utils import load_models, train_models, split_dataset_for_training\n",
    "from util import (\n",
    "    check_configs,\n",
    "    setup_log,\n",
    "    initialize_seeds,\n",
    "    create_directories,\n",
    "    load_dataset,\n",
    ")\n",
    "\n",
    "from trainers.parallel_trainer import parallel_prepare_models\n",
    "import torch.multiprocessing as mp\n",
    "if __name__ == '__main__':\n",
    "    # Required for CUDA multiprocessing\n",
    "    mp.set_start_method('spawn')\n",
    "\n",
    "# Enable benchmark mode in cudnn to improve performance when input sizes are consistent\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c129ff11",
   "metadata": {},
   "source": [
    "## Load config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9703b579",
   "metadata": {},
   "outputs": [],
   "source": [
    "configs = \"configs/config.yaml\"\n",
    "with open(configs, \"rb\") as f:\n",
    "        configs = yaml.load(f, Loader=yaml.Loader)\n",
    "\n",
    "# Validate configurations\n",
    "check_configs(configs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f957267",
   "metadata": {},
   "source": [
    "## Setting up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "70856708",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate configurations\n",
    "check_configs(configs)\n",
    "\n",
    "# Initialize seeds for reproducibility\n",
    "initialize_seeds(configs[\"run\"][\"random_seed\"])\n",
    "\n",
    "# Create necessary directories\n",
    "log_dir = configs[\"run\"][\"log_dir\"]\n",
    "directories = {\n",
    "    \"log_dir\": log_dir,\n",
    "    \"report_dir\": f\"{log_dir}/report\",\n",
    "    \"signal_dir\": f\"{log_dir}/signals\",\n",
    "    \"data_dir\": configs[\"data\"][\"data_dir\"],\n",
    "}\n",
    "create_directories(directories)\n",
    "\n",
    "# Set up logger\n",
    "logger = setup_log(\n",
    "    directories[\"report_dir\"], \"time_analysis\", configs[\"run\"][\"time_log\"]\n",
    ")\n",
    "\n",
    "start_time = time.time()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55b1e51a",
   "metadata": {},
   "source": [
    "## Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea18682d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to data/cifar10/cifar-10-python.tar.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 170498071/170498071 [00:01<00:00, 98125413.18it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data/cifar10/cifar-10-python.tar.gz to data/cifar10\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-08 16:59:02,420 INFO     Save data to data/cifar10.pkl\n",
      "2025-02-08 16:59:02,634 INFO     Save population data to data/cifar10_population.pkl\n",
      "2025-02-08 16:59:02,635 INFO     The whole dataset size: 50000\n",
      "2025-02-08 16:59:02,635 INFO     Loading dataset took 6.16908 seconds\n"
     ]
    }
   ],
   "source": [
    "baseline_time = time.time()\n",
    "dataset, population = load_dataset(configs, directories[\"data_dir\"], logger)\n",
    "logger.info(\"Loading dataset took %0.5f seconds\", time.time() - baseline_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e79b730",
   "metadata": {},
   "source": [
    "## Load or train models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "410036b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define experiment parameters\n",
    "num_experiments = configs[\"run\"][\"num_experiments\"]\n",
    "num_reference_models = configs[\"audit\"][\"num_ref_models\"]\n",
    "num_model_pairs = max(math.ceil(num_experiments / 2.0), num_reference_models + 1)\n",
    "\n",
    "# Split dataset for training\n",
    "data_splits, memberships = split_dataset_for_training(\n",
    "    len(dataset), num_model_pairs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "060fdec8-c0ca-4bbd-b5f1-9c16895c19cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-08 16:59:05,493 INFO     Training 4 models using 4 GPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU 0: 100/100 (0.0000|1.0000) | GPU 1: 100/100 (0.0000|1.0000) | GPU 2: 100/100 (0.0000|1.0000) | GPU 3: 100/100 (0.0000|1.0000) | \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ec2-user/SageMaker/privacy_meter_dev/trainers/parallel_trainer.py:257: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  saved_data = torch.load(shared_dict[idx]['model_path'])\n",
      "2025-02-08 17:08:56,723 INFO     Model parallel training took 591.2 seconds\n"
     ]
    }
   ],
   "source": [
    "# Now train models in parallel\n",
    "baseline_time = time.time()\n",
    "models_list = parallel_prepare_models(\n",
    "    log_dir, \n",
    "    dataset, \n",
    "    data_splits,  # Using the generated data_splits\n",
    "    memberships,  # Using the generated memberships\n",
    "    configs, \n",
    "    logger,\n",
    "    num_gpus=4\n",
    ")\n",
    "logger.info(\n",
    "    \"Model parallel training took %0.1f seconds\", \n",
    "    time.time() - baseline_time\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563455e3",
   "metadata": {},
   "source": [
    "## Prepare auditing dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f92be589",
   "metadata": {},
   "outputs": [],
   "source": [
    "auditing_dataset, auditing_membership = sample_auditing_dataset(\n",
    "        configs, dataset, logger, memberships\n",
    "    )\n",
    "\n",
    "# Also downsample the population set size if specified in the config\n",
    "population = Subset(\n",
    "    population,\n",
    "    np.random.choice(\n",
    "        len(population),\n",
    "        configs[\"audit\"].get(\"population_size\", len(population)),\n",
    "        replace=False,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d136fb3f",
   "metadata": {},
   "source": [
    "## Compute signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d65983a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-08 17:09:17,964 INFO     Computing signals for all models.\n",
      "Computing softmax: 100%|██████████| 10/10 [00:08<00:00,  1.12it/s]\n",
      "Computing softmax: 100%|██████████| 10/10 [00:05<00:00,  1.90it/s]\n",
      "Computing softmax: 100%|██████████| 10/10 [00:05<00:00,  1.89it/s]\n",
      "Computing softmax: 100%|██████████| 10/10 [00:05<00:00,  1.88it/s]\n",
      "2025-02-08 17:09:42,832 INFO     Signals saved to disk.\n",
      "2025-02-08 17:09:57,643 INFO     Computing signals for all models.\n",
      "Computing softmax: 100%|██████████| 2/2 [00:01<00:00,  1.89it/s]\n",
      "Computing softmax: 100%|██████████| 2/2 [00:01<00:00,  1.89it/s]\n",
      "Computing softmax: 100%|██████████| 2/2 [00:01<00:00,  1.89it/s]\n",
      "Computing softmax: 100%|██████████| 2/2 [00:01<00:00,  1.89it/s]\n",
      "2025-02-08 17:10:01,935 INFO     Signals saved to disk.\n",
      "2025-02-08 17:10:01,939 INFO     Preparing signals took 65.20681 seconds\n"
     ]
    }
   ],
   "source": [
    "baseline_time = time.time()\n",
    "signals = get_model_signals(models_list, auditing_dataset, configs, logger)\n",
    "population_signals = get_model_signals(\n",
    "        models_list, population, configs, logger, is_population=True\n",
    "    )\n",
    "logger.info(\"Preparing signals took %0.5f seconds\", time.time() - baseline_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7268e54a",
   "metadata": {},
   "source": [
    "## Audit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4b6fcf19",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-08 17:10:01,945 INFO     Fine-tuning offline_a using paired model 1\n",
      "2025-02-08 17:10:02,986 INFO     offline_a=0.00: AUC 0.7089\n",
      "2025-02-08 17:10:04,023 INFO     offline_a=0.10: AUC 0.7075\n",
      "2025-02-08 17:10:05,051 INFO     offline_a=0.20: AUC 0.7051\n",
      "2025-02-08 17:10:06,078 INFO     offline_a=0.30: AUC 0.7040\n",
      "2025-02-08 17:10:07,108 INFO     offline_a=0.40: AUC 0.7023\n",
      "2025-02-08 17:10:08,137 INFO     offline_a=0.50: AUC 0.7007\n",
      "2025-02-08 17:10:09,165 INFO     offline_a=0.60: AUC 0.6990\n",
      "2025-02-08 17:10:10,194 INFO     offline_a=0.70: AUC 0.6973\n",
      "2025-02-08 17:10:11,221 INFO     offline_a=0.80: AUC 0.6947\n",
      "2025-02-08 17:10:12,250 INFO     offline_a=0.90: AUC 0.6904\n",
      "/home/ec2-user/SageMaker/privacy_meter_dev/attacks.py:135: RuntimeWarning: overflow encountered in divide\n",
      "  ratios = prob_ratio_x[:, np.newaxis] / prob_ratio_z\n",
      "2025-02-08 17:10:13,289 INFO     offline_a=1.00: AUC 0.6604\n",
      "2025-02-08 17:10:13,290 INFO     The best offline_a is 0.0\n",
      "2025-02-08 17:10:14,319 INFO     Target Model 0: AUC 0.6954, TPR@0.1%FPR of 0.0000, TPR@0.0%FPR of 0.0000\n",
      "2025-02-08 17:10:18,576 INFO     Auditing the privacy risks of target model 0 costs 16.6 seconds\n",
      "2025-02-08 17:10:18,577 INFO     Total runtime: 685.10334 seconds\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 640x480 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Perform the privacy audit\n",
    "baseline_time = time.time()\n",
    "target_model_indices = list(range(num_experiments))\n",
    "mia_score_list, membership_list = audit_models(\n",
    "        f\"{directories['report_dir']}/exp\",\n",
    "        target_model_indices,\n",
    "        signals,\n",
    "        population_signals,\n",
    "        auditing_membership,\n",
    "        num_reference_models,\n",
    "        logger,\n",
    "        configs,\n",
    "    )\n",
    "\n",
    "if len(target_model_indices) > 1:\n",
    "    logger.info(\n",
    "        \"Auditing privacy risk took %0.1f seconds\", time.time() - baseline_time\n",
    "    )\n",
    "\n",
    "# Get average audit results across all experiments\n",
    "if len(target_model_indices) > 1:\n",
    "    get_average_audit_results(\n",
    "        directories[\"report_dir\"], mia_score_list, membership_list, logger\n",
    "    )\n",
    "\n",
    "logger.info(\"Total runtime: %0.5f seconds\", time.time() - start_time)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_privacymeter_p310",
   "language": "python",
   "name": "conda_privacymeter_p310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}