demo_linear.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RoboGraph\n",
    "This is the demo for the submission of paper\n",
    ">  __Certified Robustness of Graph Convolution Networks for Graph Classification under Topological Attacks__\n",
    "\n",
    "Before running the demo, please make sure all the required packages are installed.\n",
    "\n",
    "A detailed instruction is provided in [README.md](./README.md)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "import tempfile\n",
    "from torch_geometric.datasets import TUDataset\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.data.makedirs import makedirs\n",
    "from robograph.model.gnn import GC_NET, train, eval\n",
    "from tqdm.notebook import tqdm\n",
    "from robograph.utils import process_data, cal_logits\n",
    "\n",
    "from robograph.attack.admm import admm_solver\n",
    "from robograph.attack.cvx_env_solver import cvx_env_solver\n",
    "from robograph.attack.dual import dual_solver\n",
    "from robograph.attack.greedy_attack import Greedy_Attack\n",
    "from robograph.attack.utils import calculate_Fc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph classification with linear activation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "# prepare dataset\n",
    "ds_name = 'ENZYMES'\n",
    "path = osp.join(tempfile.gettempdir(), 'data', ds_name)\n",
    "save_path = osp.join(tempfile.gettempdir(), 'data', ds_name, 'saved')\n",
    "if not osp.isdir(save_path):\n",
    "    makedirs(save_path)\n",
    "dataset = TUDataset(path, name=ds_name, use_node_attr=True)\n",
    "dataset = dataset.shuffle()\n",
    "train_size = len(dataset) // 10 * 3\n",
    "val_size = len(dataset) // 10 * 2\n",
    "train_dataset = dataset[:train_size]\n",
    "val_dataset = dataset[train_size: train_size + val_size]\n",
    "test_dataset = dataset[train_size + val_size:]\n",
    "\n",
    "# prepare dataloader\n",
    "train_loader = DataLoader(train_dataset, batch_size=20)\n",
    "val_loader = DataLoader(val_dataset, batch_size=20)\n",
    "test_loader = DataLoader(test_dataset, batch_size=20)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# create model\n",
    "model = GC_NET(hidden=64,\n",
    "              n_features=dataset.num_features,\n",
    "              n_classes=dataset.num_classes,\n",
    "              act='linear',\n",
    "              pool='avg',\n",
    "              dropout=0.).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training a vanilla model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "best=0\n",
    "for epoch in tqdm(range(200)):\n",
    "    loss_all = train(model, train_loader)\n",
    "    train_acc = eval(model, train_loader)\n",
    "    val_acc = eval(model, val_loader)\n",
    "    if val_acc >= best:\n",
    "        best = val_acc\n",
    "        torch.save(model.state_dict(), osp.join(save_path, \"result.pk\"))\n",
    "    \n",
    "    tqdm.write(\"epoch {:03d} \".format(epoch+1) + \n",
    "              \"train_loss {:.4f} \".format(loss_all) +\n",
    "              \"train_acc {:.4f} \".format(train_acc) +\n",
    "              \"val_acc {:.4f} \".format(val_acc))\n",
    "test_acc = eval(model, test_loader, testing=True, save_path=save_path)\n",
    "print(\"test_acc {:.4f}\".format(test_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Robustness certificate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W = model.conv.weight.detach().cpu().numpy().astype(np.float64)\n",
    "U = model.lin.weight.detach().cpu().numpy().astype(np.float64)\n",
    "\n",
    "k = dataset.num_classes\n",
    "\n",
    "# counter of certifiably robust and vulnerable \n",
    "robust_dual = 0\n",
    "robust_cvx = 0\n",
    "vul_admm = 0\n",
    "vul_admm_g = 0\n",
    "vul_greedy = 0\n",
    "\n",
    "# counter of correct classification\n",
    "correct = 0\n",
    "\n",
    "# attacker settings\n",
    "strength = 3\n",
    "delta_g = 10\n",
    "\n",
    "# setting for solvers\n",
    "dual_params = dict(iter=200, nonsmooth_init='random')\n",
    "cvx_params = dict(iter=400, lr=0.3, verbose=0, constr='1+2+3', \n",
    "                 activation='linear', algo='swapping', nonsmooth_init='subgrad')\n",
    "admm_params = dict(iter=200, mu=1)\n",
    "\n",
    "for data in tqdm(test_dataset, desc='across graphs'):\n",
    "    A, X, y = process_data(data)\n",
    "    deg = A.sum(1)\n",
    "    n_nodes = A.shape[0]\n",
    "    n_edges = np.count_nonzero(A) // 2\n",
    "    \n",
    "    delta_l = np.minimum(np.maximum(deg - np.max(deg) + strength, 0), n_nodes - 1).astype(int)\n",
    "    # delta_g\n",
    "    \n",
    "    logits = cal_logits(A, X@W, U, act='linear')\n",
    "    c_pred = logits.argmax()\n",
    "    \n",
    "    if c_pred != y:\n",
    "        continue\n",
    "    correct += 1\n",
    "    fc_vals_orig = [0] * k\n",
    "    fc_vals_dual = [0] * k\n",
    "    fc_vals_cvx = [0] * k\n",
    "    fc_vals_admm = [0] * k\n",
    "    fc_vals_admm_g = [0] * k\n",
    "    fc_vals_greedy = [0] * k\n",
    "    \n",
    "    \n",
    "    for c in tqdm(range(k), desc='across labels', leave=False):\n",
    "        if c == y:\n",
    "            continue\n",
    "        u = U[y] - U[c]\n",
    "        XW = X@W\n",
    "        \n",
    "        # fc_val_orig\n",
    "        fc_vals_orig[c] = calculate_Fc(A, XW, u / n_nodes)\n",
    "        \n",
    "        # fc_val_dual\n",
    "        dual_sol = dual_solver(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g, **dual_params)\n",
    "        fc_vals_dual[c] = dual_sol['opt_f']\n",
    "        \n",
    "        # fc_val_cvx\n",
    "        cvx_sol =cvx_env_solver(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g, **cvx_params)\n",
    "        fc_vals_cvx[c] = cvx_sol['opt_f']\n",
    "        \n",
    "        # fc_val_admm\n",
    "        admm_params['init_B'] = dual_sol['opt_A']\n",
    "        admm_sol = admm_solver(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g, **admm_params)\n",
    "        fc_vals_admm[c] =  admm_sol['opt_f']\n",
    "        \n",
    "        # fc_val_admm_g: admm + greedy\n",
    "        attack = Greedy_Attack(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g)\n",
    "        if np.array_equal(admm_sol['opt_A'], admm_sol['opt_A'].T):\n",
    "            admm_A = admm_sol['opt_A']\n",
    "        else:\n",
    "            admm_A = np.minimum(admm_sol['opt_A'], admm_sol['opt_A'].T)\n",
    "        admm_g_sol = attack.attack(admm_A)  # init from admm\n",
    "        fc_vals_admm_g[c] = admm_g_sol['opt_f']\n",
    "        \n",
    "        # fc_val_greedy\n",
    "        attack = Greedy_Attack(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g)\n",
    "        greedy_sol = attack.attack(A)  # init from A\n",
    "        fc_vals_greedy[c] = greedy_sol['opt_f']\n",
    "    \n",
    "    if np.min(fc_vals_dual) >= 0:\n",
    "        robust_dual += 1\n",
    "    if np.min(fc_vals_cvx) >= 0:\n",
    "        robust_cvx += 1\n",
    "    if np.min(fc_vals_admm) < 0:\n",
    "        vul_admm += 1\n",
    "    if np.min(fc_vals_admm_g) < 0:\n",
    "        vul_admm_g += 1\n",
    "    if np.min(fc_vals_greedy) < 0:\n",
    "        vul_greedy += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('dataset {}'.format(ds_name),\n",
    "      'strength {:02d}'.format(strength),\n",
    "      'delta_g {:02d}'.format(delta_g),\n",
    "      'dual {:.2f}'.format(robust_dual / correct),\n",
    "      'cvx {:.2f}'.format(robust_cvx / correct),\n",
    "      'admm rate {:.2f}'.format(vul_admm / correct),\n",
    "      'admm_g rate {:.2f}'.format(vul_admm_g / correct),\n",
    "      'greedy rate {:.2f}'.format(vul_greedy / correct),)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Warm start from adversarial sample by greedy method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strength = 3\n",
    "for idx, data in tqdm(enumerate(train_dataset), desc='adverarial examples'):\n",
    "    A, X, y = process_data(data)\n",
    "    deg = A.sum(1)\n",
    "    n_nodes = A.shape[0]\n",
    "    delta_l = np.minimum(np.maximum(deg - np.max(deg) + strength, 0), n_nodes - 1).astype(int)\n",
    "    delta_g = n_nodes * np.max(delta_l)\n",
    "    logits = cal_logits(A, X@W, U, act='linear')\n",
    "    c_pred = logits.argmax()\n",
    "    \n",
    "    fc_vals_greedy = [0] * k\n",
    "    fc_A_greedy = [A] * k\n",
    "    for c in range(k):\n",
    "        u = U[y] - U[c]\n",
    "        XW = X@W\n",
    "        ''' greedy attack '''\n",
    "        attack = Greedy_Attack(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g)\n",
    "        greedy_sol = attack.attack(A)  # init from A\n",
    "        fc_vals_greedy[c] = greedy_sol['opt_f']\n",
    "        fc_A_greedy[c] = greedy_sol['opt_A']\n",
    "    pick_idx = np.argmin(fc_vals_greedy)\n",
    "    train_dataset[idx].edge_index = torch.tensor(fc_A_greedy[pick_idx].nonzero())\n",
    "torch.save(train_dataset, osp.join(save_path, 'adv_set.pk'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Robust linear model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GC_NET(hidden=64,\n",
    "               n_features=dataset.num_features,\n",
    "               n_classes=dataset.num_classes,\n",
    "               act='linear',\n",
    "               pool='avg',\n",
    "               dropout=0.).to(device)\n",
    "adv = torch.load(osp.join(save_path, 'adv_set.pk'))\n",
    "adv_loader = DataLoader(adv + train_dataset, batch_size=20)\n",
    "\n",
    "best = 0\n",
    "for epoch in tqdm(range(200), desc='epoch'):\n",
    "    loss_all =  train(model, train_loader, robust=True, adv_loader=adv_loader, lamb=0.5)\n",
    "    train_acc = eval(model, train_loader, robust=True)\n",
    "    val_acc = eval(model, val_loader, robust=True)\n",
    "    \n",
    "    if val_acc >= best:\n",
    "        best = val_acc\n",
    "        torch.save(model.state_dict(), osp.join(save_path, 'result_robust.pk'))\n",
    "    #     tqdm.write(\"epoch {:03d} \".format(epoch+1) + \n",
    "    #               \"train_loss {:.4f} \".format(loss_all) +\n",
    "    #               \"train_acc {:.4f} \".format(train_acc) +\n",
    "    #               \"val_acc {:.4f} \".format(val_acc))\n",
    "\n",
    "test_acc = eval(model, test_loader, testing=True, save_path=save_path, robust=True)\n",
    "print(\"test_acc {:.4f}\".format(test_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Robustness certificate with robust model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W = model.conv.weight.detach().cpu().numpy().astype(np.float64)\n",
    "U = model.lin.weight.detach().cpu().numpy().astype(np.float64)\n",
    "\n",
    "k = dataset.num_classes\n",
    "\n",
    "# counter of certifiably robust and vulnerable \n",
    "robust_dual = 0\n",
    "robust_cvx = 0\n",
    "vul_admm = 0\n",
    "vul_admm_g = 0\n",
    "vul_greedy = 0\n",
    "\n",
    "# counter of correct classification\n",
    "correct = 0\n",
    "\n",
    "# attacker settings\n",
    "strength = 3\n",
    "delta_g = 10\n",
    "\n",
    "# setting for solvers\n",
    "dual_params = dict(iter=200, nonsmooth_init='random')\n",
    "cvx_params = dict(iter=400, lr=0.3, verbose=0, constr='1+2+3', \n",
    "                 activation='linear', algo='swapping', nonsmooth_init='subgrad')\n",
    "admm_params = dict(iter=200, mu=1)\n",
    "\n",
    "for data in tqdm(test_dataset, desc='across graphs'):\n",
    "    A, X, y = process_data(data)\n",
    "    deg = A.sum(1)\n",
    "    n_nodes = A.shape[0]\n",
    "    n_edges = np.count_nonzero(A) // 2\n",
    "    \n",
    "    delta_l = np.minimum(np.maximum(deg - np.max(deg) + strength, 0), n_nodes - 1).astype(int)\n",
    "    # delta_g\n",
    "    \n",
    "    logits = cal_logits(A, X@W, U, act='linear')\n",
    "    c_pred = logits.argmax()\n",
    "    \n",
    "    if c_pred != y:\n",
    "        continue\n",
    "    correct += 1\n",
    "    fc_vals_orig = [0] * k\n",
    "    fc_vals_dual = [0] * k\n",
    "    fc_vals_cvx = [0] * k\n",
    "    fc_vals_admm = [0] * k\n",
    "    fc_vals_admm_g = [0] * k\n",
    "    fc_vals_greedy = [0] * k\n",
    "    \n",
    "    \n",
    "    for c in tqdm(range(k), desc='across labels', leave=False):\n",
    "        if c == y:\n",
    "            continue\n",
    "        u = U[y] - U[c]\n",
    "        XW = X@W\n",
    "        \n",
    "        # fc_val_orig\n",
    "        fc_vals_orig[c] = calculate_Fc(A, XW, u / n_nodes)\n",
    "        \n",
    "        # fc_val_dual\n",
    "        dual_sol = dual_solver(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g, **dual_params)\n",
    "        fc_vals_dual[c] = dual_sol['opt_f']\n",
    "        \n",
    "        # fc_val_cvx\n",
    "        cvx_sol =cvx_env_solver(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g, **cvx_params)\n",
    "        fc_vals_cvx[c] = cvx_sol['opt_f']\n",
    "        \n",
    "        # fc_val_admm\n",
    "        admm_params['init_B'] = dual_sol['opt_A']\n",
    "        admm_sol = admm_solver(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g, **admm_params)\n",
    "        fc_vals_admm[c] =  admm_sol['opt_f']\n",
    "        \n",
    "        # fc_val_admm_g: admm + greedy\n",
    "        attack = Greedy_Attack(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g)\n",
    "        if np.array_equal(admm_sol['opt_A'], admm_sol['opt_A'].T):\n",
    "            admm_A = admm_sol['opt_A']\n",
    "        else:\n",
    "            admm_A = np.minimum(admm_sol['opt_A'], admm_sol['opt_A'].T)\n",
    "        admm_g_sol = attack.attack(admm_A)  # init from admm\n",
    "        fc_vals_admm_g[c] = admm_g_sol['opt_f']\n",
    "        \n",
    "        # fc_val_greedy\n",
    "        attack = Greedy_Attack(A, XW, u / n_nodes, delta_l=delta_l, delta_g=delta_g)\n",
    "        greedy_sol = attack.attack(A)  # init from A\n",
    "        fc_vals_greedy[c] = greedy_sol['opt_f']\n",
    "    \n",
    "    if np.min(fc_vals_dual) >= 0:\n",
    "        robust_dual += 1\n",
    "    if np.min(fc_vals_cvx) >= 0:\n",
    "        robust_cvx += 1\n",
    "    if np.min(fc_vals_admm) < 0:\n",
    "        vul_admm += 1\n",
    "    if np.min(fc_vals_admm_g) < 0:\n",
    "        vul_admm_g += 1\n",
    "    if np.min(fc_vals_greedy) < 0:\n",
    "        vul_greedy += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('dataset {}'.format(ds_name),\n",
    "      'strength {:02d}'.format(strength),\n",
    "      'delta_g {:02d}'.format(delta_g),\n",
    "      'dual {:.2f}'.format(robust_dual / correct),\n",
    "      'cvx {:.2f}'.format(robust_cvx / correct),\n",
    "      'admm rate {:.2f}'.format(vul_admm / correct),\n",
    "      'admm_g rate {:.2f}'.format(vul_admm_g / correct),\n",
    "      'greedy rate {:.2f}'.format(vul_greedy / correct),)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}