GraphTransformer_for_nlp_task.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install torch torchvision\n",
        "!pip install torch-scatter\n",
        "!pip install torch-sparse\n",
        "!pip install torch-geometric\n",
        "!pip install torch-cluster\n",
        "!pip install torch-spline-conv\n",
        "!pip install torch-sparse-old"
      ],
      "metadata": {
        "id": "W7XdUCAMOY4n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import pdb\n",
        "import pickle\n",
        "\n",
        "from __future__ import division\n",
        "\n",
        "\n",
        "def accuracy(pred, target):\n",
        "    return (pred == target).sum().item() / target.numel()\n",
        "\n",
        "\n",
        "\n",
        "def true_positive(pred, target, num_classes):\n",
        "    out = []\n",
        "    for i in range(num_classes):\n",
        "        out.append(((pred == i) & (target == i)).sum())\n",
        "\n",
        "    return torch.tensor(out)\n",
        "\n",
        "\n",
        "\n",
        "def true_negative(pred, target, num_classes):\n",
        "    out = []\n",
        "    for i in range(num_classes):\n",
        "        out.append(((pred != i) & (target != i)).sum())\n",
        "\n",
        "    return torch.tensor(out)\n",
        "\n",
        "\n",
        "\n",
        "def false_positive(pred, target, num_classes):\n",
        "    out = []\n",
        "    for i in range(num_classes):\n",
        "        out.append(((pred == i) & (target != i)).sum())\n",
        "\n",
        "    return torch.tensor(out)\n",
        "\n",
        "\n",
        "\n",
        "def false_negative(pred, target, num_classes):\n",
        "    out = []\n",
        "    for i in range(num_classes):\n",
        "        out.append(((pred != i) & (target == i)).sum())\n",
        "\n",
        "    return torch.tensor(out)\n",
        "\n",
        "\n",
        "\n",
        "def precision(pred, target, num_classes):\n",
        "    tp = true_positive(pred, target, num_classes).to(torch.float)\n",
        "    fp = false_positive(pred, target, num_classes).to(torch.float)\n",
        "\n",
        "    out = tp / (tp + fp)\n",
        "    out[torch.isnan(out)] = 0\n",
        "\n",
        "    return out\n",
        "\n",
        "\n",
        "\n",
        "def recall(pred, target, num_classes):\n",
        "    tp = true_positive(pred, target, num_classes).to(torch.float)\n",
        "    fn = false_negative(pred, target, num_classes).to(torch.float)\n",
        "\n",
        "    out = tp / (tp + fn)\n",
        "    out[torch.isnan(out)] = 0\n",
        "\n",
        "    return out\n",
        "\n",
        "\n",
        "\n",
        "def f1_score(pred, target, num_classes):\n",
        "    prec = precision(pred, target, num_classes)\n",
        "    rec = recall(pred, target, num_classes)\n",
        "\n",
        "    score = 2 * (prec * rec) / (prec + rec)\n",
        "    score[torch.isnan(score)] = 0\n",
        "\n",
        "    return score\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "2jRKGUN8LDbG"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "\n",
        "\n",
        "def uniform(size, tensor):\n",
        "    bound = 1.0 / math.sqrt(size)\n",
        "    if tensor is not None:\n",
        "        tensor.data.uniform_(-bound, bound)\n",
        "\n",
        "\n",
        "def kaiming_uniform(tensor, fan, a):\n",
        "    bound = math.sqrt(6 / ((1 + a**2) * fan))\n",
        "    if tensor is not None:\n",
        "        tensor.data.uniform_(-bound, bound)\n",
        "\n",
        "\n",
        "def glorot(tensor):\n",
        "    stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))\n",
        "    if tensor is not None:\n",
        "        tensor.data.uniform_(-stdv, stdv)\n",
        "\n",
        "\n",
        "def zeros(tensor):\n",
        "    if tensor is not None:\n",
        "        tensor.data.fill_(0)\n",
        "\n",
        "\n",
        "def ones(tensor):\n",
        "    if tensor is not None:\n",
        "        tensor.data.fill_(1)\n",
        "\n",
        "\n",
        "def reset(nn):\n",
        "    def _reset(item):\n",
        "        if hasattr(item, 'reset_parameters'):\n",
        "            item.reset_parameters()\n",
        "\n",
        "    if nn is not None:\n",
        "        if hasattr(nn, 'children') and len(list(nn.children())) > 0:\n",
        "            for item in nn.children():\n",
        "                _reset(item)\n",
        "        else:\n",
        "            _reset(nn)"
      ],
      "metadata": {
        "id": "FXZPpfkAONFv"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import inspect\n",
        "\n",
        "import torch\n",
        "from torch_geometric.utils import scatter_add\n",
        "\n",
        "special_args = [\n",
        "    'edge_index', 'edge_index_i', 'edge_index_j', 'size', 'size_i', 'size_j'\n",
        "]\n",
        "__size_error_msg__ = ('All tensors which should get mapped to the same source '\n",
        "                      'or target nodes must be of same size in dimension 0.')\n",
        "\n",
        "\n",
        "class MessagePassing(torch.nn.Module):\n",
        "\n",
        "    def __init__(self, aggr='add', flow='source_to_target'):\n",
        "        super(MessagePassing, self).__init__()\n",
        "\n",
        "        self.aggr = aggr\n",
        "        assert self.aggr in ['add', 'mean', 'max']\n",
        "\n",
        "        self.flow = flow\n",
        "        assert self.flow in ['source_to_target', 'target_to_source']\n",
        "\n",
        "        self.__message_args__ = inspect.getfullargspec(self.message)[0][1:]\n",
        "        self.__special_args__ = [(i, arg)\n",
        "                                 for i, arg in enumerate(self.__message_args__)\n",
        "                                 if arg in special_args]\n",
        "        self.__message_args__ = [\n",
        "            arg for arg in self.__message_args__ if arg not in special_args\n",
        "        ]\n",
        "        self.__update_args__ = inspect.getfullargspec(self.update)[0][2:]\n",
        "\n",
        "    def propagate(self, edge_index, size=None, **kwargs):\n",
        "\n",
        "\n",
        "        size = [None, None] if size is None else list(size)\n",
        "        assert len(size) == 2\n",
        "\n",
        "        i, j = (0, 1) if self.flow == 'target_to_source' else (1, 0)\n",
        "        ij = {\"_i\": i, \"_j\": j}\n",
        "\n",
        "        message_args = []\n",
        "        for arg in self.__message_args__:\n",
        "            if arg[-2:] in ij.keys():\n",
        "                tmp = kwargs[arg[:-2]]\n",
        "                if tmp is None:  # pragma: no cover\n",
        "                    message_args.append(tmp)\n",
        "                else:\n",
        "                    idx = ij[arg[-2:]]\n",
        "                    if isinstance(tmp, tuple) or isinstance(tmp, list):\n",
        "                        assert len(tmp) == 2\n",
        "                        if size[1 - idx] is None:\n",
        "                            size[1 - idx] = tmp[1 - idx].size(0)\n",
        "                        if size[1 - idx] != tmp[1 - idx].size(0):\n",
        "                            raise ValueError(__size_error_msg__)\n",
        "                        tmp = tmp[idx]\n",
        "\n",
        "                    if size[idx] is None:\n",
        "                        size[idx] = tmp.size(0)\n",
        "                    if size[idx] != tmp.size(0):\n",
        "                        raise ValueError(__size_error_msg__)\n",
        "\n",
        "                    tmp = torch.index_select(tmp, 0, edge_index[idx])\n",
        "                    message_args.append(tmp)\n",
        "            else:\n",
        "                message_args.append(kwargs[arg])\n",
        "\n",
        "        size[0] = size[1] if size[0] is None else size[0]\n",
        "        size[1] = size[0] if size[1] is None else size[1]\n",
        "\n",
        "        kwargs['edge_index'] = edge_index\n",
        "        kwargs['size'] = size\n",
        "\n",
        "        for (idx, arg) in self.__special_args__:\n",
        "            if arg[-2:] in ij.keys():\n",
        "                message_args.insert(idx, kwargs[arg[:-2]][ij[arg[-2:]]])\n",
        "            else:\n",
        "                message_args.insert(idx, kwargs[arg])\n",
        "\n",
        "        update_args = [kwargs[arg] for arg in self.__update_args__]\n",
        "\n",
        "        out = self.message(*message_args)\n",
        "        out = scatter_add(self.aggr, out, edge_index[i], dim_size=size[i])\n",
        "        out = self.update(out, *update_args)\n",
        "\n",
        "        return out\n",
        "\n",
        "    def message(self, x_j):\n",
        "\n",
        "        return x_j\n",
        "\n",
        "    def update(self, aggr_out):\n",
        "\n",
        "        return aggr_out"
      ],
      "metadata": {
        "id": "lEPiHnIqnma7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch.nn import Parameter\n",
        "from torch_scatter import scatter_add\n",
        "from torch_geometric.nn import MessagePassing\n",
        "from torch_geometric.utils import remove_self_loops, add_self_loops\n",
        "\n",
        "import pdb\n",
        "\n",
        "class GCNConv(MessagePassing):\n",
        "    def __init__(self,\n",
        "                 in_channels,\n",
        "                 out_channels,\n",
        "                 improved=False,\n",
        "                 cached=False,\n",
        "                 bias=True):\n",
        "        super(GCNConv, self).__init__('add')\n",
        "\n",
        "        self.in_channels = in_channels\n",
        "        self.out_channels = out_channels\n",
        "        self.improved = improved\n",
        "        self.cached = cached\n",
        "        self.cached_result = None\n",
        "\n",
        "        self.weight = Parameter(torch.Tensor(in_channels, out_channels))\n",
        "\n",
        "        if bias:\n",
        "            self.bias = Parameter(torch.Tensor(out_channels))\n",
        "        else:\n",
        "            self.register_parameter('bias', None)\n",
        "\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        glorot(self.weight)\n",
        "        zeros(self.bias)\n",
        "        self.cached_result = None\n",
        "\n",
        "\n",
        "    @staticmethod\n",
        "    def norm(edge_index, num_nodes, edge_weight, improved=False, dtype=None):\n",
        "        if edge_weight is None:\n",
        "            edge_weight = torch.ones((edge_index.size(1), ),\n",
        "                                     dtype=dtype,\n",
        "                                     device=edge_index.device)\n",
        "        edge_weight = edge_weight.view(-1)\n",
        "        assert edge_weight.size(0) == edge_index.size(1)\n",
        "\n",
        "        edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)\n",
        "        edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)\n",
        "        loop_weight = torch.full((num_nodes, ),\n",
        "                                 1 if not improved else 2,\n",
        "                                 dtype=edge_weight.dtype,\n",
        "                                 device=edge_weight.device)\n",
        "        edge_weight = torch.cat([edge_weight, loop_weight], dim=0)\n",
        "\n",
        "        row, col = edge_index\n",
        "        \n",
        "        deg = scatter_add(edge_weight, col, dim=0, dim_size=num_nodes)\n",
        "        deg_inv_sqrt = deg.pow(-1)\n",
        "        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0\n",
        "\n",
        "        return edge_index, deg_inv_sqrt[col] * edge_weight\n",
        "\n",
        "\n",
        "    def forward(self, x, edge_index, edge_weight=None):\n",
        "        \"\"\"\"\"\"\n",
        "        x = torch.matmul(x, self.weight)\n",
        "\n",
        "        if not self.cached or self.cached_result is None:\n",
        "            edge_index, norm = self.norm(edge_index, x.size(0), edge_weight,\n",
        "                                         self.improved, x.dtype)\n",
        "            self.cached_result = edge_index, norm\n",
        "        edge_index, norm = self.cached_result\n",
        "\n",
        "        return self.propagate(edge_index, x=x, norm=norm)\n",
        "\n",
        "\n",
        "    def message(self, x_j, norm):\n",
        "        return norm.view(-1, 1) * x_j\n",
        "\n",
        "    def update(self, aggr_out):\n",
        "        if self.bias is not None:\n",
        "            aggr_out = aggr_out + self.bias\n",
        "        return aggr_out\n",
        "\n",
        "    def __repr__(self):\n",
        "        return '{}({}, {})'.format(self.__class__.__name__, self.in_channels,\n",
        "                                   self.out_channels)"
      ],
      "metadata": {
        "id": "9iAWKP0AN1aU"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import math\n",
        "from matplotlib import pyplot as plt\n",
        "import pdb\n",
        "\n",
        "\n",
        "class GTN(nn.Module):\n",
        "    \n",
        "    def __init__(self, num_edge, num_channels, w_in, w_out, num_class,num_layers,norm):\n",
        "        super(GTN, self).__init__()\n",
        "        self.num_edge = num_edge\n",
        "        self.num_channels = num_channels\n",
        "        self.w_in = w_in\n",
        "        self.w_out = w_out\n",
        "        self.num_class = num_class\n",
        "        self.num_layers = num_layers\n",
        "        self.is_norm = norm\n",
        "        layers = []\n",
        "        for i in range(num_layers):\n",
        "            if i == 0:\n",
        "                layers.append(GTLayer(num_edge, num_channels, first=True))\n",
        "            else:\n",
        "                layers.append(GTLayer(num_edge, num_channels, first=False))\n",
        "        self.layers = nn.ModuleList(layers)\n",
        "        self.weight = nn.Parameter(torch.Tensor(w_in, w_out))\n",
        "        self.bias = nn.Parameter(torch.Tensor(w_out))\n",
        "        self.loss = nn.CrossEntropyLoss()\n",
        "        self.linear1 = nn.Linear(self.w_out*self.num_channels, self.w_out)\n",
        "        self.linear2 = nn.Linear(self.w_out, self.num_class)\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        nn.init.xavier_uniform_(self.weight)\n",
        "        nn.init.zeros_(self.bias)\n",
        "\n",
        "    def gcn_conv(self,X,H):\n",
        "        X = torch.mm(X, self.weight)\n",
        "        H = self.norm(H, add=True)\n",
        "        return torch.mm(H.t(),X)\n",
        "\n",
        "    def normalization(self, H):\n",
        "        for i in range(self.num_channels):\n",
        "            if i==0:\n",
        "                H_ = self.norm(H[i,:,:]).unsqueeze(0)\n",
        "            else:\n",
        "                H_ = torch.cat((H_,self.norm(H[i,:,:]).unsqueeze(0)), dim=0)\n",
        "        return H_\n",
        "\n",
        "    def norm(self, H, add=False):\n",
        "        H = H.t()\n",
        "        if add == False:\n",
        "            H = H*((torch.eye(H.shape[0])==0).type(torch.FloatTensor))\n",
        "        else:\n",
        "            H = H*((torch.eye(H.shape[0])==0).type(torch.FloatTensor)) + torch.eye(H.shape[0]).type(torch.FloatTensor)\n",
        "        deg = torch.sum(H, dim=1)\n",
        "        deg_inv = deg.pow(-1)\n",
        "        deg_inv[deg_inv == float('inf')] = 0\n",
        "        deg_inv = deg_inv*torch.eye(H.shape[0]).type(torch.FloatTensor)\n",
        "        H = torch.mm(deg_inv,H)\n",
        "        H = H.t()\n",
        "        return H\n",
        "\n",
        "    def forward(self, A, X, target_x, target):\n",
        "        A = A.unsqueeze(0).permute(0,3,1,2) \n",
        "        Ws = []\n",
        "        for i in range(self.num_layers):\n",
        "            if i == 0:\n",
        "                H, W = self.layers[i](A)\n",
        "            else:\n",
        "                H = self.normalization(H)\n",
        "                H, W = self.layers[i](A, H)\n",
        "            Ws.append(W)\n",
        "        for i in range(self.num_channels):\n",
        "            if i==0:\n",
        "                X_ = F.relu(self.gcn_conv(X,H[i]))\n",
        "            else:\n",
        "                X_tmp = F.relu(self.gcn_conv(X,H[i]))\n",
        "                X_ = torch.cat((X_,X_tmp), dim=1)\n",
        "        X_ = self.linear1(X_)\n",
        "        X_ = F.relu(X_)\n",
        "        y = self.linear2(X_[target_x])\n",
        "        loss = self.loss(y, target)\n",
        "        return loss, y, Ws\n",
        "\n",
        "class GTLayer(nn.Module):\n",
        "    \n",
        "    def __init__(self, in_channels, out_channels, first=True):\n",
        "        super(GTLayer, self).__init__()\n",
        "        self.in_channels = in_channels\n",
        "        self.out_channels = out_channels\n",
        "        self.first = first\n",
        "        if self.first == True:\n",
        "            self.conv1 = GTConv(in_channels, out_channels)\n",
        "            self.conv2 = GTConv(in_channels, out_channels)\n",
        "        else:\n",
        "            self.conv1 = GTConv(in_channels, out_channels)\n",
        "    \n",
        "    def forward(self, A, H_=None):\n",
        "        if self.first == True:\n",
        "            a = self.conv1(A)\n",
        "            b = self.conv2(A)\n",
        "            H = torch.bmm(a,b)\n",
        "            W = [(F.softmax(self.conv1.weight, dim=1)).detach(),(F.softmax(self.conv2.weight, dim=1)).detach()]\n",
        "        else:\n",
        "            a = self.conv1(A)\n",
        "            H = torch.bmm(H_,a)\n",
        "            W = [(F.softmax(self.conv1.weight, dim=1)).detach()]\n",
        "        return H,W\n",
        "\n",
        "class GTConv(nn.Module):\n",
        "    \n",
        "    def __init__(self, in_channels, out_channels):\n",
        "        super(GTConv, self).__init__()\n",
        "        self.in_channels = in_channels\n",
        "        self.out_channels = out_channels\n",
        "        self.weight = nn.Parameter(torch.Tensor(out_channels,in_channels,1,1))\n",
        "        self.bias = None\n",
        "        self.scale = nn.Parameter(torch.Tensor([0.1]), requires_grad=False)\n",
        "        self.reset_parameters()\n",
        "    def reset_parameters(self):\n",
        "        n = self.in_channels\n",
        "        nn.init.constant_(self.weight, 0.1)\n",
        "        if self.bias is not None:\n",
        "            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)\n",
        "            bound = 1 / math.sqrt(fan_in)\n",
        "            nn.init.uniform_(self.bias, -bound, bound)\n",
        "\n",
        "    def forward(self, A):\n",
        "        A = torch.sum(A*F.softmax(self.weight, dim=1), dim=1)\n",
        "        return A"
      ],
      "metadata": {
        "id": "RJMYKAlXg-I-"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "epochs = 40\n",
        "node_dim = 64\n",
        "num_channels = 2\n",
        "lr = 0.005\n",
        "weight_decay = 0.001\n",
        "num_layers = 2\n",
        "norm = True\n",
        "adaptive_lr = False\n",
        "\n",
        "with open('/content/data/node_features.pkl','rb') as f:\n",
        "    node_features = pickle.load(f)\n",
        "with open('/content/data/edges.pkl','rb') as f:\n",
        "    edges = pickle.load(f)\n",
        "with open('/content/data/node_features.pkl','rb') as f:\n",
        "    labels = pickle.load(f)\n",
        "\n",
        "\n",
        "num_nodes = edges[0].shape[0]\n",
        "A = []\n",
        "\n",
        "for i,edge in enumerate(edges):\n",
        "    edge_tmp = torch.from_numpy(np.vstack((edge.nonzero()[0], edge.nonzero()[1]))).type(torch.cuda.LongTensor)\n",
        "    value_tmp = torch.ones(edge_tmp.shape[1]).type(torch.cuda.FloatTensor)\n",
        "    A.append((edge_tmp,value_tmp))\n",
        "edge_tmp = torch.stack((torch.arange(0,num_nodes),torch.arange(0,num_nodes))).type(torch.cuda.LongTensor)\n",
        "value_tmp = torch.ones(num_nodes).type(torch.cuda.FloatTensor)\n",
        "A.append((edge_tmp,value_tmp))\n",
        "\n",
        "node_features = torch.from_numpy(node_features).type(torch.cuda.FloatTensor)\n",
        "train_node = torch.from_numpy(np.array(labels[0])[:,0]).type(torch.cuda.LongTensor)\n",
        "train_target = torch.from_numpy(np.array(labels[0])[:,1]).type(torch.cuda.LongTensor)\n",
        "\n",
        "valid_node = torch.from_numpy(np.array(labels[1])[:,0]).type(torch.cuda.LongTensor)\n",
        "valid_target = torch.from_numpy(np.array(labels[1])[:,1]).type(torch.cuda.LongTensor)\n",
        "test_node = torch.from_numpy(np.array(labels[2])[:,0]).type(torch.cuda.LongTensor)\n",
        "test_target = torch.from_numpy(np.array(labels[2])[:,1]).type(torch.cuda.LongTensor)\n",
        "\n",
        "\n",
        "num_classes = torch.max(train_target).item()+1\n",
        "\n",
        "train_losses = []\n",
        "train_f1s = []\n",
        "val_losses = []\n",
        "test_losses = []\n",
        "val_f1s = []\n",
        "test_f1s = []\n",
        "final_f1 = 0\n",
        "for cnt in range(5):\n",
        "    best_val_loss = 10000\n",
        "    best_test_loss = 10000\n",
        "    best_train_loss = 10000\n",
        "    best_train_f1 = 0\n",
        "    best_val_f1 = 0\n",
        "    best_test_f1 = 0\n",
        "    model = GTN(num_edge=len(A),\n",
        "            num_channels=num_channels,\n",
        "            w_in = node_features.shape[1],\n",
        "            w_out = node_dim,\n",
        "            num_class=num_classes,\n",
        "            num_nodes = node_features.shape[0],\n",
        "            num_layers= num_layers)\n",
        "    model.cuda()\n",
        "    if adaptive_lr == 'false':\n",
        "      optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)\n",
        "    else:\n",
        "      optimizer = torch.optim.Adam([{'params':model.gcn.parameters()},\n",
        "                    {'params':model.linear1.parameters()},\n",
        "                    {'params':model.linear2.parameters()},\n",
        "                    {\"params\":model.layers.parameters(), \"lr\":0.5}\n",
        "                    ], lr=0.005, weight_decay=0.001)\n",
        "    loss = nn.CrossEntropyLoss()\n",
        "    Ws = []\n",
        "    for i in range(50):\n",
        "      print('Epoch: ',i+1)\n",
        "      for param_group in optimizer.param_groups:\n",
        "          if param_group['lr'] > 0.005:\n",
        "                param_group['lr'] = param_group['lr'] * 0.9\n",
        "    model.train()\n",
        "    model.zero_grad()\n",
        "    loss, y_train, _ = model(A, node_features, train_node, train_target)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    train_f1 = torch.mean(f1_score(torch.argmax(y_train,dim=1), train_target, num_classes=3)).cpu().numpy()\n",
        "    print('Train - Loss: {}, Macro_F1: {}'.format(loss.detach().cpu().numpy(), train_f1))\n",
        "    model.eval()\n",
        "    \n",
        "    # Valid\n",
        "    with torch.no_grad():\n",
        "        val_loss, y_valid,_ = model.forward(A, node_features, valid_node, valid_target)\n",
        "        val_f1 = torch.mean(f1_score(torch.argmax(y_valid,dim=1), valid_target, num_classes=3)).cpu().numpy()\n",
        "        print('Valid - Loss: {}, Macro_F1: {}'.format(val_loss.detach().cpu().numpy(), val_f1))\n",
        "        test_loss, y_test,W = model.forward(A, node_features, test_node, test_target)\n",
        "        test_f1 = torch.mean(f1_score(torch.argmax(y_test,dim=1), test_target, num_classes=3)).cpu().numpy()\n",
        "        test_acc = accuracy(torch.argmax(y_test,dim=1), test_target)\n",
        "        print('Test - Loss: {}, Macro_F1: {}, Acc: {}\\n'.format(test_loss.detach().cpu().numpy(), test_f1, test_acc))\n",
        "        if val_f1 > best_val_f1:\n",
        "            best_val_loss = val_loss.detach().cpu().numpy()\n",
        "            best_test_loss = test_loss.detach().cpu().numpy()\n",
        "            best_train_loss = loss.detach().cpu().numpy()\n",
        "            best_train_f1 = train_f1\n",
        "            best_val_f1 = val_f1\n",
        "            best_test_f1 = test_f1\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "    print('---------------Best Results--------------------')\n",
        "    print('Train - Loss: {}, Macro_F1: {}'.format(best_test_loss, best_train_f1))\n",
        "    print('Valid - Loss: {}, Macro_F1: {}'.format(best_val_loss, best_val_f1))\n",
        "    print('Test - Loss: {}, Macro_F1: {}'.format(best_test_loss, best_test_f1))"
      ],
      "metadata": {
        "id": "7CBMDO3FNsss"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "GraphTransformer-for-nlp-task.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}