ZO.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2ccaf0f8-2b69-43cd-bfbf-70df1547a1fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openke\n",
    "import time\n",
    "import torch\n",
    "import pandas as pd\n",
    "from openke.config import Trainer, Tester\n",
    "from openke.module.model import TransE, TransD, TransH, RotatE\n",
    "from openke.module.loss import MarginLoss\n",
    "from openke.module.strategy import NegativeSampling\n",
    "from openke.data import TrainDataLoader, TestDataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9b6665ca-453e-4fd0-a9cf-c60eb3c00346",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Files Path : ./benchmarks/FB15K237/train2id.txt\n",
      "Entity Files Path : ./benchmarks/FB15K237/entity2id.txt\n",
      "Relation Files Path : ./benchmarks/FB15K237/relation2id.txt\n",
      "The toolkit is importing datasets.\n",
      "The total of relations is 237.\n",
      "The total of entities is 14541.\n",
      "----------------------------------------------------------------------\n",
      "The total of train triples is 272115.\n",
      "Training Files Path : ./benchmarks/FB15K237/remain_0.45_unlearning.txt\n",
      "Entity Files Path : ./benchmarks/FB15K237/entity2id.txt\n",
      "Relation Files Path : ./benchmarks/FB15K237/relation2id.txt\n",
      "The toolkit is importing datasets.\n",
      "The total of relations is 237.\n",
      "The total of entities is 14541.\n"
     ]
    }
   ],
   "source": [
    "percent = \"0.45\"\n",
    "train_dataloader = TrainDataLoader(\n",
    "    in_path = None,\n",
    "    tri_file = \"./benchmarks/FB15K237/train2id.txt\",\n",
    "    ent_file = \"./benchmarks/FB15K237/entity2id.txt\",\n",
    "    rel_file = \"./benchmarks/FB15K237/relation2id.txt\",\n",
    "    nbatches = 100,\n",
    "    threads = 8,\n",
    "    sampling_mode = \"normal\",\n",
    "    bern_flag = 1,\n",
    "    filter_flag = 1,\n",
    "    neg_ent = 25,\n",
    "    neg_rel = 0)\n",
    "print('----------------------------------------------------------------------')\n",
    "retrain_dataloader = TrainDataLoader(\n",
    "    in_path = None,\n",
    "    tri_file = f'./benchmarks/FB15K237/remain_{percent}_unlearning.txt',\n",
    "    ent_file = \"./benchmarks/FB15K237/entity2id.txt\",\n",
    "    rel_file = \"./benchmarks/FB15K237/relation2id.txt\",\n",
    "    nbatches = 100,\n",
    "    threads = 8,\n",
    "    sampling_mode = \"normal\",\n",
    "    bern_flag = 1,\n",
    "    filter_flag = 1,\n",
    "    neg_ent = 25,\n",
    "    neg_rel = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d503a94d-bcf4-40c4-af3c-d8f88f49888c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The total of train triples is 149664.\n"
     ]
    }
   ],
   "source": [
    "def calculate_gradients(model, data):\n",
    "    model.eval()\n",
    "\n",
    "    loss = model.model({\n",
    "        'batch_h': torch.autograd.Variable(torch.from_numpy(data['batch_h']).cuda()),\n",
    "        'batch_t': torch.autograd.Variable(torch.from_numpy(data['batch_t']).cuda()),\n",
    "        'batch_r': torch.autograd.Variable(torch.from_numpy(data['batch_r']).cuda()),\n",
    "        'batch_y': torch.autograd.Variable(torch.from_numpy(data['batch_y']).cuda()),\n",
    "        'mode': data['mode']\n",
    "    })\n",
    "    total_memory = 0\n",
    "    for i in range(torch.cuda.device_count()):\n",
    "        total_memory += torch.cuda.max_memory_allocated(i) / (1024 ** 2)\n",
    "    # print(f\"Total LOSS memory usage across all GPUs: {total_memory} MB\")\n",
    "    loss_scalar = torch.mean(loss)\n",
    "    params_to_update = [param for name, param in model.named_parameters() if name.endswith('.weight')]\n",
    "    grads = torch.autograd.grad(loss_scalar, params_to_update, create_graph=True)\n",
    "    del loss, loss_scalar\n",
    "    torch.cuda.empty_cache()\n",
    "    return grads\n",
    "\n",
    "\n",
    "def calculate_zo_gradients(model, data, epsilon=1e-5):\n",
    "    model.eval()\n",
    "    grads = []\n",
    "    params_to_update = [param for name, param in model.named_parameters() if\n",
    "                        name.endswith('.weight') and param.requires_grad]\n",
    "    for param in params_to_update:\n",
    "        grad = torch.zeros_like(param.data).cuda()\n",
    "        perturb = torch.randn(param.data.shape[1]).cuda() * epsilon\n",
    "        perturb = perturb.unsqueeze(0).expand(param.data.shape)\n",
    "        original_param = param.data.clone()\n",
    "\n",
    "        param.data.add_(perturb)\n",
    "        with torch.no_grad():\n",
    "            loss_pos = model.model({\n",
    "                'batch_h': torch.from_numpy(data['batch_h']).to('cuda'),\n",
    "                'batch_t': torch.from_numpy(data['batch_t']).to('cuda'),\n",
    "                'batch_r': torch.from_numpy(data['batch_r']).to('cuda'),\n",
    "                'batch_y': torch.from_numpy(data['batch_y']).to('cuda'),\n",
    "                'mode': data['mode']\n",
    "            })\n",
    "        loss_pos_scalar = torch.mean(loss_pos)\n",
    "\n",
    "        param.data.copy_(original_param)\n",
    "        param.data.sub_(perturb)\n",
    "        with torch.no_grad():\n",
    "            loss_neg = model.model({\n",
    "                'batch_h': torch.from_numpy(data['batch_h']).to('cuda'),\n",
    "                'batch_t': torch.from_numpy(data['batch_t']).to('cuda'),\n",
    "                'batch_r': torch.from_numpy(data['batch_r']).to('cuda'),\n",
    "                'batch_y': torch.from_numpy(data['batch_y']).to('cuda'),\n",
    "                'mode': data['mode']\n",
    "            })\n",
    "        loss_neg_scalar = torch.mean(loss_neg)\n",
    "    \n",
    "        param.data.copy_(original_param)\n",
    "\n",
    "        grad = (loss_pos_scalar - loss_neg_scalar) / (2 * epsilon) * perturb\n",
    "        grads.append(grad.mean(dim=0))\n",
    "\n",
    "        torch.cuda.empty_cache()\n",
    "    return grads\n",
    "\n",
    "def hvps(grad_all, model_params, h_estimate):\n",
    "    element_product = 0\n",
    "    for grad_elem, v_elem in zip(grad_all, h_estimate):\n",
    "        element_product += torch.sum(grad_elem * v_elem)\n",
    "    return_grads = torch.autograd.grad(element_product, model_params, create_graph=True)\n",
    "    del element_product\n",
    "    torch.cuda.empty_cache()\n",
    "    return return_grads\n",
    "\n",
    "def woodfisher_hvps(grad_all, gamma=1.0):\n",
    "    hessian_product = tuple()\n",
    "    vTv = sum(torch.sum(grad_elem * grad_elem) for grad_elem in grad_all)\n",
    "    for grad_elem in grad_all:\n",
    "        term1 = (1 / gamma) * grad_elem\n",
    "        term2 = (gamma ** (-2) * grad_elem * vTv) / (1 + gamma ** (-1) * vTv)\n",
    "        hessian_estimate = term1 + term2\n",
    "        hessian_product += (hessian_estimate,)\n",
    "    return hessian_product\n",
    "\n",
    "\n",
    "def update_and_save_checkpoint(checkpoint_path, new_checkpoint_path, new_params):\n",
    "    weights = torch.load(checkpoint_path)\n",
    "    weights['ent_embeddings.weight'] = new_params[0]\n",
    "    weights['rel_embeddings.weight'] = new_params[1]\n",
    "    torch.save(weights, new_checkpoint_path)\n",
    "    print(f\"Updated checkpoint saved to {new_checkpoint_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d43602b4-9277-4f8e-bf1b-33e6c7688192",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GIF_unleanring(model, train_dataloader, test_dataloader, epsilon=None, iteration=100, damp=0.0, scale=50):\n",
    "    start_time = time.time()\n",
    "\n",
    "    for data in train_dataloader:\n",
    "        # grad_full = calculate_gradients(model, data)\n",
    "        grad_full = calculate_zo_gradients(model, data, epsilon=epsilon)\n",
    "        break\n",
    "    total_memory = 0\n",
    "    for i in range(torch.cuda.device_count()):\n",
    "        total_memory += torch.cuda.max_memory_allocated(i) / (1024 ** 2)\n",
    "    print(f\"Total GradsFull memory usage across all GPUs: {total_memory} MB\")\n",
    "    for data in test_dataloader:\n",
    "        # grad_removed = calculate_gradients(model, data)\n",
    "        grad_removed = calculate_zo_gradients(model, data, epsilon=epsilon)\n",
    "        break\n",
    "    total_memory = 0\n",
    "    for i in range(torch.cuda.device_count()):\n",
    "        total_memory += torch.cuda.max_memory_allocated(i) / (1024 ** 2)\n",
    "    print(f\"Total GradsRemoved memory usage across all GPUs: {total_memory} MB\")\n",
    "    grad1 = [g1 - g2 for g1, g2 in zip(grad_full, grad_removed)]\n",
    "    grad2 = grad_removed\n",
    "    res_tuple = (grad_full, grad1, grad2)\n",
    "\n",
    "    v = tuple(grad1 - grad2 for grad1, grad2 in zip(res_tuple[1], res_tuple[2]))\n",
    "    h_estimate = tuple(grad1 - grad2 for grad1, grad2 in zip(res_tuple[1], res_tuple[2]))\n",
    "    model_params = [p for p in model.parameters() if p.requires_grad]\n",
    "    \n",
    "    for i in range(iteration):\n",
    "        hv = woodfisher_hvps(res_tuple[0])\n",
    "        # hv = hvps(res_tuple[0], model_params, h_estimate)\n",
    "        with torch.no_grad():\n",
    "            h_estimate = [v1 + (1 - damp) * h_estimate1 - hv1 / scale for v1, h_estimate1, hv1 in\n",
    "                          zip(v, h_estimate, hv)]\n",
    "    print(f\"final h_estimate: {torch.cuda.max_memory_allocated() / (1024 ** 2)} MB\")\n",
    "    params_change = [h_est / scale for h_est in h_estimate]\n",
    "    params_esti = [p1 + p2 for p1, p2 in zip(params_change, model_params)]\n",
    "    total_memory = 0\n",
    "    for i in range(torch.cuda.device_count()):\n",
    "        total_memory += torch.cuda.max_memory_allocated(i) / (1024 ** 2)\n",
    "    print(f\"Total memory usage across all GPUs: {total_memory} MB\")\n",
    "    del grad_full, grad_removed, res_tuple, v, h_estimate, params_change\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    print(time.time() - start_time)\n",
    "\n",
    "    return params_esti"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "78814200-50b8-42cf-8e1d-d12058dd287b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epsilon: 1e-05\n",
      "iteration: 1\n",
      "damp: 0.0\n",
      "scale: 50\n",
      "Total GradsFull memory usage across all GPUs: 574.7001953125 MB\n",
      "Total GradsRemoved memory usage across all GPUs: 574.7001953125 MB\n",
      "final h_estimate: 318.5029296875 MB\n",
      "Total memory usage across all GPUs: 574.7001953125 MB\n",
      "0.12455511093139648\n",
      "Updated checkpoint saved to ./checkpoint/FB15K237/ZOWFGIF_0.45_TransH.ckpt\n",
      "Input Files Path : ./benchmarks/FB15K237/\n",
      "The total of test triples is 20466.\n",
      "The total of valid triples is 17535.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 20466/20466 [00:26<00:00, 785.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.3873741626739502\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no type constraint results:\n",
      "metric:\t\t\t MRR \t\t MR \t\t hit@10 \t hit@3  \t hit@1 \n",
      "l(raw):\t\t\t 0.088742 \t 573.665466 \t 0.206342 \t 0.085654 \t 0.033812 \n",
      "r(raw):\t\t\t 0.255389 \t 169.283493 \t 0.440975 \t 0.273331 \t 0.166960 \n",
      "averaged(raw):\t\t 0.172066 \t 371.474487 \t 0.323659 \t 0.179493 \t 0.100386 \n",
      "\n",
      "l(filter):\t\t 0.115251 \t 418.791168 \t 0.261507 \t 0.119955 \t 0.047835 \n",
      "r(filter):\t\t 0.297220 \t 153.827866 \t 0.513241 \t 0.325613 \t 0.195544 \n",
      "averaged(filter):\t 0.206235 \t 286.309509 \t 0.387374 \t 0.222784 \t 0.121690 \n",
      "0.387374\n"
     ]
    }
   ],
   "source": [
    "epsilon=1e-5\n",
    "iteration=1\n",
    "damp=0.00\n",
    "scale=50\n",
    "results = []\n",
    "\n",
    "embed_model = 'TransH'\n",
    "if embed_model == 'RotatE':\n",
    "    model = RotatE(\n",
    "    \tent_tot = train_dataloader.get_ent_tot(),\n",
    "    \trel_tot = train_dataloader.get_rel_tot(),\n",
    "    \tdim = 200,\n",
    "    \tmargin = 6.0,\n",
    "    \tepsilon = 2.0)\n",
    "    unlearn_model = RotatE(\n",
    "    ent_tot = train_dataloader.get_ent_tot(),\n",
    "    rel_tot = train_dataloader.get_rel_tot(),\n",
    "    dim = 200,\n",
    "    margin = 6.0,\n",
    "    epsilon = 2.0)\n",
    "elif embed_model == 'TransD':\n",
    "    model = TransD(\n",
    "    ent_tot = train_dataloader.get_ent_tot(),\n",
    "    rel_tot = train_dataloader.get_rel_tot(),\n",
    "    dim_e = 200,\n",
    "    dim_r = 200,\n",
    "    p_norm = 1,\n",
    "    norm_flag = True)\n",
    "    unlearn_model = TransD(\n",
    "    ent_tot = train_dataloader.get_ent_tot(),\n",
    "    rel_tot = train_dataloader.get_rel_tot(),\n",
    "    dim_e = 200,\n",
    "    dim_r = 200,\n",
    "    p_norm = 1,\n",
    "    norm_flag = True)\n",
    "else:\n",
    "    model = TransH(\n",
    "    ent_tot = train_dataloader.get_ent_tot(),\n",
    "    rel_tot = train_dataloader.get_rel_tot(),\n",
    "    dim = 200,\n",
    "    p_norm = 1,\n",
    "    norm_flag = True)\n",
    "    unlearn_model = TransH(\n",
    "    ent_tot = train_dataloader.get_ent_tot(),\n",
    "    rel_tot = train_dataloader.get_rel_tot(),\n",
    "    dim = 200,\n",
    "    p_norm = 1,\n",
    "    norm_flag = True)\n",
    "checkpoint_path=f\"./checkpoint/FB15K237/FB15K237_{embed_model}.ckpt\"\n",
    "model = torch.nn.DataParallel(model)\n",
    "model.to('cuda')\n",
    "model.module.load_checkpoint(checkpoint_path)\n",
    "model = NegativeSampling(\n",
    "    model = model,\n",
    "    loss = MarginLoss(margin = 5.0),\n",
    "    batch_size = train_dataloader.get_batch_size()\n",
    ")\n",
    "print('epsilon:', epsilon)\n",
    "print('iteration:', iteration)\n",
    "print('damp:', damp)\n",
    "print('scale:', scale)\n",
    "new_checkpoint_path = f\"./checkpoint/FB15K237/ZOWFGIF_{percent}_{embed_model}.ckpt\"\n",
    "# new_checkpoint_path = f\"./checkpoint/FB15K237/Delete_Edge_{embed_model}_FB15K237.ckpt\"\n",
    "params_esti = GIF_unleanring(model, train_dataloader, retrain_dataloader, epsilon=epsilon, \n",
    "                             iteration=iteration, damp=damp, scale=scale)\n",
    "update_and_save_checkpoint(checkpoint_path=checkpoint_path,\n",
    "                           new_checkpoint_path=new_checkpoint_path,\n",
    "                           new_params=params_esti)\n",
    "test_dataloader = TestDataLoader(\"./benchmarks/FB15K237/\", \"link\")\n",
    "\n",
    "# unlearn_transe = torch.nn.DataParallel(unlearn_transe)\n",
    "# test the model\n",
    "unlearn_model.load_checkpoint(new_checkpoint_path)\n",
    "unlearn_tester = Tester(model = unlearn_model, data_loader = test_dataloader, use_gpu = True)\n",
    "mrr, mr, hit10, hit3, hit1 = unlearn_tester.run_link_prediction(type_constrain = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1844134c-87e1-40a7-866b-e3d54b43139b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.2101750373840332,\n",
       " 280.8727722167969,\n",
       " 0.39367732405662537,\n",
       " 0.227279394865036,\n",
       " 0.12432815134525299)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "no type constraint results:\n",
      "metric:\t\t\t MRR \t\t MR \t\t hit@10 \t hit@3  \t hit@1 \n",
      "l(raw):\t\t\t 0.088743 \t 573.662903 \t 0.206391 \t 0.085654 \t 0.033812 \n",
      "r(raw):\t\t\t 0.255391 \t 169.284912 \t 0.440975 \t 0.273331 \t 0.166960 \n",
      "averaged(raw):\t\t 0.172067 \t 371.473907 \t 0.323683 \t 0.179493 \t 0.100386 \n",
      "\n",
      "l(filter):\t\t 0.117429 \t 408.936127 \t 0.266149 \t 0.121470 \t 0.048862 \n",
      "r(filter):\t\t 0.302921 \t 152.809433 \t 0.521206 \t 0.333089 \t 0.199795 \n",
      "averaged(filter):\t 0.210175 \t 280.872772 \t 0.393677 \t 0.227279 \t 0.124328 \n",
      "0.393677\n"
     ]
    }
   ],
   "source": [
    "mrr, mr, hit10, hit3, hit1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dd199d5-be1e-4ed8-a08b-e296f63bafd0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}