Phase:02-TRANSFORMER.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from timm.utils import AverageMeter\n",
    "from timm.models import *\n",
    "from timm.loss import SoftTargetCrossEntropy\n",
    "import timm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.model_zoo as model_zoo\n",
    "\n",
    "from torch_optimizer import Ranger\n",
    "import ttach as tta\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import glob\n",
    "import os\n",
    "import copy\n",
    "import random\n",
    "import math\n",
    "import pandas as pd\n",
    "\n",
    "from tqdm import tqdm_notebook\n",
    "\n",
    "from albumentations import (\n",
    "    HorizontalFlip, VerticalFlip, ShiftScaleRotate, Transpose, HueSaturationValue, MotionBlur, \n",
    "    RandomResizedCrop, RandomBrightnessContrast, OneOf, Compose, Normalize, Cutout, CoarseDropout,\n",
    "    CenterCrop, Resize, RandomCrop, CenterCrop\n",
    ")\n",
    "\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "from torch.nn import TransformerEncoder, TransformerEncoderLayer\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('new_split.csv')\n",
    "ratio = 768/1280"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class config:\n",
    "    BASE_LR = 5e-6\n",
    "    NUM_CLASSES = 1\n",
    "    BIN_SIZE = 10\n",
    "    NUM_EPOCHS = 10\n",
    "    MODEL_NAMES = [\"mixnet_xl\", \"dpn68b\", \"mixnet_l\", \"efficientnet_es\", \"mobilenetv3_large_100\", \"seresnext26t_32x4d\"]\n",
    "    INPUT_FEATURES = [np.load(f'class_embeddings/{i}_FOLD{0}_{df.image_id.iloc[0]}.npy').shape[0] for i in MODEL_NAMES]\n",
    "    D_MODELS = [int(((i*ratio)//2)*2) for i in INPUT_FEATURES]\n",
    "    MODEL_NAME = \"test\"\n",
    "    OPTIMIZER_NAME = \"Ranger\"\n",
    "    FILE_PREFIX = \"transformer-N=64\"\n",
    "    INPUT_FEATURE = 1280\n",
    "    SEED = 43\n",
    "    IMG_SIZE = 320\n",
    "    MEAN = [0.485, 0.456, 0.406]\n",
    "    STD = [0.229, 0.224, 0.225]\n",
    "    BATCH_SIZES = [32]*len(MODEL_NAMES)\n",
    "    BATCH_SIZE = 32\n",
    "    N = 64\n",
    "    WORKERS = 16\n",
    "    FOLD = [0,1,2,3,4]\n",
    "    DEBUG = False\n",
    "    MODE = 2 #{0: Train; 1: Val Logits; 2: Test Logits; 3: ALL}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.D_MODELS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def seed_everything(seed):\n",
    "    random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "seed_everything(config.SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_train_transforms():\n",
    "    return Compose([\n",
    "            RandomResizedCrop(config.IMG_SIZE, config.IMG_SIZE),\n",
    "            Transpose(p=0.5),\n",
    "            HorizontalFlip(p=0.5),\n",
    "            VerticalFlip(p=0.5),\n",
    "            ShiftScaleRotate(p=0.5, rotate_limit=180),\n",
    "            Normalize(mean=config.MEAN, std=config.STD, max_pixel_value=255.0, p=1.0),\n",
    "            ToTensorV2(p=1.0),\n",
    "        ], p=1.)\n",
    "  \n",
    "        \n",
    "def get_valid_transforms(tta=False):\n",
    "    aug = []\n",
    "    if not tta: aug.extend([CenterCrop(config.IMG_SIZE, config.IMG_SIZE)])\n",
    "    aug.extend([\n",
    "            Normalize(mean=config.MEAN, std=config.STD, max_pixel_value=255.0, p=1.0),\n",
    "            ToTensorV2(p=1.0),\n",
    "        ])\n",
    "    return Compose(aug)\n",
    "\n",
    "train_transform = get_train_transforms()\n",
    "valid_transform = get_valid_transforms()\n",
    "tta_test_transform = get_valid_transforms(tta=True)\n",
    "\n",
    "tta_transforms = tta.Compose(\n",
    "    [\n",
    "        tta.FiveCrops(config.IMG_SIZE, config.IMG_SIZE)  \n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimeDatasetWIND(Dataset):\n",
    "\n",
    "    def __init__(self, df, df_train, N = config.N, D = config.INPUT_FEATURE, phase='train'):\n",
    "        self.data = df.reset_index(drop=True)\n",
    "        self.data_train = df_train.reset_index(drop=True)\n",
    "        self.label = df.wind_speed\n",
    "        self.N = N\n",
    "        self.D = D\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, index): # N, D\n",
    "            images = []\n",
    "            mask = []\n",
    "            final_images = torch.zeros(self.N, self.D)\n",
    "            final_labels = self.get_labels(index).unsqueeze(0)\n",
    "            final_mask = torch.zeros((self.N))\n",
    "            images.append(torch.Tensor(np.load(f'class_embeddings/{config.MODEL_NAME}_FOLD{fold}_{self.data.image_id.iloc[index]}.npy')).unsqueeze(0))\n",
    "            mask.append(1)\n",
    "            embedding_num = int(self.data.image_id.iloc[index].split('_')[-1])\n",
    "            embedding_ocean_id = self.data.storm_id.iloc[index]\n",
    "            index_train = self.data_train[self.data_train.image_id==self.data.image_id.iloc[index]].index.values[0]\n",
    "            for i in range(1, self.N):\n",
    "                embedding_nump = int(self.data_train.image_id.iloc[index_train-i].split('_')[-1])\n",
    "                embedding_ocean_idp = self.data_train.storm_id.iloc[index_train-i]\n",
    "                if embedding_num-embedding_nump==i and embedding_ocean_id==embedding_ocean_idp:\n",
    "                    images.append(torch.Tensor(np.load(f'class_embeddings/{config.MODEL_NAME}_FOLD{fold}_{self.data_train.image_id.iloc[index_train-i]}.npy')).unsqueeze(0))\n",
    "                    mask.append(1)\n",
    "                else:\n",
    "                    break\n",
    "            images = torch.cat(images, dim=0)\n",
    "            mask = torch.tensor(mask)\n",
    "            images = torch.flip(images, [0])\n",
    "            mask = torch.flip(mask, [0])\n",
    "            l = len(images)\n",
    "            final_images[:l] = images\n",
    "            final_mask[:l] = mask\n",
    "            if self.label is not None:\n",
    "                return final_images, final_mask, final_labels\n",
    "            else:\n",
    "                return final_images, final_mask       \n",
    "    def get_labels(self, index):\n",
    "        wind_speed = self.label.iloc[index].astype(float)/config.BIN_SIZE\n",
    "        return torch.tensor(wind_speed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimeDatasetWIND_FAST(Dataset):\n",
    "\n",
    "    def __init__(self, final_images, final_mask, final_labels):\n",
    "        self.final_images = final_images\n",
    "        self.final_mask = final_mask\n",
    "        self.final_labels = final_labels  \n",
    "    def __len__(self):\n",
    "        return len(self.final_images)\n",
    "\n",
    "    def __getitem__(self, index): # N, D\n",
    "        return self.final_images[index], self.final_mask[index], self.final_labels[index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, dropout=0.1, max_len=5000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        \n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:x.size(0), :]\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BEST input_dim=1280, d_model=1024, nhead=2, dim_feedforward=2048, nlayers=3, dropout=0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerModel(nn.Module):\n",
    "    def __init__(self, num_classes=1, input_dim=1280, d_model=768, nhead=2, dim_feedforward=2048, nlayers=3, dropout=0.1):\n",
    "        super(TransformerModel, self).__init__()\n",
    "        self.model_type = 'Transformer'\n",
    "        self.pos_encoder = PositionalEncoding(d_model, dropout)\n",
    "        encoder_layers = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)\n",
    "        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, nlayers, norm=nn.ReLU())\n",
    "        self.d_model = d_model\n",
    "        self.input_dim = input_dim\n",
    "        self.decoder = nn.Linear(d_model, num_classes)\n",
    "        self.fc = nn.Linear(self.input_dim, self.d_model)\n",
    "        self.do = nn.Dropout(0.2)\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self):\n",
    "        initrange = 0.1\n",
    "        self.decoder.bias.data.zero_()\n",
    "        self.decoder.weight.data.uniform_(-initrange, initrange)\n",
    "\n",
    "    def forward(self, embedding, mask=None):\n",
    "        embedding = self.fc(embedding)\n",
    "        \n",
    "        if mask is not None:\n",
    "            embedding = embedding*mask.unsqueeze(-1)\n",
    "        embedding = self.pos_encoder(embedding.permute(1,0,2))\n",
    "        output = self.transformer_encoder(embedding)\n",
    "        output = self.decoder(self.do(output))\n",
    "        return output.permute(1,0,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_epoch(model, loss_fn, optimizer, phase, scheduler=None, num_steps=20000):\n",
    "    running_loss = AverageMeter()\n",
    "    if phase == \"train\":\n",
    "        tk1 = tqdm_notebook(dataloaders[phase], total=min(len(dataloaders[phase]), num_steps))\n",
    "        model.train()\n",
    "        for x_var, x_mask, y_var in tk1:\n",
    "            \n",
    "            x_var = x_var.to(device=device).float()\n",
    "            x_mask = x_mask.to(device=device).float()\n",
    "            y_var = y_var.to(device=device).float()\n",
    "            idx = len(x_mask[0])-np.argmax(x_mask.cpu().numpy()[:,::-1], axis=1)-1\n",
    "            optimizer.zero_grad()\n",
    "            score = model(x_var, x_mask)[:,idx,:].diagonal().t()\n",
    "            y_var = y_var\n",
    "            loss = loss_fn(score.reshape(-1,config.NUM_CLASSES), y_var.reshape(-1,config.NUM_CLASSES))  \n",
    "            running_loss.update(loss.item(), n=config.BATCH_SIZE*config.N)\n",
    "            tk1.set_postfix(loss=running_loss.avg)\n",
    "            loss.backward()\n",
    "            \n",
    "            optimizer.step()\n",
    "            \n",
    "            num_steps-=1\n",
    "            if num_steps<0:\n",
    "                break\n",
    "            if scheduler is not None:\n",
    "                scheduler.step()\n",
    "            if config.DEBUG: break\n",
    "        return running_loss.avg\n",
    "    else:\n",
    "        tk1 = tqdm_notebook(dataloaders[phase], total=len(dataloaders[phase]))\n",
    "        model.eval()\n",
    "        y_true = np.array([])\n",
    "        y_pred = np.array([])\n",
    "        running_mse = AverageMeter()\n",
    "        with torch.no_grad():\n",
    "            for x_var, x_mask, y_var in tk1:\n",
    "                \n",
    "                bs = x_var.shape[0]\n",
    "                x_var = x_var.to(device=device).float()\n",
    "                x_mask = x_mask.to(device=device).float()\n",
    "                y_var = y_var.to(device=device).float()\n",
    "                idx = len(x_mask[0])-np.argmax(x_mask.cpu().numpy()[:,::-1], axis=1)-1\n",
    "                score = model(x_var, x_mask)[:,idx,:].diagonal().t()\n",
    "                y_var = y_var\n",
    "                loss = loss_fn(score.reshape(-1,config.NUM_CLASSES), y_var.reshape(-1,config.NUM_CLASSES)) \n",
    "            \n",
    "                running_loss.update(loss.item(), n=config.BATCH_SIZE)\n",
    "                \n",
    "                y_var = y_var.cpu().detach().numpy()*config.BIN_SIZE\n",
    "                score = score.cpu().detach().numpy()*config.BIN_SIZE\n",
    "                \n",
    "                mse = np.sum((score-y_var)**2)/(len(score))\n",
    "                running_mse.update(mse, n=len(score))\n",
    "    \n",
    "                tk1.set_postfix(loss=running_loss.avg, rmse=math.sqrt(running_mse.avg))\n",
    "                if config.DEBUG: break\n",
    "        rmse = math.sqrt(running_mse.avg)\n",
    "        return running_loss.avg, rmse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import torch.multiprocessing\n",
    "torch.multiprocessing.set_sharing_strategy('file_system')\n",
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if config.MODE in [0,3]:\n",
    "    for config.MODEL_NAME, config.BATCH_SIZE, config.INPUT_FEATURE, config.D_MODEL in zip(config.MODEL_NAMES, config.BATCH_SIZES, config.INPUT_FEATURES, config.D_MODELS):\n",
    "        for fold in config.FOLD:\n",
    "\n",
    "            X = pd.read_csv('new_split.csv')\n",
    "\n",
    "            train = X[X.fold!=fold]\n",
    "            val = X[X.fold==fold]\n",
    "\n",
    "            device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "            train_losses = []\n",
    "            valid_losses = []\n",
    "            valid_rmse = []\n",
    "\n",
    "            print('saving objects ...')\n",
    "            final_images_ = []\n",
    "            final_mask_ =[]\n",
    "            final_labels_ =[]\n",
    "            val_final_images_ = []\n",
    "            val_final_mask_ =[]\n",
    "            val_final_labels_ =[]\n",
    "            trainset = TimeDatasetWIND(train, train, D = config.INPUT_FEATURE)\n",
    "            validset = TimeDatasetWIND(val, X, D = config.INPUT_FEATURE)\n",
    "            train_loader = DataLoader(trainset, config.BATCH_SIZE, num_workers=config.WORKERS, shuffle=False, pin_memory=False, drop_last=False)\n",
    "            valid_loader = DataLoader(validset, config.BATCH_SIZE, num_workers=config.WORKERS, shuffle=False, pin_memory=False, drop_last=False)\n",
    "            tk1 = tqdm_notebook(train_loader, total=len(train_loader))\n",
    "            for final_images, final_mask, final_labels in tk1:\n",
    "                final_images_.extend(copy.deepcopy(final_images))\n",
    "                final_mask_.extend(copy.deepcopy(final_mask))\n",
    "                final_labels_.extend(copy.deepcopy(final_labels))\n",
    "                del final_images, final_mask, final_labels\n",
    "\n",
    "            tk2 = tqdm_notebook(valid_loader, total=len(valid_loader))\n",
    "            for final_images, final_mask, final_labels in tk2:\n",
    "                val_final_images_.extend(copy.deepcopy(final_images))\n",
    "                val_final_mask_.extend(copy.deepcopy(final_mask))\n",
    "                val_final_labels_.extend(copy.deepcopy(final_labels))\n",
    "                del final_images, final_mask, final_labels\n",
    "\n",
    "            print('saved ...')\n",
    "\n",
    "            trainset = TimeDatasetWIND_FAST(final_images_, final_mask_, final_labels_)\n",
    "            validset = TimeDatasetWIND_FAST(val_final_images_, val_final_mask_, val_final_labels_)\n",
    "            train_loader = DataLoader(trainset, config.BATCH_SIZE, num_workers=config.WORKERS, shuffle=True, pin_memory=True, drop_last=True)\n",
    "            valid_loader = DataLoader(validset, config.BATCH_SIZE, num_workers=config.WORKERS, shuffle=False, pin_memory=True, drop_last=True)\n",
    "            dataloaders = {\n",
    "                \"train\" : train_loader,\n",
    "                \"valid\" : valid_loader,\n",
    "            }\n",
    "\n",
    "            model = TransformerModel(input_dim=config.INPUT_FEATURE, d_model=config.D_MODEL)\n",
    "            model.to(device)\n",
    "            criterion = nn.MSELoss()\n",
    "            optimizer = globals()[config.OPTIMIZER_NAME](model.parameters(), lr=config.BASE_LR)\n",
    "#             scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr, total_steps=None, epochs=None, steps_per_epoch=None, pct_start=0.3, anneal_strategy='cos', cycle_momentum=True, base_momentum=0.85, max_momentum=0.95, div_factor=25.0, final_div_factor=10000.0, last_epoch=-1, verbose=True)\n",
    "            best_rmse = 10000\n",
    "            for epoch in range(config.NUM_EPOCHS):\n",
    "                print('Starting epoch [%d / %d]' % (epoch + 1, config.NUM_EPOCHS))\n",
    "                train_loss = run_epoch(model, criterion, optimizer, \"train\")\n",
    "                valid_loss, rmse = run_epoch(model, criterion, optimizer, \"valid\")\n",
    "\n",
    "                if rmse<best_rmse:\n",
    "                    print(\"**Saving model**\")\n",
    "                    best_rmse=rmse\n",
    "                    torch.save({\n",
    "                        \"epoch\": epoch + 1,\n",
    "                        \"state_dict\" : model.state_dict(),\n",
    "                        \"rmse\" : best_rmse,\n",
    "                        \"optim_dict\" : optimizer.state_dict(),\n",
    "                        \"config_class\" : config\n",
    "                    }, f\"models/{config.MODEL_NAME}_{config.FILE_PREFIX}_FOLD{fold}.pth\")\n",
    "\n",
    "                train_losses.append(train_loss)\n",
    "                valid_losses.append(valid_loss)\n",
    "                valid_rmse.append(rmse)\n",
    "                df_data=np.array([train_losses, valid_losses, valid_rmse]).T\n",
    "                df = pd.DataFrame(df_data, columns = ['train_losses','valid_losses','valid_rmse1'])\n",
    "                df.to_csv(f'logs/{config.MODEL_NAME}_{config.FILE_PREFIX}_FOLD{fold}.csv')\n",
    "                if config.DEBUG: break\n",
    "            del dataloaders, valid_loader, train_loader, validset, trainset, val_final_images_, val_final_mask_, val_final_labels_, final_images_, final_mask_, final_labels_\n",
    "            gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation Logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mkdir transformer_regr_val_npys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if config.MODE in [1, 3]:\n",
    "    for config.MODEL_NAME, config.BATCH_SIZE, config.INPUT_FEATURE, config.D_MODEL in zip(config.MODEL_NAMES, config.BATCH_SIZES, config.INPUT_FEATURES, config.D_MODELS):\n",
    "        for fold in config.FOLD:\n",
    "            X = pd.read_csv('new_split.csv')\n",
    "            train = X[X.fold!=fold]\n",
    "            val = X[X.fold==fold]\n",
    "            validset = TimeDatasetWIND(val, X, D = config.INPUT_FEATURE)\n",
    "            valid_loader = DataLoader(validset, config.BATCH_SIZE, num_workers=config.WORKERS, shuffle=False, pin_memory=True)\n",
    "            device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "            model = TransformerModel(input_dim=config.INPUT_FEATURE, d_model=config.D_MODEL)\n",
    "            model.load_state_dict(torch.load(f\"models/{config.MODEL_NAME}_{config.FILE_PREFIX}_FOLD{fold}.pth\")[\"state_dict\"])\n",
    "            model.to(device)\n",
    "            model.eval()\n",
    "            y_logits = []\n",
    "            tk1 = tqdm_notebook(valid_loader, total=len(valid_loader))\n",
    "            with torch.no_grad():\n",
    "                for x_var, x_mask, y_var in tk1:\n",
    "                    bs = x_var.shape[0]\n",
    "                    x_var = x_var.to(device=device).float()\n",
    "                    x_mask = x_mask.to(device=device).float()\n",
    "                    idx = len(x_mask[0])-np.argmax(x_mask.cpu().numpy()[:,::-1], axis=1)-1\n",
    "                    score = model(x_var, x_mask)[:,idx,:].diagonal().t()\n",
    "                    y_logits.extend(score.detach().cpu().numpy())\n",
    "            np.save(f'transformer_class_val_npys/{config.MODEL_NAME}_{config.FILE_PREFIX}_val_preds_fold_{fold}.npy', y_logits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mkdir transformer_regr_test_npys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_test = pd.read_csv('data/test_set_features.csv')\n",
    "x_test[\"file_name\"] = \"data/test/test/\"+x_test.image_id+\".jpg\"\n",
    "x_test['wind_speed']=0\n",
    "X = pd.read_csv('new_split.csv')\n",
    "X = pd.concat([x_test,X]).sort_values(['image_id']).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "if config.MODE in [2, 3]:\n",
    "    for config.MODEL_NAME, config.BATCH_SIZE, config.INPUT_FEATURE, config.D_MODEL in zip(config.MODEL_NAMES, config.BATCH_SIZES, config.INPUT_FEATURES, config.D_MODELS):\n",
    "        testset = TimeDatasetWIND(x_test, X, D = config.INPUT_FEATURE)\n",
    "        test_loader = DataLoader(testset, config.BATCH_SIZE, num_workers=config.WORKERS, shuffle=False, pin_memory=True)\n",
    "        for fold in config.FOLD:\n",
    "            model = TransformerModel(input_dim=config.INPUT_FEATURE, d_model=config.D_MODEL)\n",
    "            model.load_state_dict(torch.load(f\"models/{config.MODEL_NAME}_{config.FILE_PREFIX}_FOLD{fold}.pth\")[\"state_dict\"])\n",
    "            model.to(device)\n",
    "            model.eval()\n",
    "            y_logits = []\n",
    "            tk1 = tqdm_notebook(test_loader, total=len(test_loader))\n",
    "            with torch.no_grad():\n",
    "                for x_var, x_mask, y_var in tk1:\n",
    "                    bs = x_var.shape[0]\n",
    "                    x_var = x_var.to(device=device).float()\n",
    "                    x_mask = x_mask.to(device=device).float()\n",
    "                    idx = len(x_mask[0])-np.argmax(x_mask.cpu().numpy()[:,::-1], axis=1)-1\n",
    "                    score = model(x_var, x_mask)[:,idx,:].diagonal().t()\n",
    "                    y_logits.extend(score.detach().cpu().numpy())\n",
    "            np.save(f'transformer_class_test_npys/{config.MODEL_NAME}_{config.FILE_PREFIX}_preds_fold_test_{fold}.npy', y_logits)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.8"
  },
  "nbpresent": {
   "slides": {},
   "themes": {
    "default": "22135883-8335-46d0-97c3-18db6312258e",
    "theme": {
     "18781ab6-62e8-44b7-8d7f-afb05e0b8291": {
      "backgrounds": {
       "backgroundColor": {
        "background-color": "backgroundColor",
        "id": "backgroundColor"
       }
      },
      "id": "18781ab6-62e8-44b7-8d7f-afb05e0b8291",
      "palette": {
       "backgroundColor": {
        "id": "backgroundColor",
        "rgb": [
         34,
         34,
         34
        ]
       },
       "headingColor": {
        "id": "headingColor",
        "rgb": [
         238,
         238,
         238
        ]
       },
       "linkColor": {
        "id": "linkColor",
        "rgb": [
         170,
         34,
         51
        ]
       },
       "mainColor": {
        "id": "mainColor",
        "rgb": [
         238,
         238,
         238
        ]
       }
      },
      "rules": {
       "a": {
        "color": "linkColor"
       },
       "h1": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 7
       },
       "h2": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 5
       },
       "h3": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 3.75
       },
       "h4": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 3
       },
       "h5": {
        "color": "headingColor",
        "font-family": "Ubuntu"
       },
       "h6": {
        "color": "headingColor",
        "font-family": "Ubuntu"
       },
       "h7": {
        "color": "headingColor",
        "font-family": "Ubuntu"
       },
       "li": {
        "color": "mainColor",
        "font-family": "Ubuntu",
        "font-size": 5
       },
       "p": {
        "color": "mainColor",
        "font-family": "Ubuntu",
        "font-size": 5
       }
      },
      "text-base": {
       "color": "mainColor",
       "font-family": "Ubuntu",
       "font-size": 5
      }
     },
     "22135883-8335-46d0-97c3-18db6312258e": {
      "backgrounds": {
       "backgroundColor": {
        "background-color": "backgroundColor",
        "id": "backgroundColor"
       }
      },
      "id": "22135883-8335-46d0-97c3-18db6312258e",
      "palette": {
       "backgroundColor": {
        "id": "backgroundColor",
        "rgb": [
         34,
         34,
         34
        ]
       },
       "headingColor": {
        "id": "headingColor",
        "rgb": [
         238,
         238,
         238
        ]
       },
       "linkColor": {
        "id": "linkColor",
        "rgb": [
         170,
         34,
         51
        ]
       },
       "mainColor": {
        "id": "mainColor",
        "rgb": [
         238,
         238,
         238
        ]
       }
      },
      "rules": {
       "a": {
        "color": "linkColor"
       },
       "h1": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 7
       },
       "h2": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 5
       },
       "h3": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 3.75
       },
       "h4": {
        "color": "headingColor",
        "font-family": "Ubuntu",
        "font-size": 3
       },
       "h5": {
        "color": "headingColor",
        "font-family": "Ubuntu"
       },
       "h6": {
        "color": "headingColor",
        "font-family": "Ubuntu"
       },
       "h7": {
        "color": "headingColor",
        "font-family": "Ubuntu"
       },
       "li": {
        "color": "mainColor",
        "font-family": "Ubuntu",
        "font-size": 5
       },
       "p": {
        "color": "mainColor",
        "font-family": "Ubuntu",
        "font-size": 5
       }
      },
      "text-base": {
       "color": "mainColor",
       "font-family": "Ubuntu",
       "font-size": 5
      }
     },
     "fc1bfafc-cbc1-43d9-a6b2-d31b1be1fe23": {
      "backgrounds": {
       "backgroundColor": {
        "background-color": "backgroundColor",
        "id": "backgroundColor"
       }
      },
      "id": "fc1bfafc-cbc1-43d9-a6b2-d31b1be1fe23",
      "palette": {
       "backgroundColor": {
        "id": "backgroundColor",
        "rgb": [
         0,
         43,
         54
        ]
       },
       "headingColor": {
        "id": "headingColor",
        "rgb": [
         238,
         232,
         213
        ]
       },
       "linkColor": {
        "id": "linkColor",
        "rgb": [
         38,
         139,
         210
        ]
       },
       "mainColor": {
        "id": "mainColor",
        "rgb": [
         147,
         161,
         161
        ]
       }
      },
      "rules": {
       "a": {
        "color": "linkColor"
       },
       "h1": {
        "color": "headingColor",
        "font-family": "Oswald",
        "font-size": 7
       },
       "h2": {
        "color": "headingColor",
        "font-family": "Oswald",
        "font-size": 5
       },
       "h3": {
        "color": "headingColor",
        "font-family": "Oswald",
        "font-size": 3.75
       },
       "h4": {
        "color": "headingColor",
        "font-family": "Oswald",
        "font-size": 3
       },
       "h5": {
        "color": "headingColor",
        "font-family": "Oswald"
       },
       "h6": {
        "color": "headingColor",
        "font-family": "Oswald"
       },
       "h7": {
        "color": "headingColor",
        "font-family": "Oswald"
       },
       "li": {
        "color": "mainColor",
        "font-family": "Lato",
        "font-size": 5
       },
       "p": {
        "color": "mainColor",
        "font-family": "Lato",
        "font-size": 5
       }
      },
      "text-base": {
       "color": "mainColor",
       "font-family": "Lato",
       "font-size": 5
      }
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}