{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%capture\n!pip install /kaggle/input/torch-geometric/torch_geometric-2.3.1-py3-none-any.whl /kaggle/input/torch-geometric/torch_scatter-2.1.1-cp310-cp310-linux_x86_64.whl","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:39:23.516150Z","iopub.execute_input":"2023-11-10T19:39:23.516534Z","iopub.status.idle":"2023-11-10T19:39:40.573196Z","shell.execute_reply.started":"2023-11-10T19:39:23.516501Z","shell.execute_reply":"2023-11-10T19:39:40.571803Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor\n\nimport numpy as np\nimport pandas as pd\nimport os\nfrom tqdm import tqdm \n\nimport sklearn,sklearn.model_selection\nimport torch\nfrom torch import nn\nfrom torch import Tensor\nfrom torch_geometric.nn import GCNConv, SAGEConv, MLP\nfrom torch_geometric.datasets import Planetoid\nfrom torch_geometric.loader import DataLoader\nfrom torch_geometric.data import Data\nfrom torch.utils.data import Dataset \nfrom torch.optim.lr_scheduler import OneCycleLR\nimport torchmetrics as tm\n\nimport matplotlib.pyplot as plt\nimport warnings\n\nwarnings.filterwarnings(\"ignore\", category=UserWarning)\ndevice = 'cuda:0' if torch.cuda.is_available() else 'cpu'\ntorch.backends.cudnn.deterministic = True","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:39:40.575957Z","iopub.execute_input":"2023-11-10T19:39:40.576427Z","iopub.status.idle":"2023-11-10T19:39:49.631985Z","shell.execute_reply.started":"2023-11-10T19:39:40.576381Z","shell.execute_reply":"2023-11-10T19:39:49.630851Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TileDataset(Dataset):\n    def __init__(self, directory: str, split: str = \"test\", tile: bool = False):\n        self.file_path_list = self.__get_path(directory, split)\n        self.tile = tile\n        self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n        self.load_all_data()\n    \n    def __get_path(self, directory: str, split: str = \"test\"):\n        \n        path = os.path.join(directory, split)\n        files = os.listdir(path)\n        list_path = [os.path.join(path,file) for file in files]\n        \n        return list_path\n        \n    def __len__(self):\n        return len(self.file_path_list)\n    \n\n    def __tile(self, file):\n        config_feat = torch.tensor(file.get(\"config_feat\").astype(np.float32)) * 100\n        target = file['config_runtime'].astype(np.float32) / file['config_runtime_normalizers'].astype(np.float32)\n        \n        return config_feat, target\n    \n    def  __layout(self, file):\n        config_feat = torch.tensor(file.get(\"node_config_feat\").astype(np.float32)).mean(dim = 1)\n        target = file['config_runtime'].astype(np.float32)\n        \n        return config_feat, target\n    \n    def load_all_data(self):\n        \n        self.all_data = []\n        \n        for file in tqdm(self.file_path_list):\n            file = dict(np.load(file))\n            node_feat = torch.tensor(file['node_feat'].astype(np.float32))\n            node_opcode = torch.tensor(file['node_opcode'].astype(np.int64))\n            edge_index = torch.tensor(np.swapaxes(file['edge_index'],0,1).astype(np.int64))\n\n            config_feat, target_raw = self.__tile(file) if self.tile else self.__layout(file)\n\n            # minmax scale the target, we only care about order\n            target = (target_raw-(min(target_raw)-1.0e-7))/(max(target_raw) -(min(target_raw)-1.0e-7))\n            target = torch.tensor(target)\n\n            self.all_data.append(Data(\n                node_feat=node_feat, \n                node_opcode=node_opcode, \n                edge_index=edge_index,\n                config_feat=config_feat, \n                target=target, \n                target_raw = torch.tensor(target_raw)\n            ).to(device))\n            \n\n    def __getitem__(self, idx):\n        return self.all_data[idx] ","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:39:49.633780Z","iopub.execute_input":"2023-11-10T19:39:49.634351Z","iopub.status.idle":"2023-11-10T19:39:49.651478Z","shell.execute_reply.started":"2023-11-10T19:39:49.634321Z","shell.execute_reply":"2023-11-10T19:39:49.650313Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Evaluation Metrics","metadata":{}},{"cell_type":"code","source":"class EvaluationMetric(tm.Metric):\n    \n    higher_is_better = True\n    \n    def __init__(self, k:int=5) -> None:\n        super().__init__()\n        self.add_state(\"runtimes\", default=[], dist_reduce_fx=None)\n        self.k = k\n        \n    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:\n        \"\"\"\n        Update the metric state\n        Args:\n            preds: Tensor of shape (bs, seq_len) with the predicted runtimes orders\n            target: Tensor of shape (bs, seq_len) with the target runtimes\n        \"\"\"\n        best_runtimes = target.min()\n        best_predicted_runtimes = target[preds.argsort()[:5]].min()\n        self.runtimes.append(best_predicted_runtimes/ best_runtimes)\n        \n    def compute(self) -> torch.Tensor:\n        return (2-torch.tensor(self.runtimes)).mean()","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:39:49.653985Z","iopub.execute_input":"2023-11-10T19:39:49.654402Z","iopub.status.idle":"2023-11-10T19:39:49.667390Z","shell.execute_reply.started":"2023-11-10T19:39:49.654365Z","shell.execute_reply":"2023-11-10T19:39:49.666606Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"# function from here https://github.com/allegro/allRank\ndef listMLE(y_pred, y_true, eps=1e-10, padded_value_indicator=-1):\n    \"\"\"\n    ListMLE loss introduced in \"Listwise Approach to Learning to Rank - Theory and Algorithm\".\n    :param y_pred: predictions from the model, shape [batch_size, slate_length]\n    :param y_true: ground truth labels, shape [batch_size, slate_length]\n    :param eps: epsilon value, used for numerical stability\n    :param padded_value_indicator: an indicator of the y_true index containing a padded item, e.g. -1\n    :return: loss value, a torch.Tensor\n    \"\"\"\n    \n    # shuffle for randomised tie resolution\n    random_indices = torch.randperm(y_pred.shape[-1])\n    y_pred_shuffled = y_pred[random_indices]\n    y_true_shuffled = y_true[random_indices]\n\n    y_true_sorted, indices = y_true_shuffled.sort(descending=True, dim=-1)\n\n    mask = y_true_sorted == padded_value_indicator\n\n    preds_sorted_by_true = torch.gather(y_pred_shuffled, dim=0, index=indices)\n    preds_sorted_by_true[mask] = float(\"-inf\")\n\n    max_pred_values, _ = preds_sorted_by_true.max(dim=0, keepdim=True)\n\n    preds_sorted_by_true_minus_max = preds_sorted_by_true - max_pred_values\n\n    cumsums = torch.cumsum(preds_sorted_by_true_minus_max.exp().flip(dims=[0]), dim=0).flip(dims=[0])\n\n    observation_loss = torch.log(cumsums + eps) - preds_sorted_by_true_minus_max\n\n    observation_loss[mask] = 0.0\n    return torch.mean(torch.sum(observation_loss, dim=0))","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:39:49.668464Z","iopub.execute_input":"2023-11-10T19:39:49.669401Z","iopub.status.idle":"2023-11-10T19:39:49.685448Z","shell.execute_reply.started":"2023-11-10T19:39:49.669372Z","shell.execute_reply":"2023-11-10T19:39:49.684391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LightningWrapper(pl.LightningModule):\n    def __init__(self, model:nn.Module):\n        super().__init__()\n        self.model = model\n        self.val_eval = EvaluationMetric()\n        self.train_eval = EvaluationMetric()\n        self.loss_function = listMLE\n        self.start = False\n        \n    def forward(self, x):\n        return self.model(x)\n\n    def training_step(self, batch, batch_idx):\n        outputs = self.model(batch)\n        loss = self.loss_function(outputs, batch.target)\n        self.train_eval.update(outputs, batch.target_raw)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        outputs = self.model(batch)\n        loss = self.loss_function(outputs, batch.target)\n        self.log(\"val_loss\", loss, prog_bar=True)\n        if not self.start:\n            self.log(\"val_accuracy\", 0)\n            self.start = True\n        self.val_eval.update(outputs, batch.target_raw)\n        return loss\n    \n    def on_validation_epoch_end(self) -> None:\n        topk = self.val_eval.compute()\n        topk_train = self.train_eval.compute()\n        \n        self.log(\"train_accuracy\", topk_train, prog_bar=True)\n        self.train_eval.reset()\n        \n        self.log(\"val_accuracy\", topk, prog_bar=True)\n        self.val_eval.reset()\n        return super().on_validation_epoch_end()\n\n    def predict_step(self, batch, batch_idx):\n        return self.model(batch)\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(\n            model.parameters(),\n            lr=0.001,\n        )\n        \n        lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=4, min_lr=0.000001)\n        \n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": lr_scheduler,\n                \"interval\": \"epoch\",\n                \"frequency\": 1,\n                \"monitor\": \"val_accuracy\",\n                \"name\": \"LR\"\n            },\n        }\n        return optimizer","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:42:12.205618Z","iopub.execute_input":"2023-11-10T19:42:12.206049Z","iopub.status.idle":"2023-11-10T19:42:12.223053Z","shell.execute_reply.started":"2023-11-10T19:42:12.206016Z","shell.execute_reply":"2023-11-10T19:42:12.221060Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TileModel(torch.nn.Module):\n    def __init__(self, hidden_channels, graph_feats, hidden_dim, dense_input):\n        super().__init__()\n        op_embedding_dim = 128 \n        \n        self.embedding = torch.nn.Embedding(\n            120, # op code\n            op_embedding_dim,\n        )\n        \n        assert len(hidden_channels)>0\n        in_channels = op_embedding_dim+140\n        self.convs = torch.nn.ModuleList()\n        last_dim = hidden_channels[0]\n        self.convs.append(GCNConv(in_channels, hidden_channels[0]))\n        for i in range(len(hidden_channels)-1):\n            self.convs.append(GCNConv(hidden_channels[i], hidden_channels[i+1]))\n            last_dim = hidden_channels[i+1]\n        self.convs.append(GCNConv(last_dim, graph_feats))\n        \n        self.dense = MLP(channel_list=[dense_input, 64, 32, 16, 1], dropout=0.1, norm=None)\n        \n        \n    \n    def forward(self, data: Data) -> Tensor:\n        \n        x = torch.concat(\n            [\n                data.node_feat, \n                self.embedding(data.node_opcode)\n            ],\n            dim = 1\n        )\n        \n        #pass though conv layers\n        for conv in self.convs:\n            x = conv(x, data.edge_index).relu()\n            \n        # get 1d graph embedding using average pooling\n        x_graph = torch.mean(x,0)\n        \n        #put graph data into config data\n        x = torch.concat([data.config_feat, x_graph.repeat((len(data.config_feat),1))],axis=1) #torch.Size([10528, 225])\n\n        x = torch.flatten(self.dense(x))\n        \n        x = (x - x.mean()) / (x.std() + 1e-8)\n        \n        return x","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:42:12.672203Z","iopub.execute_input":"2023-11-10T19:42:12.672960Z","iopub.status.idle":"2023-11-10T19:42:12.685271Z","shell.execute_reply.started":"2023-11-10T19:42:12.672920Z","shell.execute_reply":"2023-11-10T19:42:12.683755Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = TileModel(hidden_channels = [32,64,64,32],graph_feats = 128,hidden_dim=128, dense_input = 152).to(device)\nmodel = LightningWrapper(model)","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:43:03.674616Z","iopub.execute_input":"2023-11-10T19:43:03.675048Z","iopub.status.idle":"2023-11-10T19:43:03.687817Z","shell.execute_reply.started":"2023-11-10T19:43:03.675014Z","shell.execute_reply":"2023-11-10T19:43:03.686861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset_tile_xla = TileDataset(\"/kaggle/input/predict-ai-model-runtime/npz_all/npz/tile/xla/\", split = \"train\", tile = True)\nvalidation_dataset_tile_xla = TileDataset(\"/kaggle/input/predict-ai-model-runtime/npz_all/npz/tile/xla/\", split = \"valid\", tile = True)\ntest_dataset_tile_xla = TileDataset(\"/kaggle/input/predict-ai-model-runtime/npz_all/npz/tile/xla/\", split = \"test\", tile = True)","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:39:49.869587Z","iopub.execute_input":"2023-11-10T19:39:49.870338Z","iopub.status.idle":"2023-11-10T19:41:12.934450Z","shell.execute_reply.started":"2023-11-10T19:39:49.870306Z","shell.execute_reply":"2023-11-10T19:41:12.933429Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NUM_EPOCH = 10\n\ncheckpoint_callback = ModelCheckpoint(\n    monitor='val_accuracy',\n    dirpath='tile_xla',\n    filename='model-{epoch:02d}-val_acc-{val_accuracy:.3f}',\n    auto_insert_metric_name=False,\n    save_on_train_epoch_end = False,  # run checkpointing after validation\n    save_top_k=5,\n    mode = 'max',\n    verbose = True,\n)\n\nlr_monitor = LearningRateMonitor(logging_interval='epoch')\n\ntrainer_config = dict(\n    max_epochs= NUM_EPOCH,\n    precision= 32,\n    accumulate_grad_batches= 1,\n    check_val_every_n_epoch= 1,\n    callbacks = [checkpoint_callback, lr_monitor]\n)","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:43:06.165134Z","iopub.execute_input":"2023-11-10T19:43:06.165538Z","iopub.status.idle":"2023-11-10T19:43:06.173384Z","shell.execute_reply.started":"2023-11-10T19:43:06.165508Z","shell.execute_reply":"2023-11-10T19:43:06.172270Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer = pl.Trainer(**trainer_config,)\ntrainer.fit(model, train_dataset_tile_xla, validation_dataset_tile_xla)","metadata":{"execution":{"iopub.status.busy":"2023-11-10T19:43:06.739239Z","iopub.execute_input":"2023-11-10T19:43:06.740361Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df = pd.read_csv('../input/predict-ai-model-runtime/sample_submission.csv')\n\ntile_xla_predictions = trainer.predict(dataloaders = test_dataset_tile_xla, ckpt_path='best')\n        \ntile_xla_predictions = [torch.argsort(pred) for pred in tile_xla_predictions]\ntile_xla_predictions[0]\nsub = submission_df\n\nfor i,filename in enumerate(test_dataset_tile_xla.file_path_list):\n    id = 'tile:xla:' + filename[:-4].split(\"/\")[-1]\n    sub.loc[sub.ID == id,'TopConfigs'] = ';'.join(tile_xla_predictions[i][0:5].numpy().astype(str))\n    \nsub.to_csv('submission.csv',index=False)\nsub","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}