train_graph_regression.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SrOiGmrTQBmv"
      },
      "source": [
        "# Environmental Setup\n",
        "\n",
        "\n",
        "Install necessary modules and clone the Github repo."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "s1YuxS8qTA17",
        "outputId": "bcd922bd-eee0-4bc9-df6f-0748acb35101"
      },
      "outputs": [],
      "source": [
        "!pip install wget\n",
        "\n",
        "import os\n",
        "import shutil\n",
        "import wget\n",
        "from urllib.parse import urlparse\n",
        "\n",
        "%cd /content/\n",
        "\n",
        "repo_path = \"https://github.com/KyawYeThu-11/DyGLib.git\"\n",
        "repo_name = os.path.splitext(os.path.basename(urlparse(repo_path).path))[0]\n",
        "\n",
        "if not os.path.exists(repo_name):\n",
        "  !git clone $repo_path\n",
        "  !pip install -r DyGLib/requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "u4Qx50hKysbZ",
        "outputId": "31e4ed0b-b31f-41f1-c9ff-0cfdb2cc7888"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/content/DyGLib\n"
          ]
        }
      ],
      "source": [
        "%cd /content/DyGLib"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "VYknFecnBLzG"
      },
      "outputs": [],
      "source": [
        "import logging\n",
        "import time\n",
        "import sys\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import warnings\n",
        "import shutil\n",
        "import json\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch_geometric.nn import global_mean_pool\n",
        "\n",
        "from models.DyGFormerGraph import DyGFormerGraph\n",
        "from models.DyGFormer import DyGFormer\n",
        "from models.modules import MergeLayer, GraphRegressor\n",
        "from utils.utils import set_random_seed, convert_to_gpu, get_parameter_sizes, create_optimizer\n",
        "from utils.utils import get_neighbor_sampler, NegativeEdgeSampler\n",
        "from utils.metrics import get_link_prediction_metrics\n",
        "from utils.DataLoader import get_idx_data_loader, get_graph_regression_data\n",
        "from utils.EarlyStopping import EarlyStopping\n",
        "from utils.load_configs import get_link_prediction_args"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "URXZEgpZvhwg"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ztBbqeSsgqu3"
      },
      "source": [
        "### Constants and Helper Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hN_4YQGW4Hyd"
      },
      "outputs": [],
      "source": [
        "class Args:\n",
        "    def __init__(self, dataset_name):\n",
        "        self.dataset_name = dataset_name\n",
        "        self.batch_size = 1000\n",
        "        self.model_name = 'DyGFormer'\n",
        "        self.gpu = 0\n",
        "        self.num_neighbors = 20\n",
        "        self.sample_neighbor_strategy = 'recent'\n",
        "        self.time_scaling_factor = 1e-6\n",
        "        self.num_walk_heads = 8\n",
        "        self.num_folds = 5\n",
        "        self.num_heads = 2\n",
        "        self.num_layers = 2\n",
        "        self.walk_length = 1\n",
        "        self.time_gap = 2000\n",
        "        self.time_feat_dim = 100\n",
        "        self.position_feat_dim = 172\n",
        "        self.edge_bank_memory_mode = 'unlimited_memory'\n",
        "        self.time_window_mode = 'fixed_proportion'\n",
        "        self.patch_size = 2\n",
        "        self.channel_embedding_dim = 50\n",
        "        self.max_input_sequence_length = 64\n",
        "        self.learning_rate = 0.005\n",
        "        self.load_checkpoint = False\n",
        "        self.dropout = 0.05\n",
        "        self.scheduler_period = 2\n",
        "        self.num_epochs = 10\n",
        "        self.optimizer = 'Adam'\n",
        "        self.weight_decay = 0.0\n",
        "        self.patience = 20\n",
        "        self.test_interval_epochs = 10\n",
        "        self.negative_sample_strategy = 'random'\n",
        "        self.load_best_configs = False\n",
        "\n",
        "    def __str__(self):\n",
        "        properties = [f\"{key}={value}\" for key, value in self.__dict__.items()]\n",
        "        return f\"Args({', '.join(properties)})\"\n",
        "\n",
        "    def __repr__(self):\n",
        "        return self.__str__()\n",
        "\n",
        "# Create an instance of Args with the loaded configuration\n",
        "args = Args(dataset_name = 'connectome')\n",
        "args.device = f'cuda:{args.gpu}' if torch.cuda.is_available() and args.gpu >= 0 else 'cpu'\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "515H-s4pJ7LJ"
      },
      "outputs": [],
      "source": [
        "def set_up_logger(args):\n",
        "        # set up logger\n",
        "        logging.basicConfig(level=logging.INFO)\n",
        "        logger = logging.getLogger()\n",
        "        logger.setLevel(logging.DEBUG)\n",
        "        os.makedirs(f\"./logs/{args.model_name}/{args.dataset_name}/\", exist_ok=True)\n",
        "        # create file handler that logs debug and higher level messages\n",
        "        fh = logging.FileHandler(f\"./logs/{args.model_name}/{args.dataset_name}/{str(time.time())}.log\")\n",
        "        fh.setLevel(logging.DEBUG)\n",
        "        # create console handler with a higher log level\n",
        "        ch = logging.StreamHandler()\n",
        "        ch.setLevel(logging.WARNING)\n",
        "        # create formatter and add it to the handlers\n",
        "        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')\n",
        "        fh.setFormatter(formatter)\n",
        "        ch.setFormatter(formatter)\n",
        "        # add the handlers to logger\n",
        "        logger.addHandler(fh)\n",
        "        logger.addHandler(ch)\n",
        "\n",
        "        return logger, fh, ch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b0CV9WpdI9_E"
      },
      "outputs": [],
      "source": [
        "def step_forward(dataset_dir, index, subject_id, label, loss_func, mode):\n",
        "  node_raw_features, edge_raw_features, full_data = \\\n",
        "        get_graph_regression_data(dataset_dir=dataset_dir, index=index)\n",
        "\n",
        "  idx_data_loader = get_idx_data_loader(indices_list=list(range(len(full_data.src_node_ids))), batch_size=args.batch_size, shuffle=False)\n",
        "  idx_data_loader_tqdm = tqdm(idx_data_loader, ncols=120)\n",
        "\n",
        "  # initialize validation and test neighbor sampler to retrieve temporal graph\n",
        "  full_neighbor_sampler = get_neighbor_sampler(data=full_data, sample_neighbor_strategy=args.sample_neighbor_strategy,\n",
        "                                                  time_scaling_factor=args.time_scaling_factor, seed=1)\n",
        "  # training, only use training graph\n",
        "  model[0].set_neighbor_sampler(full_neighbor_sampler)\n",
        "\n",
        "  graph_embedding_list = []\n",
        "\n",
        "  for batch_idx, train_data_indices in enumerate(idx_data_loader_tqdm):\n",
        "\n",
        "\n",
        "    train_data_indices = train_data_indices.numpy()\n",
        "    batch_src_node_ids, batch_dst_node_ids, batch_node_interact_times, batch_edge_ids = \\\n",
        "                      full_data.src_node_ids[train_data_indices], full_data.dst_node_ids[train_data_indices], \\\n",
        "                      full_data.node_interact_times[train_data_indices], full_data.edge_ids[train_data_indices]\n",
        "\n",
        "    graph_embedding = model[0](node_raw_features, edge_raw_features, batch_src_node_ids, batch_dst_node_ids, batch_node_interact_times)\n",
        "    graph_embedding_list.append(graph_embedding)\n",
        "\n",
        "    if mode == 'train':\n",
        "      idx_data_loader_tqdm.set_description(f'Epoch: {epoch + 1}, Subject: {subject_id}, training for the {batch_idx + 1}-th batch.')\n",
        "\n",
        "  graph_embeddings_tensor = torch.tensor(np.stack(graph_embedding_list).squeeze()).to(args.device)\n",
        "\n",
        "  predict = model[1](graph_embeddings_tensor)\n",
        "  loss = loss_func(input=predict, target=label)\n",
        "\n",
        "  return predict, loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3tAEIbw47HKM"
      },
      "source": [
        "### Main"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "collapsed": true,
        "id": "GRR-DP0RKJj6",
        "outputId": "5722b5f3-2199-4bce-b196-a325e416955a"
      },
      "outputs": [],
      "source": [
        "percentile_folder = '5-percentile'\n",
        "save_model_folder = f\"saved_models/graph_regression/{args.dataset_name}\"\n",
        "os.makedirs(save_model_folder, exist_ok=True)\n",
        "checkpoint_path = os.path.join(save_model_folder, f'{args.model_name}.pth')\n",
        "\n",
        "logger, fh, ch = set_up_logger(args)\n",
        "\n",
        "tart_time = time.time()\n",
        "\n",
        "logger.info(f'configuration is {args}')\n",
        "\n",
        "# Initialize the model\n",
        "dynamic_backbone = DyGFormerGraph(node_feat_dim=args.position_feat_dim, edge_feat_dim=args.position_feat_dim, time_feat_dim=args.time_feat_dim,\n",
        "                                 channel_embedding_dim=args.channel_embedding_dim, patch_size=args.patch_size,\n",
        "                                 num_layers=args.num_layers, num_heads=args.num_heads, dropout=args.dropout,\n",
        "                                 max_input_sequence_length=args.max_input_sequence_length, device=args.device)\n",
        "\n",
        "graph_regressor = GraphRegressor(in_channels=args.position_feat_dim, hidden_channels=int(args.position_feat_dim/2), out_channels=int(args.position_feat_dim/4))\n",
        "model = nn.Sequential(dynamic_backbone, graph_regressor)\n",
        "\n",
        "# log the model structure\n",
        "logger.info(f'model -> {model}')\n",
        "logger.info(f'model name: {args.model_name}, #parameters: {get_parameter_sizes(model) * 4} B, '\n",
        "                    f'{get_parameter_sizes(model) * 4 / 1024} KB, {get_parameter_sizes(model) * 4 / 1024 / 1024} MB.')\n",
        "\n",
        "# Create the optimizer and scheduler with specified parameters\n",
        "optimizer = create_optimizer(model=model, optimizer_name=args.optimizer, learning_rate=args.learning_rate, weight_decay=args.weight_decay)\n",
        "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.scheduler_period)\n",
        "\n",
        "# Convert the model to GPU if available\n",
        "model = convert_to_gpu(model, device=args.device)\n",
        "\n",
        "\n",
        "epoch_resumed = 0\n",
        "index_resumed = 0\n",
        "# Load checkpoint if specified and exists\n",
        "if args.load_checkpoint == True and os.path.exists(checkpoint_path):\n",
        "    checkpoint = torch.load(checkpoint_path, map_location=args.device)\n",
        "    logger.info(f\"load model {checkpoint_path}\")\n",
        "\n",
        "    model.load_state_dict(checkpoint['model'])\n",
        "    optimizer.load_state_dict(checkpoint['optimizer'])\n",
        "    scheduler.load_state_dict(checkpoint['scheduler'])\n",
        "\n",
        "    epoch_resumed = checkpoint['epoch']\n",
        "    index_resumed = checkpoint['subject_index'] + 1\n",
        "    print(f\"Epoch resumed: {epoch_resumed}\")\n",
        "    print(f\"Index resumed: {index_resumed}\")\n",
        "\n",
        "# Define the mean square loss function\n",
        "loss_func = nn.MSELoss()\n",
        "\n",
        "# Training loop\n",
        "for epoch in range(epoch_resumed, args.num_epochs, 1):\n",
        "\n",
        "    # Training for an epoch starts\n",
        "    model.train()\n",
        "\n",
        "    train_losses = []\n",
        "\n",
        "    language_score_df = pd.read_csv(f'DG_data/{args.dataset_name}/{percentile_folder}/Train_Data_csv/Language_Task_Acc.csv')\n",
        "\n",
        "    # Iterate over each subject in the training data\n",
        "    for index, row in language_score_df.iterrows():\n",
        "        if index < index_resumed:\n",
        "          continue\n",
        "        subject_id = int(row.iloc[2])\n",
        "        score = torch.Tensor([row.iloc[3]]).to(args.device)\n",
        "        subject_folder = os.path.join(f'processed_data/{args.dataset_name}/{percentile_folder}/Train', str(index))\n",
        "\n",
        "        # process every batch for a single subject and compute predictions and loss\n",
        "        predict, loss = step_forward(subject_folder, index, subject_id, score, loss_func, 'train')\n",
        "\n",
        "        train_losses.append(loss.item())\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        scheduler.step()\n",
        "\n",
        "        print(f\"For subject {subject_id}, label: {score.item():.4f}, predict: {predict.item():.4f}, and loss: {loss:.4f}\")\n",
        "\n",
        "        # Save the checkpoint\n",
        "        torch.save({\n",
        "                  'subject_index': index,\n",
        "                  'epoch': epoch,\n",
        "                  'model': model.state_dict(),\n",
        "                  'optimizer': optimizer.state_dict(),\n",
        "                  'scheduler': scheduler.state_dict()},\n",
        "                  checkpoint_path)\n",
        "\n",
        "    # validation\n",
        "    print(\"---------------Validation Starts--------------\")\n",
        "\n",
        "    model.eval()\n",
        "\n",
        "    val_losses = []\n",
        "\n",
        "    with torch.no_grad():\n",
        "        language_score_df = pd.read_csv(f'DG_data/{args.dataset_name}/{percentile_folder}/Val_Data_csv/Language_Task_Acc.csv')\n",
        "\n",
        "        for index, row in language_score_df.iterrows():\n",
        "            subject_id = int(row.iloc[2])\n",
        "            score = torch.Tensor([row.iloc[3]]).to(args.device)\n",
        "            subject_folder = os.path.join(f'processed_data/{args.dataset_name}/{percentile_folder}/Val', str(index))\n",
        "\n",
        "            predict, loss = step_forward(subject_folder, index, subject_id, score, loss_func, 'test')\n",
        "\n",
        "            val_losses.append(loss.item())\n",
        "\n",
        "            print(f\"For subject {subject_id}, label: {score.item():.4f}, predict: {predict.item():.4f}, and loss: {loss:.4f}\")\n",
        "\n",
        "    print(\"-----------------------------------\")\n",
        "    logger.info(f'Epoch: {epoch}, learning rate: {optimizer.param_groups[0][\"lr\"]}, train loss: {np.mean(train_losses):.4f} and val loss: {np.mean(val_losses):.4f}')\n",
        "    print(\"-----------------------------------\")\n",
        "\n",
        "    # Save the results of the current run to a JSON file\n",
        "    result_json = {\n",
        "            \"train losses\": [f'{loss:.4f}' for loss in train_losses],\n",
        "            'val losses': [f'{loss:.4f}' for loss in val_losses]\n",
        "    }\n",
        "\n",
        "    result_json = json.dumps(result_json, indent=4)\n",
        "\n",
        "    save_result_folder = f\"saved_results/graph_regression/{args.dataset_name}\"\n",
        "    os.makedirs(save_result_folder, exist_ok=True)\n",
        "    save_result_path = os.path.join(save_result_folder, f\"{args.save_model_name}.json\")\n",
        "\n",
        "    with open(save_result_path, 'w') as file:\n",
        "        file.write(result_json)\n",
        "\n",
        "torch.save(model.state_dict(), os.path.join(save_model_folder, f'{args.model_name}_final.pth'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vFsUY6EmS7FK"
      },
      "source": [
        "## Testing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mlwy-bdVTLBT"
      },
      "source": [
        "Load the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Rk9ZgOgWTDnS",
        "outputId": "1993de26-1b6e-421c-a065-ae13de31bf7c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Initialize the model\n",
        "dynamic_backbone = DyGFormerGraph(node_feat_dim=args.position_feat_dim, edge_feat_dim=args.position_feat_dim, time_feat_dim=args.time_feat_dim,\n",
        "                                 channel_embedding_dim=args.channel_embedding_dim, patch_size=args.patch_size,\n",
        "                                 num_layers=args.num_layers, num_heads=args.num_heads, dropout=args.dropout,\n",
        "                                 max_input_sequence_length=args.max_input_sequence_length, device=args.device)\n",
        "\n",
        "graph_regressor = GraphRegressor(in_channels=args.position_feat_dim, hidden_channels=int(args.position_feat_dim/2), out_channels=int(args.position_feat_dim/4))\n",
        "model = nn.Sequential(dynamic_backbone, graph_regressor)\n",
        "\n",
        "# Load the pretrained model from the specified path\n",
        "load_model_path = f\"saved_models/graph_regression/{args.dataset_name}/DyGFormer.pth\"\n",
        "checkpoint = torch.load(load_model_path, map_location=args.device)\n",
        "model.load_state_dict(checkpoint['model'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mGX3fGf4TNpb"
      },
      "source": [
        "Test the model with the testing set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 390
        },
        "collapsed": true,
        "id": "C8iturmKTO7T",
        "outputId": "4651366d-b76b-4418-b97d-1e2cfd826e2e"
      },
      "outputs": [],
      "source": [
        "percentile_folder = '5-percentile'\n",
        "test_losses = []\n",
        "loss_func = nn.MSELoss()\n",
        "\n",
        "with torch.no_grad():\n",
        "    language_score_df = pd.read_csv(f'DG_data/{args.dataset_name}/{percentile_folder}/Test_Data_csv/Language_Task_Acc.csv')\n",
        "\n",
        "    for index, row in language_score_df.iterrows():\n",
        "        subject_id = int(row.iloc[2])\n",
        "        score = torch.Tensor([row.iloc[3]]).to(args.device)\n",
        "        subject_folder = os.path.join(f'processed_data/{args.dataset_name}/{percentile_folder}/Test', str(index))\n",
        "\n",
        "        predict, loss = step_forward(subject_folder, index, subject_id, score, loss_func, 'test')\n",
        "\n",
        "        test_losses.append(loss.item())\n",
        "\n",
        "        print(f\"For subject {subject_id}, label: {score.item():.4f}, predict: {predict.item():.4f}, and loss: {loss:.4f}\")\n",
        "\n",
        "    print(f'test loss: {np.mean(test_losses):.4f}')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}