GCN_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocess the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import torch\n",
    "# torch.__version__\n",
    "# torch.version.cuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# ! pip install torch==1.7.0+cu110 torchvision==0.8.1+cu110 torchaudio===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html\n",
    "# ! pip install --upgrade --force-reinstall torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html\n",
    "# ! pip install --upgrade --force-reinstall torch-sparse -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html\n",
    "# ! pip install --upgrade --force-reinstall torch-cluster -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html\n",
    "# ! pip install --upgrade --force-reinstall torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.7.0+cu110.html\n",
    "# ! pip install torch-geometric\n",
    "# ! pip install numpy==1.18.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0,1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Sequential as Seq, Linear, ReLU\n",
    "\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "from collections import Counter\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from sklearn import metrics\n",
    "from sklearn.metrics import classification_report, f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from torch_geometric.nn import MessagePassing\n",
    "from torch_geometric.nn import GraphConv, TopKPooling, GatedGraphConv\n",
    "from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp\n",
    "\n",
    "from torch_geometric.data import DataLoader, InMemoryDataset, Data\n",
    "\n",
    "from torch_geometric.utils import remove_self_loops, add_self_loops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conll_data = pickle.load(open('conll_graph_all.pickle', 'rb'))\n",
    "vocabulary = pickle.load(open('vocabulary_all.pickle', 'rb'))\n",
    "voc2id = {key:{l: i for i, l in enumerate(vocabulary[key])} for key in vocabulary}\n",
    "labels = pickle.load(open('labels.pickle', 'rb'))\n",
    "label2id = {l: i for i, l in enumerate(labels)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CoNLLDataset(InMemoryDataset):\n",
    "    def __init__(self, data, split, voc2id, labels, window_size=3, root='.', transform=None, pre_transform=None):\n",
    "        self.dataset = data[split]\n",
    "        self.voc2id = voc2id\n",
    "        self.labels = labels\n",
    "        self.label2id = {l: i for i, l in enumerate(labels)}\n",
    "        \n",
    "        self.window_size = window_size\n",
    "        self.split = split\n",
    "        \n",
    "        super(CoNLLDataset, self).__init__(root, transform, pre_transform)\n",
    "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
    "\n",
    "        \n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return []\n",
    "    \n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return ['/data/graphner_embeddings/node_selfloop_simple_features/'+self.split]\n",
    "\n",
    "    def download(self):\n",
    "        pass\n",
    "    \n",
    "    def process(self):\n",
    "        data_list = []        \n",
    "    \n",
    "        for doc_i, doc in tqdm(enumerate(self.dataset), total=len(self.dataset)):\n",
    "            features = ['word', 'pos', 'chunk', 'extra'] # to keep the same order\n",
    "\n",
    "            word_id  = [self.voc2id['word'][doc['word'][0]]]\n",
    "            len_so_far = len(self.voc2id['word'])\n",
    "            pos_id   = [len_so_far + self.voc2id['pos'][doc['pos'][0]]]\n",
    "            len_so_far += len(self.voc2id['pos']) \n",
    "            chunk_id = [len_so_far + self.voc2id['chunk'][doc['chunk'][0]]]\n",
    "            len_so_far += len(self.voc2id['chunk'])\n",
    "\n",
    "            extra_ids = []\n",
    "            for v in doc['extra']:\n",
    "                extra_ids.append(len_so_far + self.voc2id['extra'][v])\n",
    "            len_so_far += len(self.voc2id['extra'])\n",
    "\n",
    "            window_size = max(len(doc['left_context']), len(doc['right_context'])) if self.window_size == 'all' \\\n",
    "                        else self.window_size\n",
    "\n",
    "            ids = word_id + pos_id + chunk_id + extra_ids\n",
    "            edges = [(i+1, 0) for i, n in enumerate(ids[1:])]\n",
    "\n",
    "            context_word_ids = []\n",
    "            for window  in [doc['left_context'][-window_size:], doc['right_context'][:window_size]]:\n",
    "                for i, w in enumerate(window):\n",
    "\n",
    "                    # All nodes relate to the central node\n",
    "                    # edges.append((len(ids), 0)) \n",
    "\n",
    "                    # Nodes linked by order in the sentence\n",
    "                    edges.append((len(ids), 0 if i == 0 else (len(ids) - 1)))\n",
    "                    ids.append(self.voc2id['word'][w])\n",
    "            \n",
    "            # add self loop\n",
    "            edges.append((0, 0))\n",
    "\n",
    "            x = torch.LongTensor(ids).unsqueeze(1)\n",
    "            y = torch.tensor([self.label2id[doc['label']]])\n",
    "            edge_index = torch.tensor(list(zip(*edges)))\n",
    "\n",
    "            try:\n",
    "                data = Data(x=x, y=y, edge_index=edge_index)\n",
    "                data_list.append(data)\n",
    "            except:\n",
    "                if len(edge_index) == 0:\n",
    "                    continue\n",
    "                print(x, y, edge_index)\n",
    "                break\n",
    "\n",
    "        data, slices = self.collate(data_list)\n",
    "        torch.save((data, slices), '/data/graphner_embeddings/node_selfloop_simple_features/'+self.split)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = CoNLLDataset(data=conll_data, split='train', voc2id=voc2id, labels=labels, window_size='all')\n",
    "val_dataset = CoNLLDataset(data=conll_data, split='validation', voc2id=voc2id, labels=labels, window_size='all')\n",
    "test_dataset = CoNLLDataset(data=conll_data, split='test', voc2id=voc2id, labels=labels, window_size=('all'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAGEConv(MessagePassing):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(SAGEConv, self).__init__(aggr='max') #  \"Max\" aggregation.\n",
    "        self.lin = torch.nn.Linear(in_channels, out_channels)\n",
    "        self.act = torch.nn.ReLU()\n",
    "        self.update_lin = torch.nn.Linear(in_channels + out_channels, in_channels, bias=False)\n",
    "        self.update_act = torch.nn.ReLU()\n",
    "        \n",
    "    def forward(self, x, edge_index):\n",
    "        # x has shape [N, in_channels]\n",
    "        # edge_index has shape [2, E]\n",
    "        \n",
    "        \n",
    "        edge_index, _ = remove_self_loops(edge_index)\n",
    "        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))\n",
    "        \n",
    "        \n",
    "        return self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x)\n",
    "\n",
    "    def message(self, x_j):\n",
    "        # x_j has shape [E, in_channels]\n",
    "\n",
    "        x_j = self.lin(x_j)\n",
    "        x_j = self.act(x_j)\n",
    "        \n",
    "        return x_j\n",
    "\n",
    "    def update(self, aggr_out, x):\n",
    "        # aggr_out has shape [N, out_channels]\n",
    "\n",
    "        new_embedding = torch.cat([aggr_out, x], dim=1)\n",
    "        \n",
    "        new_embedding = self.update_lin(new_embedding)\n",
    "        new_embedding = self.update_act(new_embedding)\n",
    "        \n",
    "        return new_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_dim = 512\n",
    "dim_input = sum(len(voc2id[k]) for k in ['word', 'pos', 'chunk', 'extra', 'classes'])\n",
    "n_labels = len(labels)\n",
    "dropout_rate = 0.5\n",
    "\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.conv1 = SAGEConv(embed_dim, embed_dim)\n",
    "        self.pool1 = TopKPooling(embed_dim, ratio=0.8)\n",
    "        self.conv2 = SAGEConv(embed_dim, embed_dim)\n",
    "        self.pool2 = TopKPooling(embed_dim, ratio=0.8)\n",
    "        self.conv3 = SAGEConv(embed_dim, embed_dim)\n",
    "        self.pool3 = TopKPooling(embed_dim, ratio=0.8)\n",
    "        self.item_embedding = torch.nn.Embedding(num_embeddings=dim_input, embedding_dim=embed_dim)\n",
    "        # self.lin0 = torch.nn.Linear(dim_input, embed_dim)\n",
    "        self.lin1 = torch.nn.Linear(embed_dim*2, embed_dim)\n",
    "        self.lin2 = torch.nn.Linear(embed_dim, embed_dim//2)\n",
    "        self.lin3 = torch.nn.Linear(embed_dim//2, n_labels)\n",
    "        self.bn1 = torch.nn.BatchNorm1d(embed_dim)\n",
    "        self.bn2 = torch.nn.BatchNorm1d(embed_dim //2)\n",
    "        self.act1 = torch.nn.ReLU()\n",
    "        self.act2 = torch.nn.ReLU()        \n",
    "  \n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch = data.x, data.edge_index, data.batch\n",
    "        # x = self.lin0(x.float())\n",
    "        x = self.item_embedding(x)\n",
    "        x = x.squeeze(1)\n",
    "\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "        \n",
    "        z = self.pool1(x, edge_index, None, batch)\n",
    "        x, edge_index, _, batch, _, _ = z\n",
    "        x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)\n",
    "\n",
    "        x = F.relu(self.conv2(x, edge_index))\n",
    "     \n",
    "        x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)\n",
    "        x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)\n",
    "\n",
    "        x = F.relu(self.conv3(x, edge_index))\n",
    "\n",
    "        x, edge_index, _, batch, _, _ = self.pool3(x, edge_index, None, batch)\n",
    "        x3 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)\n",
    "\n",
    "        x = x1 + x2 + x3\n",
    "\n",
    "        x = self.lin1(x)\n",
    "        x = self.act1(x)\n",
    "        x = self.lin2(x)\n",
    "        x = self.act2(x)      \n",
    "        x = F.dropout(x, p=dropout_rate, training=self.training)\n",
    "\n",
    "        # x = torch.sigmoid(self.lin3(x))\n",
    "        x = self.lin3(x).squeeze(1)\n",
    "        # print(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Net().to(device)\n",
    "\n",
    "writer = SummaryWriter(comment='gcn-justids-winall-wei2-lr1e3-mom0.9-wd5e4-embdim512-dr0.5-bs64-notall', log_dir=None,)\n",
    "logs = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)\n",
    "test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_params = {'lr': 5e-4,\n",
    "                    'momentum': 0.9, \n",
    "                    'weight_decay': 5e-4,\n",
    "                   }\n",
    "\n",
    "device = torch.device('cuda')\n",
    "weights1 = [0.5530, 1.0000, 0.0317, 0.4590, 0.4120]\n",
    "weights2 = [0.7436, 1.0000, 0.1780, 0.6775, 0.6419]\n",
    "\n",
    "weights = torch.Tensor(weights2).to(device)\n",
    "\n",
    "optimizer = torch.optim.SGD(model.parameters(),  **optimizer_params)\n",
    "loss_fn = torch.nn.CrossEntropyLoss(weight=weights)\n",
    "# print()#weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(loader):\n",
    "    model.train()\n",
    "\n",
    "    loss_all = 0\n",
    "    print_loss_every = int(len(loader) / 5)\n",
    "    for ii, data in enumerate(tqdm(loader, total=len(loader))):\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        label = data.y.to(device)\n",
    "        loss = loss_fn(output, label)\n",
    "        loss.backward()\n",
    "        loss_all += data.num_graphs * loss.item()\n",
    "        optimizer.step()\n",
    "        if ii % print_loss_every == 0:\n",
    "            print('== Loss at', ii, ' : ', data.num_graphs * loss.item())\n",
    "\n",
    "    return loss_all / len(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(loader, split=''):\n",
    "    print('==== Evaluation on the', split.upper(), ' set ====')\n",
    "    model.eval()\n",
    "\n",
    "    predictions = []\n",
    "    labels = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for data in tqdm(loader, total=len(loader)):\n",
    "\n",
    "            data = data.to(device)\n",
    "            pred = torch.nn.Softmax(dim=1)(model(data)).detach().cpu().numpy()\n",
    "            label = data.y.detach().cpu().numpy()\n",
    "            \n",
    "            predictions.append(np.argmax(pred, axis=1))\n",
    "            labels.append(label)\n",
    "\n",
    "    predictions = np.hstack(predictions)\n",
    "    labels = np.hstack(labels)\n",
    "    \n",
    "    print(classification_report(predictions, labels, digits=4))\n",
    "    \n",
    "    micro_F1 = metrics.f1_score(labels, predictions, average='micro')\n",
    "    macro_F1 = metrics.f1_score(labels, predictions, average='macro')\n",
    "    weighted_F1 = metrics.f1_score(labels, predictions, average='weighted')\n",
    "\n",
    "        \n",
    "    return (micro_F1, macro_F1, weighted_F1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for epoch in range(5):\n",
    "    loss = train(train_loader)\n",
    "    val_acc = evaluate(val_loader, 'val')\n",
    "    logs.append((loss, val_acc))\n",
    "    \n",
    "    writer.add_scalar(\"Learning_rate\", optimizer_params['lr'], len(logs))\n",
    "    writer.add_scalar(\"Loss/train\", loss, len(logs))\n",
    "    writer.add_scalar(\"micro_F1/dev\", val_acc[0], len(logs))\n",
    "    writer.add_scalar(\"macro_F1/dev\", val_acc[1], len(logs))\n",
    "    writer.add_scalar(\"weighted_F1/dev\", val_acc[2], len(logs))\n",
    "\n",
    "    print(f'Epoch: {len(logs):03d}, Loss: {loss:.5f}')\n",
    "    # print('train_acc', train_acc)\n",
    "    print('val_acc', val_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_acc = evaluate(test_loader, 'test')\n",
    "print('test_acc', test_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}