pytorch_lightning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "nasty-interval",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "stuck-fourth",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from glob import glob\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import cv2\n",
    "from datetime import datetime\n",
    "\n",
    "from omegaconf import OmegaConf, DictConfig\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from efficientnet_pytorch import EfficientNet\n",
    "\n",
    "from torch.optim import SGD, Adam, AdamW\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from pytorch_lightning.utilities.seed import seed_everything\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\n",
    "from adamp import AdamP\n",
    "\n",
    "from albumentations import *\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "\n",
    "now = datetime.now()\n",
    "now = now.strftime(\"%y%m%d/%H%M%S\")\n",
    "%config Completer.use_jedi = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "danish-above",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "seed: 42\n",
      "device_list:\n",
      "- 0\n",
      "data_dir: /opt/ml/input/data/train\n",
      "test_dir: /opt/ml/input/data/eval\n",
      "train:\n",
      "  n_fold: 5\n",
      "  use_fold:\n",
      "  - 0\n",
      "  n_epochs: 20\n",
      "dataset:\n",
      "  image_size:\n",
      "  - 320\n",
      "  - 256\n",
      "  batch_size: 128\n",
      "  num_workers: 0\n",
      "network:\n",
      "  model_name: efficientnet-b0\n",
      "  optimizer: AdamP\n",
      "  learning_rate: 0.0003\n",
      "run_test: true\n",
      "debug: false\n",
      "tb: true\n",
      "log_dir: /opt/ml/logs/210408/070814\n",
      "\n"
     ]
    }
   ],
   "source": [
    "s = \\\n",
    "'''\n",
    "seed: 42\n",
    "device_list: [0]\n",
    "data_dir: /opt/ml/input/data/train\n",
    "test_dir: /opt/ml/input/data/eval\n",
    "\n",
    "train:\n",
    "  n_fold: 5\n",
    "  use_fold: [0]\n",
    "  n_epochs: 20\n",
    "  \n",
    "dataset:\n",
    "  image_size: [320, 256]\n",
    "  batch_size: 128\n",
    "  num_workers: 0\n",
    "\n",
    "network:\n",
    "  model_name: efficientnet-b0\n",
    "  optimizer: AdamP  # ['SGD', 'Adam', 'AdamW']\n",
    "  learning_rate: 3e-4\n",
    "\n",
    "run_test: true\n",
    "debug: false\n",
    "tb: true\n",
    "'''\n",
    "conf = OmegaConf.create(s)\n",
    "conf.log_dir = f'/opt/ml/logs/{now}'  # 학습 설정 및 결과는 이 경로에 저장됩니다.\n",
    "\n",
    "if not os.path.isdir(conf.log_dir):\n",
    "    os.makedirs(conf.log_dir, exist_ok=True)\n",
    "OmegaConf.save(conf, os.path.join(conf.log_dir, 'config.yaml'))\n",
    "\n",
    "if conf.debug:\n",
    "    conf.train.n_epochs = 2\n",
    "    conf.dataset.batch_size = 4\n",
    "    conf.tb = False\n",
    "\n",
    "print(OmegaConf.to_yaml(conf))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "international-fields",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mask_label(image_name):\n",
    "    \"\"\"\n",
    "    이미지 파일 이름을 통해 mask label을 구합니다.\n",
    "\n",
    "    :param image_name: 학습 이미지 파일 이름\n",
    "    :return: mask label\n",
    "    \"\"\"\n",
    "    if 'incorrect_mask' in image_name:\n",
    "        return 1\n",
    "    elif 'normal' in image_name:\n",
    "        return 2\n",
    "    elif 'mask' in image_name:\n",
    "        return 0\n",
    "    else:\n",
    "        raise ValueError(f'No class for {image_name}')\n",
    "\n",
    "\n",
    "def get_gender_label(gender):\n",
    "    \"\"\"\n",
    "    gender label을 구하는 함수입니다.\n",
    "    :param gender: `male` or `female`\n",
    "    :return: gender label\n",
    "    \"\"\"\n",
    "    return 0 if gender == 'male' else 1\n",
    "\n",
    "\n",
    "def get_age_label(age):\n",
    "    \"\"\"\n",
    "    age label을 구하는 함수입니다.\n",
    "    :param age: 나이를 나타내는 int.\n",
    "    :return: age label\n",
    "    \"\"\"\n",
    "    return 0 if int(age) < 30 else 1 if int(age) < 58 else 2\n",
    "\n",
    "def convert_gender_age(gender, age):\n",
    "    \"\"\"\n",
    "    gender와 age label을 조합하여 고유한 레이블을 만듭니다.\n",
    "    이를 구하는 이유는 train/val의 성별 및 연령 분포를 맞추기 위함입니다. (by Stratified K-Fold)\n",
    "    :param gender: `male` or `female`\n",
    "    :param age: 나이를 나타내는 int.\n",
    "    :return: gender & age label을 조합한 레이블\n",
    "    \"\"\"\n",
    "    gender_label = get_gender_label(gender)\n",
    "    age_label = get_age_label(age)\n",
    "    return gender_label * 3 + age_label\n",
    "\n",
    "\n",
    "def convert_label(image_path, sep=False):\n",
    "    \"\"\"\n",
    "    이미지의 label을 구하는 함수입니다.\n",
    "    :param image_path: 이미지 경로를 나타내는 str\n",
    "    :param sep: 마스크, 성별, 연령 label을 따로 반환할건지 합쳐서 할지 나타내는 bool 인수입니다. 참일 경우 따로 반환합니다.\n",
    "    :return: 이미지의 label (int or list)\n",
    "    \"\"\"\n",
    "    image_name = image_path.split('/')[-1]\n",
    "    mask_label = get_mask_label(image_name)\n",
    "\n",
    "    profile = image_path.split('/')[-2]\n",
    "    image_id, gender, race, age = profile.split(\"_\")\n",
    "    gender_label = get_gender_label(gender)\n",
    "    age_label = get_age_label(age)\n",
    "    if sep:\n",
    "        return mask_label, gender_label, age_label\n",
    "    else:\n",
    "        return mask_label * 6 + gender_label * 3 + age_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "hydraulic-excuse",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)\n",
    "\n",
    "\n",
    "def get_transforms(need=('train', 'val'), img_size=(512, 384)):\n",
    "    \"\"\"\n",
    "    Augmentation 함수를 반환합니다.\n",
    "    \"\"\"\n",
    "    transformations = {}\n",
    "    if 'train' in need:\n",
    "        transformations['train'] = Compose([\n",
    "            CenterCrop(448, 336, p=1.0),\n",
    "            RandomResizedCrop(img_size[0], img_size[1], p=1.0),\n",
    "            HorizontalFlip(p=0.5),\n",
    "            ShiftScaleRotate(p=0.3),\n",
    "            HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.3),\n",
    "#             RandomBrightnessContrast(brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1), p=0.3),\n",
    "            Normalize(mean=mean, std=std, max_pixel_value=255.0, p=1.0),\n",
    "            CoarseDropout(p=0.3),\n",
    "            Cutout(p=0.3),\n",
    "            ToTensorV2(p=1.0),\n",
    "        ], p=1.0)\n",
    "    if 'val' in need:\n",
    "        transformations['val'] = Compose([\n",
    "            CenterCrop(448, 336, p=1.0),\n",
    "            Resize(img_size[0], img_size[1], p=1.0),\n",
    "            Normalize(mean=mean, std=std, max_pixel_value=255.0, p=1.0),\n",
    "            ToTensorV2(p=1.0),\n",
    "        ], p=1.0)\n",
    "    return transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "possible-brisbane",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_EXTENSIONS = [\n",
    "    \".jpg\", \".JPG\", \".jpeg\", \".JPEG\", \".png\",\n",
    "    \".PNG\", \".ppm\", \".PPM\", \".bmp\", \".BMP\",\n",
    "]\n",
    "\n",
    "\n",
    "def is_image_file(filepath):\n",
    "    \"\"\"\n",
    "    해당 파일이 이미지 파일인지 확인합니다.\n",
    "    \"\"\"\n",
    "    return any(filepath.endswith(extension) for extension in IMG_EXTENSIONS)\n",
    "\n",
    "\n",
    "def remove_hidden_file(filepath):\n",
    "    \"\"\"\n",
    "    `._`로 시작하는 숨김 파일일 경우 False를 반환합니다.\n",
    "    \"\"\"\n",
    "    filename = filepath.split('/')[-1]\n",
    "    return False if filename.startswith('._') else True\n",
    "\n",
    "def get_img(path):\n",
    "    \"\"\"\n",
    "    이미지를 불러옵니다.\n",
    "    \"\"\"\n",
    "    im_bgr = cv2.imread(path)\n",
    "    im_rgb = im_bgr[:, :, ::-1]\n",
    "    return im_rgb\n",
    "\n",
    "class MaskDataset(Dataset):\n",
    "    def __init__(self, image_dir, info, transform=None):\n",
    "        self.image_dir = image_dir\n",
    "        self.info = info\n",
    "        self.transform = transform\n",
    "        \n",
    "        self.image_paths = [path for name in info.path.values for path in glob(os.path.join(image_dir, name, '*'))]\n",
    "        self.image_paths = list(filter(is_image_file, self.image_paths))\n",
    "        self.image_paths = list(filter(remove_hidden_file, self.image_paths))\n",
    "        \n",
    "        self.labels = [convert_label(path, sep=False) for path in self.image_paths]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image_path = self.image_paths[idx]\n",
    "        label = self.labels[idx]\n",
    "        image = get_img(image_path)\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image=image)['image']\n",
    "        label = torch.eye(18)[label]\n",
    "        return image, label\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "legitimate-official",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TestDataset(Dataset):\n",
    "    def __init__(self, img_paths, transform):\n",
    "        self.img_paths = img_paths\n",
    "        self.transform = transform\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        image_path = self.img_paths[index]\n",
    "        image = get_img(image_path)\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image=image)['image']\n",
    "        return image\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "civil-midnight",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "# import timm\n",
    "\n",
    "class MyModel(nn.Module):\n",
    "    def __init__(self, model_name):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.model = EfficientNet.from_pretrained('efficientnet-b0', num_classes=18)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        y_hat = self.model(x)\n",
    "        return y_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "photographic-relaxation",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LabelSmoothing(nn.Module):\n",
    "    def __init__(self, smoothing=0.1, average='mean', log_softmax=True):\n",
    "        super(LabelSmoothing, self).__init__()\n",
    "        self.confidence = 1.0 - smoothing\n",
    "        self.smoothing = smoothing\n",
    "        self.average = average\n",
    "        self.log_softmax = log_softmax\n",
    "\n",
    "    def forward(self, x, target):\n",
    "        x = x.float()\n",
    "        target = target.float()\n",
    "        if self.log_softmax: x = x.log_softmax(dim=-1)\n",
    "        nll_loss = -x * target\n",
    "        nll_loss = nll_loss.sum(-1)\n",
    "        smooth_loss = -x.mean(dim=-1)\n",
    "        loss = self.confidence * nll_loss + self.smoothing * smooth_loss\n",
    "        if self.average == 'mean':\n",
    "            return loss.mean()\n",
    "        else:\n",
    "            return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "loaded-garbage",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.core.decorators import auto_move_data\n",
    "from pytorch_lightning.metrics import Accuracy\n",
    "from adamp import AdamP\n",
    "\n",
    "class LitTrainer(pl.LightningModule):\n",
    "    def __init__(self, model_conf):\n",
    "        super(LitTrainer, self).__init__()\n",
    "        self.save_hyperparameters('model_conf')\n",
    "        self.conf = model_conf\n",
    "        self.model = MyModel(self.conf.model_name)\n",
    "        self.criterion = LabelSmoothing()\n",
    "        self.evaluator = Accuracy()\n",
    "        # Bugfix\n",
    "        ##################################################\n",
    "        self.test_result = []\n",
    "        ##################################################\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.model(x)\n",
    "        return x\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "            \n",
    "        y_hat = self(x)\n",
    "        train_loss = self.criterion(y_hat, y)\n",
    "        self.log('train_loss', train_loss)\n",
    "        \n",
    "        y = y.argmax(dim=1)\n",
    "        y_hat = y_hat.argmax(dim=1)\n",
    "        train_score = self.evaluator(y_hat, y)\n",
    "        self.log('train_score', train_score)\n",
    "        return train_loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "            \n",
    "        y_hat = self(x)\n",
    "        valid_loss = self.criterion(y_hat, y)\n",
    "        self.log('valid_loss', valid_loss)\n",
    "        \n",
    "        y = y.argmax(dim=1)\n",
    "        y_hat = y_hat.argmax(dim=1)\n",
    "        valid_score = self.evaluator(y_hat, y)\n",
    "        self.log('valid_score', valid_score, on_epoch=True, prog_bar=True)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x = batch\n",
    "        y_hat= self(x)\n",
    "\n",
    "        score = torch.nn.functional.softmax(y_hat, dim=-1)\n",
    "        \n",
    "        y_hat = self(torch.flip(x, dims=[-1]))\n",
    "        score1 = (score + torch.nn.functional.softmax(y_hat, dim=-1)) / 2.\n",
    "        return {\"score\": score}\n",
    "\n",
    "    def test_epoch_end(self, output_results):\n",
    "        all_score = torch.cat([out[\"score\"] for out in output_results], dim=0).cpu().numpy()\n",
    "        # Bugfix\n",
    "        ##################################################\n",
    "        self.test_result = {'all_score': all_score}\n",
    "#         return {'all_score': all_score}\n",
    "        ##################################################\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        if self.conf.optimizer == 'AdamP':\n",
    "            optimizer = AdamP(self.parameters(), lr=self.conf.learning_rate)\n",
    "        elif self.conf.optimizer == 'Adam':\n",
    "            optimizer = Adam(self.parameters(), lr=self.conf.learning_rate)\n",
    "        elif self.conf.optimizer == 'AdamW':\n",
    "            optimizer = AdamW(self.parameters(), lr=self.conf.learning_rate)\n",
    "        else:\n",
    "            raise NotImplementedError(f'{self.conf.optimizer} is not used!')\n",
    "        scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=20, T_mult=1)\n",
    "        return {\n",
    "            'optimizer': optimizer,\n",
    "            'lr_scheduler': scheduler,\n",
    "        }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "prostate-tuning",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 42\n"
     ]
    },
    {
     "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>gender</th>\n",
       "      <th>race</th>\n",
       "      <th>age</th>\n",
       "      <th>path</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>000001</td>\n",
       "      <td>female</td>\n",
       "      <td>Asian</td>\n",
       "      <td>45</td>\n",
       "      <td>000001_female_Asian_45</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>000002</td>\n",
       "      <td>female</td>\n",
       "      <td>Asian</td>\n",
       "      <td>52</td>\n",
       "      <td>000002_female_Asian_52</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>000004</td>\n",
       "      <td>male</td>\n",
       "      <td>Asian</td>\n",
       "      <td>54</td>\n",
       "      <td>000004_male_Asian_54</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>000005</td>\n",
       "      <td>female</td>\n",
       "      <td>Asian</td>\n",
       "      <td>58</td>\n",
       "      <td>000005_female_Asian_58</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>000006</td>\n",
       "      <td>female</td>\n",
       "      <td>Asian</td>\n",
       "      <td>59</td>\n",
       "      <td>000006_female_Asian_59</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       id  gender   race  age                    path\n",
       "0  000001  female  Asian   45  000001_female_Asian_45\n",
       "1  000002  female  Asian   52  000002_female_Asian_52\n",
       "2  000004    male  Asian   54    000004_male_Asian_54\n",
       "3  000005  female  Asian   58  000005_female_Asian_58\n",
       "4  000006  female  Asian   59  000006_female_Asian_59"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seed_everything(conf.seed)\n",
    "\n",
    "info = pd.read_csv(f'{conf.data_dir}/train.csv')\n",
    "if conf.debug: \n",
    "    info = info.iloc[:40]\n",
    "info.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "featured-reserve",
   "metadata": {},
   "outputs": [],
   "source": [
    "info['gender_age'] = info.apply(lambda x: convert_gender_age(x.gender, x.age), axis=1)\n",
    "\n",
    "skf = StratifiedKFold(n_splits=conf.train.n_fold, shuffle=True)\n",
    "info.loc[:, 'fold'] = 0\n",
    "for fold_num, (train_index, val_index) in enumerate(skf.split(X=info.index, y=info.gender_age.values)):\n",
    "    info.loc[info.iloc[val_index].index, 'fold'] = fold_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "imperial-agent",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "Using native 16bit precision.\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name      | Type           | Params\n",
      "---------------------------------------------\n",
      "0 | model     | MyModel        | 4.0 M \n",
      "1 | criterion | LabelSmoothing | 0     \n",
      "2 | evaluator | Accuracy       | 0     \n",
      "---------------------------------------------\n",
      "4.0 M     Trainable params\n",
      "0         Non-trainable params\n",
      "4.0 M     Total params\n",
      "16.122    Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained weights for efficientnet-b0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:52: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:52: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6aec2f46bda4c339b1e45f03bbec210",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validating', layout=Layout(flex='2'), m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "image_dir = os.path.join(conf.data_dir, 'images')\n",
    "\n",
    "for fold_idx in conf.train.use_fold:\n",
    "    train = info[info.fold != fold_idx].reset_index(drop=True)\n",
    "    val = info[info.fold == fold_idx].reset_index(drop=True)\n",
    "\n",
    "    data_conf = conf.dataset\n",
    "    transforms = get_transforms(img_size=data_conf.image_size)\n",
    "    train_dataset = MaskDataset(image_dir, train, transforms['train'])\n",
    "    val_dataset = MaskDataset(image_dir, val, transforms['val'])\n",
    "    train_loader = DataLoader(train_dataset, batch_size=data_conf.batch_size, shuffle=True, num_workers=data_conf.num_workers)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=data_conf.batch_size, shuffle=False, num_workers=data_conf.num_workers)\n",
    "    \n",
    "    fold_log_dir = os.path.join(conf.log_dir, f'fold{fold_idx}')\n",
    "    os.makedirs(fold_log_dir, exist_ok=True)\n",
    "    tb_logger = TensorBoardLogger(save_dir=conf.log_dir, version=f'fold{fold_idx}')\n",
    "    early_stop_callback = EarlyStopping(monitor='valid_score', mode='max',\n",
    "                                        patience=conf.train.n_epochs//5, verbose=False)\n",
    "\n",
    "    checkpoint_callback = ModelCheckpoint(dirpath=fold_log_dir,\n",
    "                                          filename=\"{epoch:02d}-{valid_score:.4f}\",\n",
    "                                          monitor='valid_score', mode='max', verbose=False)\n",
    "    model = LitTrainer(conf.network)\n",
    "    \n",
    "    logger = tb_logger if conf.tb else None\n",
    "    trainer = pl.Trainer(gpus=len(conf.device_list), precision=16,\n",
    "                         max_epochs=conf.train.n_epochs,\n",
    "                         progress_bar_refresh_rate=1,\n",
    "                         logger=logger, callbacks=[checkpoint_callback, early_stop_callback])\n",
    "\n",
    "    trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "intense-performance",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained weights for efficientnet-b0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:52: UserWarning: The dataloader, test dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b073d8dcc2bc4adcbf32f678521da09b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Testing', layout=Layout(flex='2'), max=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "test inference is done!\n"
     ]
    }
   ],
   "source": [
    "if conf.run_test:\n",
    "    submission = pd.read_csv(os.path.join(conf.test_dir, 'info.csv'))\n",
    "    if conf.debug: \n",
    "        submission = submission.iloc[:100]\n",
    "    image_dir = os.path.join(conf.test_dir, 'images')\n",
    "\n",
    "    image_paths = [os.path.join(image_dir, img_id) for img_id in submission.ImageID]\n",
    "    dataset = TestDataset(image_paths, transforms['val'])\n",
    "\n",
    "    loader = DataLoader(\n",
    "        dataset,\n",
    "        num_workers=0,\n",
    "        batch_size=48,\n",
    "        shuffle=False\n",
    "    )\n",
    "    \n",
    "    score = []\n",
    "    for fold_idx in conf.train.use_fold:\n",
    "        ckpt_path = glob(os.path.join(conf.log_dir, f'fold{fold_idx}/*.ckpt'))[0]\n",
    "        model = LitTrainer.load_from_checkpoint(ckpt_path, model_conf=conf.network)\n",
    "        tester = pl.Trainer(gpus=1, auto_select_gpus=True)\n",
    "        # Bugfix\n",
    "        ##################################################\n",
    "        tester.test(model, loader, verbose=False)[0]\n",
    "        # prob = tester.test(model, loader, verbose=False)[0]\n",
    "        score.append(model.test_result['all_score'])\n",
    "        # score.append(prob['all_score'])\n",
    "        ##################################################\n",
    "        \n",
    "    output = np.mean(score, axis=0).argmax(axis=-1)\n",
    "    submission['ans'] = output\n",
    "\n",
    "    submission.to_csv(os.path.join(conf.log_dir, 'submission.csv'), index=False)\n",
    "    print('test inference is done!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "departmental-greeting",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plain-walker",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "neither-longitude",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}