baseline1_tile-xla-end-to-end-train-infer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.009203,
     "end_time": "2023-09-01T16:50:47.024469",
     "exception": false,
     "start_time": "2023-09-01T16:50:47.015266",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Attention!!!\n",
    "\n",
    "This is a very simple but bad quality notebook. \n",
    " - I do not use any sort of ranking loss, which would be better.\n",
    " - My strategy instead is to min-max scale the times and apply L1-loss\n",
    " - My model is also not optimized. It is a relatively simple GNN that embeds the graph and only processes 1 datapoint at a time and is only trained on 1 epoch.\n",
    " - The public score would be much better if you paired this submission with a trained model for layout. Since this only contributes to half of the score.\n",
    " - Have fun playing around with it!\n",
    " \n",
    " \n",
    " # CHANGES\n",
    " - V5 - normalized train and infer targets, use MSE loss, changed evaluation metric to perform top5 mean instead of top5 max for robustness, 5-fold CV\n",
    " - V6 - use SAGEConv instead of GCN, add dropout layer, increase number of paramters, changed evaluation metric to perform top50 mean, 10->20 epochs.\n",
    " - V13 - fixed problem where model weights weren't being reset leading to heavy overfitting...oops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_kg_hide-output": true,
    "papermill": {
     "duration": 264.245878,
     "end_time": "2023-09-01T16:55:11.276628",
     "exception": false,
     "start_time": "2023-09-01T16:50:47.03075",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: http://mirrors.aliyun.com/pypi/simple\n",
      "Requirement already satisfied: torch-geometric in /root/miniconda3/lib/python3.8/site-packages (2.3.1)\n",
      "Requirement already satisfied: torch-scatter in /root/miniconda3/lib/python3.8/site-packages (2.1.1)\n",
      "Requirement already satisfied: jinja2 in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (3.1.2)\n",
      "Requirement already satisfied: scipy in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (1.10.1)\n",
      "Requirement already satisfied: psutil>=5.8.0 in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (5.9.4)\n",
      "Requirement already satisfied: requests in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (2.28.2)\n",
      "Requirement already satisfied: numpy in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (1.24.2)\n",
      "Requirement already satisfied: pyparsing in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (3.0.9)\n",
      "Requirement already satisfied: tqdm in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (4.61.2)\n",
      "Requirement already satisfied: scikit-learn in /root/miniconda3/lib/python3.8/site-packages (from torch-geometric) (1.3.1)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /root/miniconda3/lib/python3.8/site-packages (from jinja2->torch-geometric) (2.1.2)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /root/miniconda3/lib/python3.8/site-packages (from requests->torch-geometric) (1.26.6)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /root/miniconda3/lib/python3.8/site-packages (from requests->torch-geometric) (3.1.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /root/miniconda3/lib/python3.8/site-packages (from requests->torch-geometric) (2021.5.30)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /root/miniconda3/lib/python3.8/site-packages (from requests->torch-geometric) (2.10)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /root/miniconda3/lib/python3.8/site-packages (from scikit-learn->torch-geometric) (1.3.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /root/miniconda3/lib/python3.8/site-packages (from scikit-learn->torch-geometric) (3.2.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install torch-geometric torch-scatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "papermill": {
     "duration": 4.819384,
     "end_time": "2023-09-01T16:55:16.104784",
     "exception": false,
     "start_time": "2023-09-01T16:55:11.2854",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "from tqdm import tqdm \n",
    "\n",
    "import sklearn,sklearn.model_selection\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import Tensor\n",
    "from torch_geometric.nn import GCNConv,SAGEConv\n",
    "from torch_geometric.datasets import Planetoid\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from timm.scheduler import CosineLRScheduler\n",
    "import matplotlib.pyplot as plt\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "papermill": {
     "duration": 0.020227,
     "end_time": "2023-09-01T16:55:16.152594",
     "exception": false,
     "start_time": "2023-09-01T16:55:16.132367",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_df(directory):\n",
    "    splits = [\"train\", \"valid\", \"test\"]\n",
    "    dfs = dict()\n",
    "    \n",
    "    for split in splits:\n",
    "        path = os.path.join(directory, split)\n",
    "        files = os.listdir(path)\n",
    "        list_df = []\n",
    "        \n",
    "        for file in files:\n",
    "            d = dict(np.load(os.path.join(path,file)))\n",
    "            d['file'] = file\n",
    "            list_df.append(d)\n",
    "        dfs[split] = pd.DataFrame.from_dict(list_df)\n",
    "    return dfs\n",
    "tile_xla = load_df(\"/root/autodl-tmp/npz_all/npz/tile/xla/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.008874,
     "end_time": "2023-09-01T16:56:21.968592",
     "exception": false,
     "start_time": "2023-09-01T16:56:21.959718",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Define Dataset and Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "papermill": {
     "duration": 0.020329,
     "end_time": "2023-09-01T16:56:21.997734",
     "exception": false,
     "start_time": "2023-09-01T16:56:21.977405",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class TileDataset(Dataset):\n",
    "    def __init__(self, df):\n",
    "        self.df = df\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        row = self.df.iloc[idx]\n",
    "        config_feat = torch.tensor(row['config_feat'].astype(np.float32))\n",
    "        node_feat = torch.tensor(row['node_feat'].astype(np.float32))\n",
    "        node_opcode = torch.tensor(row['node_opcode'].astype(np.int64))\n",
    "        edge_index = torch.tensor(np.swapaxes(row['edge_index'],0,1).astype(np.int64))\n",
    "        target = (row['config_runtime']/(row['config_runtime_normalizers']+1e-5)).astype(np.float32) #/row['config_runtime_normalizers']\n",
    "        # minmax scale the target, we only care about order\n",
    "        target = (target-np.mean(target))/(np.std(target)+1e-5)\n",
    "\n",
    "#         target = (target-np.mean(target))/(np.std(target))\n",
    "        target = torch.tensor(target)\n",
    "        return config_feat,node_feat,node_opcode,edge_index,target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleModel(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels, graph_in, graph_out, hidden_dim,dropout=0.0):\n",
    "        super().__init__()\n",
    "        op_embedding_dim = 4 # I choose 4-dimensional embedding\n",
    "        self.embedding = torch.nn.Embedding(120, #120 different op-codes\n",
    "                                            op_embedding_dim,\n",
    "                                           )\n",
    "        assert len(hidden_channels)>0\n",
    "        \n",
    "        self.linear = nn.Linear(op_embedding_dim+140,graph_in)\n",
    "        in_channels=graph_in\n",
    "        self.convs = torch.nn.ModuleList()\n",
    "        last_dim = hidden_channels[0]\n",
    "        conv = SAGEConv\n",
    "        self.convs.append(conv(in_channels, hidden_channels[0]))\n",
    "        for i in range(len(hidden_channels)-1):\n",
    "            self.convs.append(conv(hidden_channels[i], hidden_channels[i+1]))\n",
    "            last_dim = hidden_channels[i+1]\n",
    "        self.convs.append(conv(last_dim, graph_out))\n",
    "        \n",
    "        \n",
    "        \n",
    "        self.dense = torch.nn.Sequential(nn.Linear(graph_out*2+24, hidden_dim),\n",
    "                                         nn.Dropout(p=dropout),\n",
    "                                         nn.ReLU(),\n",
    "                                         nn.Linear(hidden_dim, hidden_dim),\n",
    "                                         nn.Dropout(p=dropout),\n",
    "                                         nn.ReLU(),\n",
    "                                         nn.Linear(hidden_dim, 1),\n",
    "                                        )\n",
    "#         self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, x_cfg: Tensor,x_feat: Tensor, x_op: Tensor, edge_index: Tensor) -> Tensor:\n",
    "        \n",
    "        #get graph features\n",
    "        x = torch.concat([x_feat,self.embedding(x_op)],dim = 1)\n",
    "        x = self.linear(x)\n",
    "        #pass though conv layers\n",
    "        for conv in self.convs:\n",
    "            x = conv(x, edge_index).relu()\n",
    "        # get 1d graph embedding using average pooling\n",
    "        x_mean = x.mean(0)\n",
    "        x_max = x.max(0).values\n",
    "        \n",
    "        #put graph data into config data\n",
    "        x = torch.concat([x_cfg,x_max.repeat((len(x_cfg),1)),x_mean.repeat((len(x_cfg),1))],axis=1)\n",
    "        #put into dense nn\n",
    "        x = torch.flatten(self.dense(x))\n",
    "        x = (x-torch.mean(x))/(torch.std(x)+1e-5)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 0.008439,
     "end_time": "2023-09-01T16:56:22.088164",
     "exception": false,
     "start_time": "2023-09-01T16:56:22.079725",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Train One Epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "papermill": {
     "duration": 0.021726,
     "end_time": "2023-09-01T16:56:22.11899",
     "exception": false,
     "start_time": "2023-09-01T16:56:22.097264",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = pd.concat((tile_xla[\"train\"],tile_xla[\"valid\"]),axis=0).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "papermill": {
     "duration": 9138.052878,
     "end_time": "2023-09-01T19:28:40.180685",
     "exception": false,
     "start_time": "2023-09-01T16:56:22.127807",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 0, comp_score = 0.578, mean_score = 0.244,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 1, comp_score = 0.834, mean_score = 0.676,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 2, comp_score = 0.972, mean_score = 0.329,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 3, comp_score = 0.969, mean_score = 0.413,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 4, comp_score = 0.956, mean_score = 0.644,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 5, comp_score = 0.945, mean_score = 0.683,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 6, comp_score = 0.963, mean_score = 0.584,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 7, comp_score = 0.964, mean_score = 0.535,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 8, comp_score = 0.956, mean_score = 0.633,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 0 epoch 9, comp_score = 0.959, mean_score = 0.498,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 0, comp_score = 0.547, mean_score = 0.183,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 1, comp_score = 0.824, mean_score = 0.641,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 2, comp_score = 0.959, mean_score = 0.523,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 3, comp_score = 0.958, mean_score = 0.461,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 4, comp_score = 0.945, mean_score = 0.699,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 5, comp_score = 0.617, mean_score = 0.454,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 6, comp_score = 0.950, mean_score = 0.596,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 7, comp_score = 0.316, mean_score = 0.178,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 8, comp_score = 0.945, mean_score = 0.520,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 1 epoch 9, comp_score = 0.958, mean_score = 0.427,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 0, comp_score = 0.505, mean_score = 0.220,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 1, comp_score = 0.860, mean_score = 0.706,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 2, comp_score = 0.954, mean_score = 0.620,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 3, comp_score = 0.955, mean_score = 0.492,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 4, comp_score = 0.938, mean_score = 0.673,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 5, comp_score = 0.954, mean_score = 0.545,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 6, comp_score = 0.957, mean_score = 0.501,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 7, comp_score = 0.720, mean_score = 0.512,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 8, comp_score = 0.960, mean_score = 0.518,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 2 epoch 9, comp_score = 0.954, mean_score = 0.547,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 0, comp_score = 0.585, mean_score = 0.188,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 1, comp_score = 0.892, mean_score = 0.780,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 2, comp_score = 0.956, mean_score = 0.586,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 3, comp_score = 0.962, mean_score = 0.454,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 4, comp_score = 0.962, mean_score = 0.435,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 5, comp_score = 0.951, mean_score = 0.471,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 6, comp_score = 0.960, mean_score = 0.483,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 7, comp_score = 0.958, mean_score = 0.533,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 8, comp_score = 0.956, mean_score = 0.596,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 3 epoch 9, comp_score = 0.935, mean_score = 0.723,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 0, comp_score = 0.555, mean_score = 0.203,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 1, comp_score = 0.853, mean_score = 0.718,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 2, comp_score = 0.967, mean_score = 0.434,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 3, comp_score = 0.946, mean_score = 0.743,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 4, comp_score = 0.956, mean_score = 0.663,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 5, comp_score = 0.960, mean_score = 0.519,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                          \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 6, comp_score = 0.952, mean_score = 0.692,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 7, comp_score = 0.949, mean_score = 0.694,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 8, comp_score = 0.958, mean_score = 0.680,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                           \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fold 4 epoch 9, comp_score = 0.950, mean_score = 0.638,\n",
      "comp_score = 0.9175187261750045, mean_score = 0.7223694792745212,\n"
     ]
    }
   ],
   "source": [
    "kfold = sklearn.model_selection.KFold(n_splits=5,shuffle=True,random_state=0)\n",
    "score_means = []\n",
    "score_maxs = []\n",
    "for fold,(tr_idx,va_idx) in enumerate(kfold.split(df)):\n",
    "    model = SimpleModel(hidden_channels = [32,48,64,84],graph_in = 64,graph_out = 64,hidden_dim=128,dropout = 0.2).to(device)\n",
    "    train_dataset = TileDataset(df.iloc[tr_idx])\n",
    "    val_dataset = TileDataset(df.iloc[va_idx])\n",
    "    criterion = torch.nn.MSELoss()\n",
    "    steps = len(train_dataset)*20\n",
    "    warmup_steps = int(steps*0.2)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4,weight_decay = 1e-4)\n",
    "    scheduler = CosineLRScheduler(optimizer,t_initial= steps,warmup_t=warmup_steps, warmup_lr_init=1e-6,lr_min=2e-8,)\n",
    "    \n",
    "    def score_tile_mean(predictions, df):\n",
    "        score = 0\n",
    "        for i in range(len(df)):\n",
    "            predbest = np.mean(df.iloc[i]['config_runtime'][predictions[i]])\n",
    "            best = np.mean(np.sort(df.iloc[i]['config_runtime'])[:50])\n",
    "            score += 2-predbest/best\n",
    "        score /= len(df)\n",
    "        return score\n",
    "    def score_tile_max(predictions, df):\n",
    "        score = 0\n",
    "        for i in range(len(df)):\n",
    "            predbest = np.min(df.iloc[i]['config_runtime'][predictions[i][:5]])\n",
    "            best = np.min(df.iloc[i]['config_runtime'])\n",
    "    #         print(best,predbest)\n",
    "            score += 2 - predbest/best\n",
    "        score /= len(df)\n",
    "        return score\n",
    "\n",
    "    best_score = 0\n",
    "    best_score_max = 0\n",
    "    for epoch in range(10):\n",
    "        model.train()\n",
    "        pbar = tqdm(range(len(train_dataset)),leave=False)\n",
    "        loss_sum = 0\n",
    "        n = 0\n",
    "        for i in pbar:\n",
    "            cfg_ft,nd_ft,nd_op,ind,target = train_dataset[i]\n",
    "            cfg_ft,nd_ft,nd_op,ind,target = cfg_ft.to(device),nd_ft.to(device),nd_op.to(device),ind.to(device),target.to(device)\n",
    "\n",
    "            out = model(cfg_ft,nd_ft,nd_op,ind)\n",
    "            loss = criterion(out, target)\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1e-2)\n",
    "            scheduler.step(i+len(train_dataset)*epoch)\n",
    "            optimizer.step()\n",
    "            loss_sum+=loss.item()\n",
    "            n+=1\n",
    "            pbar.set_description(f'running loss: {(loss_sum/n):.2f},current loss: {(loss.item()):.2f}')\n",
    "        pbar.close()\n",
    "        model.eval()\n",
    "\n",
    "        tile_xla_predictions = []\n",
    "        pbar = tqdm(range(len(val_dataset)),leave=False)\n",
    "        for i in pbar:\n",
    "            cfg_ft,nd_ft,nd_op,ind,target = val_dataset[i]\n",
    "            cfg_ft,nd_ft,nd_op,ind,target = cfg_ft.to(device),nd_ft.to(device),nd_op.to(device),ind.to(device),target.to(device)\n",
    "\n",
    "            out = model(cfg_ft,nd_ft,nd_op,ind)\n",
    "            # tile_xla_predictions.append(np.argsort(out.detach().numpy())[:50])\n",
    "            tile_xla_predictions.append(np.argsort(out.detach().cpu().numpy())[:50])\n",
    "        pbar.close()\n",
    "        score_mean = score_tile_mean(tile_xla_predictions, val_dataset.df)\n",
    "        score_max = score_tile_max(tile_xla_predictions, val_dataset.df)\n",
    "        print(f'fold {fold} epoch {epoch}, comp_score = {score_max:.3f}, mean_score = {score_mean:.3f},')\n",
    "        if score_mean>best_score:\n",
    "            best_score = score_mean\n",
    "            best_score_max = score_max\n",
    "            torch.save(model.state_dict(), f'best_model_{fold}.pth')\n",
    "    score_means.append(best_score)\n",
    "    score_maxs.append(best_score_max)\n",
    "print(f'comp_score = {np.mean(score_maxs)}, mean_score = {np.mean(score_means)},')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 20.903162,
     "end_time": "2023-09-01T19:29:21.465429",
     "exception": false,
     "start_time": "2023-09-01T19:29:00.562267",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Evaluate on Validation Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 20.540632,
     "end_time": "2023-09-01T19:30:02.340577",
     "exception": false,
     "start_time": "2023-09-01T19:29:41.799945",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "**0.31 is not bad considering that this model only trained on 1 epoch and is not on a ranking loss!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "papermill": {
     "duration": 20.552961,
     "end_time": "2023-09-01T19:30:43.664106",
     "exception": false,
     "start_time": "2023-09-01T19:30:23.111145",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Predict and Submit (only tile:xla predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "papermill": {
     "duration": 42.864288,
     "end_time": "2023-09-01T19:31:46.765649",
     "exception": false,
     "start_time": "2023-09-01T19:31:03.901361",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 844/844 [00:02<00:00, 295.00it/s]\n",
      "100%|██████████| 844/844 [00:02<00:00, 324.86it/s]\n",
      "100%|██████████| 844/844 [00:02<00:00, 294.45it/s]\n",
      "100%|██████████| 844/844 [00:02<00:00, 318.67it/s]\n",
      "100%|██████████| 844/844 [00:03<00:00, 280.40it/s]\n"
     ]
    }
   ],
   "source": [
    "dataset = TileDataset(tile_xla[\"test\"])\n",
    "tile_xla_predictions = [[] for i in range(len(dataset))]\n",
    "for fold in range(5):\n",
    "    model.load_state_dict(torch.load(f'/root/autodl-tmp/best_model_{fold}.pth'))\n",
    "    model.eval()\n",
    "    pbar = tqdm(range(len(dataset)))\n",
    "    for i in pbar:\n",
    "        cfg_ft,nd_ft,nd_op,ind,target = dataset[i]\n",
    "        cfg_ft,nd_ft,nd_op,ind,target = cfg_ft.to(device),nd_ft.to(device),nd_op.to(device),ind.to(device),target.to(device)\n",
    "\n",
    "        out = model(cfg_ft,nd_ft,nd_op,ind)\n",
    "        tile_xla_predictions[i].append(out.detach().cpu().numpy())\n",
    "tile_xla_predictions = [np.argsort(np.mean(pred,axis=0))[:5] for pred in tile_xla_predictions]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "papermill": {
     "duration": 20.880172,
     "end_time": "2023-09-01T19:32:28.307392",
     "exception": false,
     "start_time": "2023-09-01T19:32:07.42722",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ID</th>\n",
       "      <th>TopConfigs</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>tile:xla:d6f5f54247bd1e58a10b9e7062c636ab</td>\n",
       "      <td>0;22;21;20;19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>tile:xla:e3a655daa38e34ec240df959b650ac16</td>\n",
       "      <td>1041;1301;124;207;711</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>tile:xla:f8c2c1a1098b2a361c26df668b286c87</td>\n",
       "      <td>41;116;101;202;44</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>tile:xla:4dd1716853ed46ee4e7d09ede1732de8</td>\n",
       "      <td>4556;9142;1045;3026;4859</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>tile:xla:d0a69155b6340748c36724e4bfc34be3</td>\n",
       "      <td>655;460;692;215;650</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>889</th>\n",
       "      <td>layout:nlp:random:60880ed76de53f4d7a1b960b24f2...</td>\n",
       "      <td>0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>890</th>\n",
       "      <td>layout:nlp:random:23559853d9702baaaacbb0c83fd3...</td>\n",
       "      <td>0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>891</th>\n",
       "      <td>layout:nlp:random:f6c146fc5cf10be4f3accbaca989...</td>\n",
       "      <td>0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>892</th>\n",
       "      <td>layout:nlp:random:32531d07a084b319dce484f53a4c...</td>\n",
       "      <td>0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>893</th>\n",
       "      <td>layout:nlp:random:3a0c5517a87df8d82fd637b83298...</td>\n",
       "      <td>0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>894 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    ID  \\\n",
       "0            tile:xla:d6f5f54247bd1e58a10b9e7062c636ab   \n",
       "1            tile:xla:e3a655daa38e34ec240df959b650ac16   \n",
       "2            tile:xla:f8c2c1a1098b2a361c26df668b286c87   \n",
       "3            tile:xla:4dd1716853ed46ee4e7d09ede1732de8   \n",
       "4            tile:xla:d0a69155b6340748c36724e4bfc34be3   \n",
       "..                                                 ...   \n",
       "889  layout:nlp:random:60880ed76de53f4d7a1b960b24f2...   \n",
       "890  layout:nlp:random:23559853d9702baaaacbb0c83fd3...   \n",
       "891  layout:nlp:random:f6c146fc5cf10be4f3accbaca989...   \n",
       "892  layout:nlp:random:32531d07a084b319dce484f53a4c...   \n",
       "893  layout:nlp:random:3a0c5517a87df8d82fd637b83298...   \n",
       "\n",
       "                                            TopConfigs  \n",
       "0                                        0;22;21;20;19  \n",
       "1                                1041;1301;124;207;711  \n",
       "2                                    41;116;101;202;44  \n",
       "3                             4556;9142;1045;3026;4859  \n",
       "4                                  655;460;692;215;650  \n",
       "..                                                 ...  \n",
       "889  0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...  \n",
       "890  0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...  \n",
       "891  0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...  \n",
       "892  0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...  \n",
       "893  0;1;2;3;4;5;6;7;8;9;10;11;12;13;14;15;16;17;18...  \n",
       "\n",
       "[894 rows x 2 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub = pd.read_csv('/root/autodl-tmp/sample_submission.csv')\n",
    "for i,filename in enumerate(tile_xla[\"test\"]['file'].values):\n",
    "    id = 'tile:xla:' +filename[:-4]\n",
    "    sub.loc[sub.ID == id,'TopConfigs'] = ';'.join(tile_xla_predictions[i].astype(str))\n",
    "sub.to_csv('submission.csv',index=False)\n",
    "sub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.10"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 9735.485807,
   "end_time": "2023-09-01T19:32:52.176348",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2023-09-01T16:50:36.690541",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}