Modelling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4714VKOKo8cK"
   },
   "source": [
    "# Modeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 9124,
     "status": "ok",
     "timestamp": 1620765396200,
     "user": {
      "displayName": "Zekun Wang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh140SMb0c80XwwVdi_E7-Lh2XNkXKVWOakF5qL=s64",
      "userId": "05904147758498762863"
     },
     "user_tz": 300
    },
    "id": "E6TREv-bfXKZ",
    "outputId": "27c66779-c91f-4b46-ac2a-9f5c29002c9b"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import glob\n",
    "\n",
    "os.environ['DGLBACKEND'] = 'pytorch'\n",
    "\n",
    "import torch as th\n",
    "import dgl\n",
    "import numpy as np\n",
    "\n",
    "from gnn.estimator_fns import *\n",
    "from gnn.graph_utils import *\n",
    "from gnn.data import *\n",
    "from gnn.utils import *\n",
    "from gnn.pytorch_model import *\n",
    "from train import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E881gvwobfSz"
   },
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "BmqGrWm-HXdk"
   },
   "outputs": [],
   "source": [
    "file_list = glob.glob('./data/*edgelist.csv')\n",
    "\n",
    "edges = \",\".join(map(lambda x: x.split(\"/\")[-1], [file for file in file_list if \"relation\" in file]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BDQKAwpYo8cW"
   },
   "source": [
    "### Generate graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 318,
     "status": "ok",
     "timestamp": 1620765400708,
     "user": {
      "displayName": "Zekun Wang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh140SMb0c80XwwVdi_E7-Lh2XNkXKVWOakF5qL=s64",
      "userId": "05904147758498762863"
     },
     "user_tz": 300
    },
    "id": "23aFsML_bfS1",
    "outputId": "d584ec26-25d7-4c80-a41e-42785674428f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "numpy version:1.20.3 PyTorch version:1.10.1 DGL version:0.7.2\n",
      "Namespace(training_dir='./data', model_dir='./model/2022_01_24_17_24_08', output_dir='./output', nodes='features.csv', target_ntype='TransactionID', edges='relation*', labels='tags.csv', new_accounts='test.csv', compute_metrics=True, threshold=0, num_gpus=0, optimizer='adam', lr=0.01, n_epochs=1000, n_hidden=32, n_layers=6, weight_decay=0.0005, dropout=0.2, embedding_size=360)\n"
     ]
    }
   ],
   "source": [
    "print('numpy version:{} PyTorch version:{} DGL version:{}'.format(np.__version__,\n",
    "                                                                    th.__version__,\n",
    "                                                                    dgl.__version__))\n",
    "\n",
    "args = parse_args()\n",
    "print(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 233720,
     "status": "ok",
     "timestamp": 1620765634870,
     "user": {
      "displayName": "Zekun Wang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh140SMb0c80XwwVdi_E7-Lh2XNkXKVWOakF5qL=s64",
      "userId": "05904147758498762863"
     },
     "user_tz": 300
    },
    "id": "C4PfzLQKbfS1",
    "outputId": "a00d8856-b3c7-49b4-d3ca-6f98d20620a5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting relation graphs from the following edge lists : ['relation_addr1_edgelist.csv', 'relation_addr2_edgelist.csv', 'relation_card1_edgelist.csv', 'relation_card2_edgelist.csv', 'relation_card3_edgelist.csv', 'relation_card4_edgelist.csv', 'relation_card5_edgelist.csv', 'relation_card6_edgelist.csv', 'relation_DeviceInfo_edgelist.csv', 'relation_DeviceType_edgelist.csv', 'relation_id_01_edgelist.csv', 'relation_id_02_edgelist.csv', 'relation_id_03_edgelist.csv', 'relation_id_04_edgelist.csv', 'relation_id_05_edgelist.csv', 'relation_id_06_edgelist.csv', 'relation_id_07_edgelist.csv', 'relation_id_08_edgelist.csv', 'relation_id_09_edgelist.csv', 'relation_id_10_edgelist.csv', 'relation_id_11_edgelist.csv', 'relation_id_12_edgelist.csv', 'relation_id_13_edgelist.csv', 'relation_id_14_edgelist.csv', 'relation_id_15_edgelist.csv', 'relation_id_16_edgelist.csv', 'relation_id_17_edgelist.csv', 'relation_id_18_edgelist.csv', 'relation_id_19_edgelist.csv', 'relation_id_20_edgelist.csv', 'relation_id_21_edgelist.csv', 'relation_id_22_edgelist.csv', 'relation_id_23_edgelist.csv', 'relation_id_24_edgelist.csv', 'relation_id_25_edgelist.csv', 'relation_id_26_edgelist.csv', 'relation_id_27_edgelist.csv', 'relation_id_28_edgelist.csv', 'relation_id_29_edgelist.csv', 'relation_id_30_edgelist.csv', 'relation_id_31_edgelist.csv', 'relation_id_32_edgelist.csv', 'relation_id_33_edgelist.csv', 'relation_id_34_edgelist.csv', 'relation_id_35_edgelist.csv', 'relation_id_36_edgelist.csv', 'relation_id_37_edgelist.csv', 'relation_id_38_edgelist.csv', 'relation_ProductCD_edgelist.csv', 'relation_P_emaildomain_edgelist.csv', 'relation_R_emaildomain_edgelist.csv', 'relation_TransactionID_edgelist.csv'] \n",
      "Read edges for target<addr1> from edgelist: ./data\\relation_addr1_edgelist.csv\n",
      "Read edges for target<addr2> from edgelist: ./data\\relation_addr2_edgelist.csv\n",
      "Read edges for target<card1> from edgelist: ./data\\relation_card1_edgelist.csv\n",
      "Read edges for target<card2> from edgelist: ./data\\relation_card2_edgelist.csv\n",
      "Read edges for target<card3> from edgelist: ./data\\relation_card3_edgelist.csv\n",
      "Read edges for target<card4> from edgelist: ./data\\relation_card4_edgelist.csv\n",
      "Read edges for target<card5> from edgelist: ./data\\relation_card5_edgelist.csv\n",
      "Read edges for target<card6> from edgelist: ./data\\relation_card6_edgelist.csv\n",
      "Read edges for target<DeviceInfo> from edgelist: ./data\\relation_DeviceInfo_edgelist.csv\n",
      "Read edges for target<DeviceType> from edgelist: ./data\\relation_DeviceType_edgelist.csv\n",
      "Read edges for target<id_01> from edgelist: ./data\\relation_id_01_edgelist.csv\n",
      "Read edges for target<id_02> from edgelist: ./data\\relation_id_02_edgelist.csv\n",
      "Read edges for target<id_03> from edgelist: ./data\\relation_id_03_edgelist.csv\n",
      "Read edges for target<id_04> from edgelist: ./data\\relation_id_04_edgelist.csv\n",
      "Read edges for target<id_05> from edgelist: ./data\\relation_id_05_edgelist.csv\n",
      "Read edges for target<id_06> from edgelist: ./data\\relation_id_06_edgelist.csv\n",
      "Read edges for target<id_07> from edgelist: ./data\\relation_id_07_edgelist.csv\n",
      "Read edges for target<id_08> from edgelist: ./data\\relation_id_08_edgelist.csv\n",
      "Read edges for target<id_09> from edgelist: ./data\\relation_id_09_edgelist.csv\n",
      "Read edges for target<id_10> from edgelist: ./data\\relation_id_10_edgelist.csv\n",
      "Read edges for target<id_11> from edgelist: ./data\\relation_id_11_edgelist.csv\n",
      "Read edges for target<id_12> from edgelist: ./data\\relation_id_12_edgelist.csv\n",
      "Read edges for target<id_13> from edgelist: ./data\\relation_id_13_edgelist.csv\n",
      "Read edges for target<id_14> from edgelist: ./data\\relation_id_14_edgelist.csv\n",
      "Read edges for target<id_15> from edgelist: ./data\\relation_id_15_edgelist.csv\n",
      "Read edges for target<id_16> from edgelist: ./data\\relation_id_16_edgelist.csv\n",
      "Read edges for target<id_17> from edgelist: ./data\\relation_id_17_edgelist.csv\n",
      "Read edges for target<id_18> from edgelist: ./data\\relation_id_18_edgelist.csv\n",
      "Read edges for target<id_19> from edgelist: ./data\\relation_id_19_edgelist.csv\n",
      "Read edges for target<id_20> from edgelist: ./data\\relation_id_20_edgelist.csv\n",
      "Read edges for target<id_21> from edgelist: ./data\\relation_id_21_edgelist.csv\n",
      "Read edges for target<id_22> from edgelist: ./data\\relation_id_22_edgelist.csv\n",
      "Read edges for target<id_23> from edgelist: ./data\\relation_id_23_edgelist.csv\n",
      "Read edges for target<id_24> from edgelist: ./data\\relation_id_24_edgelist.csv\n",
      "Read edges for target<id_25> from edgelist: ./data\\relation_id_25_edgelist.csv\n",
      "Read edges for target<id_26> from edgelist: ./data\\relation_id_26_edgelist.csv\n",
      "Read edges for target<id_27> from edgelist: ./data\\relation_id_27_edgelist.csv\n",
      "Read edges for target<id_28> from edgelist: ./data\\relation_id_28_edgelist.csv\n",
      "Read edges for target<id_29> from edgelist: ./data\\relation_id_29_edgelist.csv\n",
      "Read edges for target<id_30> from edgelist: ./data\\relation_id_30_edgelist.csv\n",
      "Read edges for target<id_31> from edgelist: ./data\\relation_id_31_edgelist.csv\n",
      "Read edges for target<id_32> from edgelist: ./data\\relation_id_32_edgelist.csv\n",
      "Read edges for target<id_33> from edgelist: ./data\\relation_id_33_edgelist.csv\n",
      "Read edges for target<id_34> from edgelist: ./data\\relation_id_34_edgelist.csv\n",
      "Read edges for target<id_35> from edgelist: ./data\\relation_id_35_edgelist.csv\n",
      "Read edges for target<id_36> from edgelist: ./data\\relation_id_36_edgelist.csv\n",
      "Read edges for target<id_37> from edgelist: ./data\\relation_id_37_edgelist.csv\n",
      "Read edges for target<id_38> from edgelist: ./data\\relation_id_38_edgelist.csv\n",
      "Read edges for target<ProductCD> from edgelist: ./data\\relation_ProductCD_edgelist.csv\n",
      "Read edges for target<P_emaildomain> from edgelist: ./data\\relation_P_emaildomain_edgelist.csv\n",
      "Read edges for target<R_emaildomain> from edgelist: ./data\\relation_R_emaildomain_edgelist.csv\n",
      "Will add self loop for target later......\n",
      "Read in features for target nodes\n",
      "Constructed heterograph with the following metagraph structure: Node types ['DeviceInfo', 'DeviceType', 'P_emaildomain', 'ProductCD', 'R_emaildomain', 'addr1', 'addr2', 'card1', 'card2', 'card3', 'card4', 'card5', 'card6', 'id_01', 'id_02', 'id_03', 'id_04', 'id_05', 'id_06', 'id_07', 'id_08', 'id_09', 'id_10', 'id_11', 'id_12', 'id_13', 'id_14', 'id_15', 'id_16', 'id_17', 'id_18', 'id_19', 'id_20', 'id_21', 'id_22', 'id_23', 'id_24', 'id_25', 'id_26', 'id_27', 'id_28', 'id_29', 'id_30', 'id_31', 'id_32', 'id_33', 'id_34', 'id_35', 'id_36', 'id_37', 'id_38', 'target'], Edge types[('DeviceInfo', 'DeviceInfo<>target', 'target'), ('DeviceType', 'DeviceType<>target', 'target'), ('P_emaildomain', 'P_emaildomain<>target', 'target'), ('ProductCD', 'ProductCD<>target', 'target'), ('R_emaildomain', 'R_emaildomain<>target', 'target'), ('addr1', 'addr1<>target', 'target'), ('addr2', 'addr2<>target', 'target'), ('card1', 'card1<>target', 'target'), ('card2', 'card2<>target', 'target'), ('card3', 'card3<>target', 'target'), ('card4', 'card4<>target', 'target'), ('card5', 'card5<>target', 'target'), ('card6', 'card6<>target', 'target'), ('id_01', 'id_01<>target', 'target'), ('id_02', 'id_02<>target', 'target'), ('id_03', 'id_03<>target', 'target'), ('id_04', 'id_04<>target', 'target'), ('id_05', 'id_05<>target', 'target'), ('id_06', 'id_06<>target', 'target'), ('id_07', 'id_07<>target', 'target'), ('id_08', 'id_08<>target', 'target'), ('id_09', 'id_09<>target', 'target'), ('id_10', 'id_10<>target', 'target'), ('id_11', 'id_11<>target', 'target'), ('id_12', 'id_12<>target', 'target'), ('id_13', 'id_13<>target', 'target'), ('id_14', 'id_14<>target', 'target'), ('id_15', 'id_15<>target', 'target'), ('id_16', 'id_16<>target', 'target'), ('id_17', 'id_17<>target', 'target'), ('id_18', 'id_18<>target', 'target'), ('id_19', 'id_19<>target', 'target'), ('id_20', 'id_20<>target', 'target'), ('id_21', 'id_21<>target', 'target'), ('id_22', 'id_22<>target', 'target'), ('id_23', 'id_23<>target', 'target'), ('id_24', 'id_24<>target', 'target'), ('id_25', 'id_25<>target', 'target'), ('id_26', 'id_26<>target', 'target'), ('id_27', 'id_27<>target', 'target'), ('id_28', 'id_28<>target', 'target'), ('id_29', 'id_29<>target', 'target'), ('id_30', 'id_30<>target', 'target'), ('id_31', 'id_31<>target', 'target'), ('id_32', 'id_32<>target', 'target'), ('id_33', 'id_33<>target', 'target'), ('id_34', 'id_34<>target', 'target'), ('id_35', 'id_35<>target', 'target'), ('id_36', 'id_36<>target', 'target'), ('id_37', 'id_37<>target', 'target'), ('id_38', 'id_38<>target', 'target'), ('target', 'self_relation', 'target'), ('target', 'target<>DeviceInfo', 'DeviceInfo'), ('target', 'target<>DeviceType', 'DeviceType'), ('target', 'target<>P_emaildomain', 'P_emaildomain'), ('target', 'target<>ProductCD', 'ProductCD'), ('target', 'target<>R_emaildomain', 'R_emaildomain'), ('target', 'target<>addr1', 'addr1'), ('target', 'target<>addr2', 'addr2'), ('target', 'target<>card1', 'card1'), ('target', 'target<>card2', 'card2'), ('target', 'target<>card3', 'card3'), ('target', 'target<>card4', 'card4'), ('target', 'target<>card5', 'card5'), ('target', 'target<>card6', 'card6'), ('target', 'target<>id_01', 'id_01'), ('target', 'target<>id_02', 'id_02'), ('target', 'target<>id_03', 'id_03'), ('target', 'target<>id_04', 'id_04'), ('target', 'target<>id_05', 'id_05'), ('target', 'target<>id_06', 'id_06'), ('target', 'target<>id_07', 'id_07'), ('target', 'target<>id_08', 'id_08'), ('target', 'target<>id_09', 'id_09'), ('target', 'target<>id_10', 'id_10'), ('target', 'target<>id_11', 'id_11'), ('target', 'target<>id_12', 'id_12'), ('target', 'target<>id_13', 'id_13'), ('target', 'target<>id_14', 'id_14'), ('target', 'target<>id_15', 'id_15'), ('target', 'target<>id_16', 'id_16'), ('target', 'target<>id_17', 'id_17'), ('target', 'target<>id_18', 'id_18'), ('target', 'target<>id_19', 'id_19'), ('target', 'target<>id_20', 'id_20'), ('target', 'target<>id_21', 'id_21'), ('target', 'target<>id_22', 'id_22'), ('target', 'target<>id_23', 'id_23'), ('target', 'target<>id_24', 'id_24'), ('target', 'target<>id_25', 'id_25'), ('target', 'target<>id_26', 'id_26'), ('target', 'target<>id_27', 'id_27'), ('target', 'target<>id_28', 'id_28'), ('target', 'target<>id_29', 'id_29'), ('target', 'target<>id_30', 'id_30'), ('target', 'target<>id_31', 'id_31'), ('target', 'target<>id_32', 'id_32'), ('target', 'target<>id_33', 'id_33'), ('target', 'target<>id_34', 'id_34'), ('target', 'target<>id_35', 'id_35'), ('target', 'target<>id_36', 'id_36'), ('target', 'target<>id_37', 'id_37'), ('target', 'target<>id_38', 'id_38')]\n",
      "Number of nodes of type target : 590540\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feature mean shape:torch.Size([390]), std shape:torch.Size([390])\n"
     ]
    }
   ],
   "source": [
    "args.edges = edges\n",
    "\n",
    "args.edges = get_edgelists('relation*', args.training_dir)\n",
    "\n",
    "g, features, target_id_to_node, id_to_node = construct_graph(args.training_dir,\n",
    "                                                                args.edges,\n",
    "                                                                args.nodes,\n",
    "                                                                args.target_ntype)\n",
    "\n",
    "mean, stdev, features = normalize(th.from_numpy(features))\n",
    "\n",
    "print('feature mean shape:{}, std shape:{}'.format(mean.shape, stdev.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 234609,
     "status": "ok",
     "timestamp": 1620765636388,
     "user": {
      "displayName": "Zekun Wang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh140SMb0c80XwwVdi_E7-Lh2XNkXKVWOakF5qL=s64",
      "userId": "05904147758498762863"
     },
     "user_tz": 300
    },
    "id": "8Bw1ceNebfS2",
    "outputId": "f56605cc-8069-4363-fdd5-df07dd4ac626"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting labels\n",
      "Got labels\n",
      "----Data statistics------'\n",
      "            #Nodes: 726345\n",
      "            #Edges: 19518802\n",
      "            #Features Shape: torch.Size([590540, 390])\n",
      "            #Labeled Test samples: 118108.0\n"
     ]
    }
   ],
   "source": [
    "g.nodes['target'].data['features'] = features\n",
    "\n",
    "print(\"Getting labels\")\n",
    "n_nodes = g.number_of_nodes('target')\n",
    "\n",
    "labels, _, test_mask = get_labels(target_id_to_node,\n",
    "                                            n_nodes,\n",
    "                                            args.target_ntype,\n",
    "                                            os.path.join(args.training_dir, args.labels),\n",
    "                                            os.path.join(args.training_dir, args.new_accounts))\n",
    "print(\"Got labels\")\n",
    "\n",
    "labels = th.from_numpy(labels).float()\n",
    "test_mask = th.from_numpy(test_mask).float()\n",
    "\n",
    "n_nodes = th.sum(th.tensor([g.number_of_nodes(n_type) for n_type in g.ntypes]))\n",
    "n_edges = th.sum(th.tensor([g.number_of_edges(e_type) for e_type in g.etypes]))\n",
    "\n",
    "print(\"\"\"----Data statistics------'\n",
    "            #Nodes: {}\n",
    "            #Edges: {}\n",
    "            #Features Shape: {}\n",
    "            #Labeled Test samples: {}\"\"\".format(n_nodes,\n",
    "                                                    n_edges,\n",
    "                                                    features.shape,\n",
    "                                                    test_mask.sum()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FjhjypJEo8cX"
   },
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "UETWzSN6bfS3"
   },
   "outputs": [],
   "source": [
    "if args.num_gpus:\n",
    "    cuda = True\n",
    "    device = th.device('cuda:0')\n",
    "else:\n",
    "    cuda = False\n",
    "    device = th.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 5097409,
     "status": "ok",
     "timestamp": 1620799361346,
     "user": {
      "displayName": "Zekun Wang",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh140SMb0c80XwwVdi_E7-Lh2XNkXKVWOakF5qL=s64",
      "userId": "05904147758498762863"
     },
     "user_tz": 300
    },
    "id": "AWQp2MxLHXdo",
    "outputId": "5a0c315c-0e7e-4391-ec98-0411d0a4c59a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing Model\n",
      "Initialized Model\n",
      "Starting Model training\n",
      "Epoch 00000, Time(s) 2019.3499, Loss 0.5230, F1 0.0000 \n",
      "Epoch 00001, Time(s) 2408.0788, Loss 17.7347, F1 0.1790 \n",
      "Epoch 00002, Time(s) 2036.9546, Loss 0.2719, F1 0.0000 \n",
      "Epoch 00003, Time(s) 2106.3953, Loss 0.4772, F1 0.0000 \n",
      "Epoch 00004, Time(s) 2024.7838, Loss 0.2795, F1 0.0000 \n",
      "Epoch 00005, Time(s) 2075.7785, Loss 0.2205, F1 0.0000 \n",
      "Epoch 00006, Time(s) 2099.2977, Loss 0.1873, F1 0.0000 \n",
      "Epoch 00007, Time(s) 2134.7374, Loss 0.1464, F1 0.1388 \n",
      "Epoch 00008, Time(s) 2146.4198, Loss 1.4278, F1 0.0000 \n",
      "Epoch 00009, Time(s) 2140.8130, Loss 0.2177, F1 0.0000 \n",
      "Epoch 00010, Time(s) 2187.4841, Loss 0.2550, F1 0.0000 \n",
      "Epoch 00011, Time(s) 2259.3356, Loss 0.2419, F1 0.0000 \n",
      "Epoch 00012, Time(s) 2241.1635, Loss 0.2051, F1 0.0000 \n"
     ]
    }
   ],
   "source": [
    "print(\"Initializing Model\")\n",
    "in_feats = features.shape[1]\n",
    "n_classes = 2\n",
    "\n",
    "ntype_dict = {n_type: g.number_of_nodes(n_type) for n_type in g.ntypes}\n",
    "\n",
    "model = get_model(ntype_dict, g.etypes, vars(args), in_feats, n_classes, device)\n",
    "print(\"Initialized Model\")\n",
    "\n",
    "features = features.to(device)\n",
    "\n",
    "labels = labels.long().to(device)\n",
    "test_mask = test_mask.to(device)\n",
    "# g = g.to(device)\n",
    "\n",
    "loss = th.nn.CrossEntropyLoss()\n",
    "\n",
    "# print(model)\n",
    "optim = th.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)\n",
    "\n",
    "print(\"Starting Model training\")\n",
    "\n",
    "initial_record()\n",
    "\n",
    "model, class_preds, pred_proba = train_fg(model, optim, loss, features, labels, g, g,\n",
    "                                            test_mask, device, args.n_epochs,\n",
    "                                            args.threshold,  args.compute_metrics)\n",
    "print(\"Finished Model training\")\n",
    "\n",
    "print(\"Saving model\") \n",
    "\n",
    "if not os.path.exists(args.model_dir):\n",
    "    os.makedirs(args.model_dir)\n",
    "\n",
    "save_model(g, model, args.model_dir, id_to_node, mean, stdev)\n",
    "print(\"Model and metadata saved\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "20_modeling_pytoch.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.9.7"
  },
  "metadata": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}