run_dp_ntk_regression_10classes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/nvme2/xuhaiyang/softwares/anaconda3/envs/galore/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/mnt/nvme2/xuhaiyang/softwares/anaconda3/envs/galore/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet18_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet18_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50000/50000 [00:01<00:00, 37756.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total train idx len 10000\n",
      "total test idx len 1000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [03:07<00:00, 53.35it/s]\n",
      "100%|██████████| 1000/1000 [00:19<00:00, 52.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hdis shape torch.Size([10000, 10000])\n",
      "alpha shape torch.Size([10000, 10])\n",
      "tensor(0.8662) tensor(0.8380)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "from torch.autograd import Variable\n",
    "import torchvision\n",
    "from tqdm import tqdm\n",
    "import math\n",
    "from ntk_utils import gen_h_dis, gen_alpha, gen_z_embed, process_query\n",
    "from truncated_laplace_utils import add_truncated_laplace_noise, add_laplace_noise\n",
    "\n",
    "# we use CPU as our device. \n",
    "# GPU doesnt have enough memory\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "# the following code is adapted from\n",
    "# https://github.com/josharnoldjosh/Resnet-Extract-Image-Feature-Pytorch-Python \n",
    "\n",
    "# Load the pretrained resnet18 model\n",
    "model = models.resnet18(pretrained=True)\n",
    "model = model.to(device)\n",
    "\n",
    "# Use the model object to select the desired layer\n",
    "layer = model._modules.get('avgpool')\n",
    "\n",
    "# Set model to evaluation mode\n",
    "model.eval()\n",
    "\n",
    "# Image transforms\n",
    "scaler = transforms.Resize((224, 224))\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "to_tensor = transforms.ToTensor()\n",
    "\n",
    "def get_vector(img):\n",
    "    # 2. Create a PyTorch Variable with the transformed image\n",
    "    t_img = Variable(normalize(to_tensor(scaler(img))).unsqueeze(0))\n",
    "    # 3. Create a vector of zeros that will hold our feature vector\n",
    "    #    The 'avgpool' layer has an output size of 512\n",
    "    my_embedding = torch.zeros(512).to(device)\n",
    "    # 4. Define a function that will copy the output of a layer\n",
    "    def copy_data(m, i, o):\n",
    "        # my_embedding.copy_(o.data)\n",
    "        my_embedding.copy_(o.data.reshape(o.data.size(1)))\n",
    "    # 5. Attach that function to our selected layer\n",
    "    h = layer.register_forward_hook(copy_data)\n",
    "    # 6. Run the model on our transformed image\n",
    "    t_img = t_img.to(device)\n",
    "    model(t_img)\n",
    "    # 7. Detach our copy function from the layer\n",
    "    h.remove()\n",
    "\n",
    "    return my_embedding\n",
    "\n",
    "# download cifar10 dataset\n",
    "ds = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)\n",
    "\n",
    "# mapping from class name to index\n",
    "class_to_idx_dict = ds.class_to_idx\n",
    "# mapping from index to class name\n",
    "idx_to_cls_dict = {}\n",
    "\n",
    "# collect all class name\n",
    "class_name_list = []\n",
    "# collect all class index\n",
    "class_idx_list = []\n",
    "for key, value in list(class_to_idx_dict.items()):\n",
    "    class_name_list.append(key)\n",
    "    class_idx_list.append(value)\n",
    "    idx_to_cls_dict[value] = key\n",
    "\n",
    "# collect data according to its class\n",
    "ds_train_idx_by_cls_list = [] \n",
    "for cls_idx in class_idx_list:\n",
    "    ds_train_idx_by_cls_list.append([])\n",
    "\n",
    "# start collecting\n",
    "for i in tqdm(range(len(ds))):\n",
    "    cur_cls_idx = ds[i][1]\n",
    "    ds_train_idx_by_cls_list[cur_cls_idx].append(i)\n",
    "\n",
    "\n",
    "# %%\n",
    "######### normalization x data start ###########\n",
    "def calculate_norm(input_data):\n",
    "    # input_data: n * d\n",
    "    square_data = input_data * input_data\n",
    "    norm_data = square_data.sum(dim=1)\n",
    "    norm_data = torch.sqrt(norm_data)\n",
    "    return norm_data\n",
    "\n",
    "def data_normalization(input_data):\n",
    "    # input data: n * d\n",
    "    # print(\"data normalized\")\n",
    "    x_norm = calculate_norm(input_data)\n",
    "    x_norm = x_norm[..., None]\n",
    "    ball_data = input_data / x_norm\n",
    "\n",
    "    return ball_data\n",
    "######### normalization x data end ###########\n",
    "\n",
    "\n",
    "train_num = 1000\n",
    "test_num = 100\n",
    "label_num = 10\n",
    "\n",
    "total_train_idx_list = []\n",
    "total_test_idx_list = []\n",
    "for train_idx_list in ds_train_idx_by_cls_list:\n",
    "    total_train_idx_list += train_idx_list[:train_num]\n",
    "    total_test_idx_list += train_idx_list[-test_num:]\n",
    "\n",
    "print(\"total train idx len\", len(total_train_idx_list))\n",
    "print(\"total test idx len\", len(total_test_idx_list))\n",
    "\n",
    "############ get img and label tensor according to idx start ###############\n",
    "def get_img_label_tensor(idx_list):\n",
    "    img_ts_list = []\n",
    "    label_list = []\n",
    "    for idx in tqdm(idx_list):\n",
    "        image, label = ds[idx]\n",
    "        img_ts = get_vector(image)\n",
    "        img_ts_list.append(img_ts)\n",
    "        label_list.append(label)\n",
    "\n",
    "    # concat image tensor\n",
    "    img_ts = torch.stack(img_ts_list, dim=0)\n",
    "\n",
    "    label_ts = torch.zeros((len(idx_list), label_num), dtype=torch.float32)\n",
    "    # set the negative label to -1\n",
    "    label_ts -= 1.0\n",
    "    for i, label in enumerate(label_list):\n",
    "        # set corresponding label to 1\n",
    "        label_ts[i][label] = 1.0\n",
    "\n",
    "    cls_index_label_ts = torch.tensor(label_list, dtype=torch.int64)\n",
    "\n",
    "    return img_ts, label_ts, cls_index_label_ts\n",
    "############ get img and label tensor according to idx end ###############\n",
    "\n",
    "# img_ts: n * 512\n",
    "# label_ts: n * 10\n",
    "# cls_index_label_ts: n * 1\n",
    "train_img_ts, train_label_ts, train_cls_index_label_ts = get_img_label_tensor(total_train_idx_list)\n",
    "test_img_ts, test_label_ts, test_cls_index_label_ts = get_img_label_tensor(total_test_idx_list)\n",
    "\n",
    "train_img_ts = data_normalization(train_img_ts)\n",
    "test_img_ts = data_normalization(test_img_ts)\n",
    "\n",
    "m = 256\n",
    "reg_lambda = 10.0\n",
    "\n",
    "x_data = train_img_ts.to(device)\n",
    "y_data = train_label_ts.to(device)\n",
    "n, d = x_data.shape\n",
    "\n",
    "# generate w_r\n",
    "w_r = torch.randn((m, d), dtype=torch.float32).to(device)\n",
    "# h_dis: n * n\n",
    "h_dis = gen_h_dis(w_r, x_data)\n",
    "\n",
    "print(\"hdis shape\", h_dis.shape)\n",
    "\n",
    "# calculate NTK Regression alpha\n",
    "alpha = gen_alpha(h_dis, reg_lambda, y_data)\n",
    "\n",
    "print(\"alpha shape\", alpha.shape)\n",
    "    \n",
    "\n",
    "################# test ntk regression accuracy start #####################\n",
    "def process_10_cls_query(z, w_r, x_data, alpha):\n",
    "    # z denote the query, nz denote the query num\n",
    "    # z: nz * d\n",
    "    # w_r: m * d\n",
    "    # x_data: n * d\n",
    "    # alpha: n * 10\n",
    "    # return: pred: nz * 1\n",
    "\n",
    "    # nz * n\n",
    "    query_embed = gen_z_embed(z, x_data, w_r)\n",
    "\n",
    "    # nz * 10\n",
    "    query_pred = query_embed @ alpha\n",
    "\n",
    "    query_result = torch.argmax(query_pred, dim=1)\n",
    "\n",
    "    # nz * 1\n",
    "    return query_result\n",
    "\n",
    "\n",
    "def test_accuracy_for_10_cls(test_dataset, gt_label, w_r, x_data, alpha):\n",
    "    # test_dataset: n * 512\n",
    "    # gt_label: n * 1\n",
    "    \n",
    "    # pred: nz * 1\n",
    "    pred = process_10_cls_query(test_dataset, w_r, x_data, alpha)\n",
    "    nz = pred.shape[0]\n",
    "    succ_cnt = torch.sum(pred == gt_label)\n",
    "    test_acc = succ_cnt / nz\n",
    "    return test_acc\n",
    "################# test ntk regression accuracy end #####################\n",
    "\n",
    "\n",
    "unprivate_train_acc = test_accuracy_for_10_cls(train_img_ts, train_cls_index_label_ts, w_r, train_img_ts, alpha)\n",
    "unprivate_test_acc = test_accuracy_for_10_cls(test_img_ts, test_cls_index_label_ts, w_r, train_img_ts, alpha)\n",
    "\n",
    "# print accuracy for unprivate version \n",
    "print(unprivate_train_acc, unprivate_test_acc)\n",
    "\n",
    "# create the log folder\n",
    "os.makedirs(\"logs\", exist_ok=True)\n",
    "\n",
    "# write the unprivate acc to log\n",
    "with open(os.path.join(\"logs\", \"unprivate_log.txt\"), \"w\") as fw:\n",
    "    fw.write(f\"unprivate_train_acc {unprivate_train_acc}\\n\")\n",
    "    fw.write(f\"unprivate_test_acc {unprivate_test_acc}\\n\")\n",
    "\n",
    "\n",
    "# function for calculating k according to eps, delta and beta\n",
    "def cal_k(eps, delta, beta):\n",
    "    eta = 7e-3\n",
    "    n = 1e3\n",
    "    k_bound = (eps * eps * eta * eta) / (8 * math.log(1 / delta) * n * n * beta * beta)\n",
    "    k = int(math.floor(k_bound))\n",
    "    return k\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n",
      "eps exponent 1.4, k 559\n",
      "truncated laplace noise added\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 1/5 [01:16<05:04, 76.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8298) tensor(0.7900)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 2/5 [02:35<03:53, 78.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8232) tensor(0.7920)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 3/5 [03:54<02:36, 78.30s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8130) tensor(0.7840)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 4/5 [05:15<01:19, 79.34s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8081) tensor(0.7768)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [06:33<00:00, 78.67s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8119) tensor(0.7816)\n",
      "test acc tensor(0.8119)\n",
      "train acc tensor(0.7816)\n",
      "--------------------------------------------------\n",
      "--------------------------------------------------\n",
      "eps exponent 1.5, k 886\n",
      "truncated laplace noise added\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 1/5 [01:40<06:43, 101.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8405) tensor(0.8050)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 2/5 [03:22<05:04, 101.54s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8313) tensor(0.8025)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 3/5 [05:07<03:26, 103.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8298) tensor(0.8043)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 4/5 [06:56<01:45, 105.41s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8294) tensor(0.8030)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [08:41<00:00, 104.38s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.8249) tensor(0.7934)\n",
      "test acc tensor(0.8249)\n",
      "train acc tensor(0.7934)\n",
      "--------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# we will repeat adding noise for several times,\n",
    "# to reduce the variance of our results\n",
    "private_repeat_time = 5\n",
    "\n",
    "\n",
    "def gaussain_sampling_on_k(h_dis, y_data, reg_lambda, wt_train_img_ts=None, k=None):\n",
    "    \n",
    "    # make sure we have the private wtX \n",
    "    assert wt_train_img_ts is not None\n",
    "\n",
    "    test_acc_list = []\n",
    "    train_acc_list = []\n",
    "\n",
    "    for _ in tqdm(range(private_repeat_time)):\n",
    "        if k is None:\n",
    "            wt_h_dis = h_dis\n",
    "        else:\n",
    "            # setup gaussian sampler\n",
    "            gaussian_sampler = torch.distributions.MultivariateNormal(\n",
    "                loc=torch.zeros(n).to(device), covariance_matrix=h_dis\n",
    "            )\n",
    "\n",
    "            wt_h_dis = torch.zeros(n, n).to(device)\n",
    "\n",
    "            # Gaussian sample k times, then we take the average to get the pirvate version wtHdis\n",
    "            for i in range(k):\n",
    "                sample_vec = gaussian_sampler.sample()\n",
    "                wt_h_dis += sample_vec[..., None] @ sample_vec[None, ...]\n",
    "\n",
    "            # take mean over dim k\n",
    "            # n * n\n",
    "            wt_h_dis = wt_h_dis / k\n",
    "\n",
    "        # calculating the alpha according to wtHdis\n",
    "        alpha = gen_alpha(wt_h_dis, reg_lambda, y_data)\n",
    "        alpha = alpha / n\n",
    "\n",
    "        cur_train_acc = test_accuracy_for_10_cls(train_img_ts, train_cls_index_label_ts, w_r, wt_train_img_ts, alpha)\n",
    "        cur_test_acc = test_accuracy_for_10_cls(test_img_ts, test_cls_index_label_ts, w_r, wt_train_img_ts, alpha)\n",
    "\n",
    "        train_acc_list.append(cur_train_acc)\n",
    "        test_acc_list.append(cur_test_acc)\n",
    "\n",
    "    # calculat the average private train acc and private test acc\n",
    "    final_train_acc = sum(train_acc_list) / len(train_acc_list)\n",
    "    final_test_acc = sum(test_acc_list) / len(test_acc_list)\n",
    "\n",
    "    return final_train_acc, final_test_acc\n",
    "\n",
    "\n",
    "# fix beta and delta\n",
    "beta = 1e-6\n",
    "delta = 1e-3\n",
    "\n",
    "# we run different eps exponent, from (0.5, 1.5)\n",
    "eps_exponent_list = [0.5 + i * 0.1 for i in range(11)]\n",
    "\n",
    "for eps_exponent in eps_exponent_list:\n",
    "    eps = 10 ** eps_exponent\n",
    "\n",
    "    # calculate number of Gaussian Samples according to eps\n",
    "    k = cal_k(eps, delta, beta)\n",
    "\n",
    "    print(\"-\" * 50)\n",
    "\n",
    "    print(f\"eps exponent {eps_exponent}, k {k}\")\n",
    "\n",
    "    # add truncated laplace noise on X\n",
    "    wt_train_img_ts = add_truncated_laplace_noise(beta, eps, delta, train_img_ts)\n",
    "\n",
    "    # get private test acc and train acc\n",
    "    final_test_acc, final_train_acc = gaussain_sampling_on_k(h_dis, y_data, reg_lambda, wt_train_img_ts, k)\n",
    "    \n",
    "    print(\"test acc\", final_test_acc)\n",
    "    print(\"train acc\", final_train_acc)\n",
    "\n",
    "    # write the result to corresponding logs\n",
    "    cur_log_path = os.path.join(\"logs\", f\"eps_{eps_exponent}.txt\")\n",
    "    with open(cur_log_path, \"w\") as fw:\n",
    "        # fw.write(f\"eps_exponent {str(eps_exponent)}\\n\")\n",
    "        fw.write(f\"test_acc {str(final_test_acc)}\\n\")\n",
    "        fw.write(f\"train_acc {str(final_train_acc)}\\n\")\n",
    "    \n",
    "\n",
    "    print(\"-\" * 50)\n",
    "\n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "galore",
   "language": "python",
   "name": "python3"
  },
  "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": 2
}