{"cells":[{"metadata":{},"cell_type":"markdown","source":"Train-validation split strategy\n\nWe have 75k images to which three steganography algorithms (JMiPOD, JUNIWARD, UERD) are applied - positive examples,\n and the same unaltered images - negative ones. If to reduce the task to binary classification problem our training set contains 300k images.\n\nIf we split 75k into 3 parts, there could be 3 global train steps, mapping as:\n\n- 25k [0: JMiPOD], 25k [1: JUNIWARD], 25k [2: UERD];\n- 25k [1: JMiPOD], 25k [2: JUNIWARD], 25k [0: UERD];\n- 25k [2: JMiPOD], 25k [0: JUNIWARD], 25k [1: UERD].\n\nEach global step training independently.\n\nWhile traing each part splits into another 5 [number of folds]. \nSo our one global train step contains: \n1.  fold 0: 5k [0: JMiPOD] + 5k [0: JUNIWARD] + 5k [0: UERD] + 15k [Cover] = 30k [15k-pos / 15k-neg] \n \n(neg images are taking with the same names as pos ones only from Cover fold path)\n\n2.  fold 1: 5k [1: JMiPOD] + 5k [1: JUNIWARD] + 5k [1: UERD] + 15k [Cover] \n3.  ... \n4.  ...\n5.  fold 4: 5k [4: JMiPOD] + 5k [4: JUNIWARD] + 5k [4: UERD] + 15k [Cover] \n\nSo on each folder we get balanced neg classes and pos/neg targets.\nFor the next global step just change pos mapping in alg_mapping.","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"!pip install -q efficientnet_pytorch > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport gc\nimport copy\nimport random\nimport time\nimport traceback\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom glob import glob\nfrom pathlib import Path\nfrom datetime import datetime\nfrom sklearn.metrics import roc_auc_score\nfrom typing import Callable, List, Dict, Tuple, Union, Any\n\nimport cv2\nimport torch\nimport torch.nn as nn\nfrom PIL import Image\nfrom torch import Tensor\nimport albumentations as A\nfrom dataclasses import dataclass\nfrom torchvision import transforms\nfrom efficientnet_pytorch import EfficientNet\nfrom torch.utils.data import Dataset, DataLoader\nfrom albumentations.pytorch.transforms import ToTensorV2\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@dataclass\nclass Project:\n    \"\"\"\n    This class stores information about the paths.\n    \"\"\"\n\n    kaggle_path: Path = Path(\".\").absolute().parent\n    data_dir = kaggle_path / \"input/alaska2-image-steganalysis\"\n    output_dir = kaggle_path / \"working\"\n    checkpoint_dir = output_dir / \"checkpoint\"\n\n    def __post_init__(self):\n        # create the directories if they don't exist\n        self.checkpoint_dir.mkdir(exist_ok=True)\n\n\nalaska = Project()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_val_split() -> pd.DataFrame:\n\n    alg_mapping = {\"JMiPOD\": 0, \"JUNIWARD\": 1, \"UERD\": 2, \"Cover\": 3}\n\n    train_files = os.listdir(alaska.data_dir / \"Cover\")\n\n    df_pos_train_split = pd.DataFrame(train_files)\n    df_pos_train_split.rename(columns={0: \"image\"}, inplace=True)\n    df_pos_train_split[\"alg_idx\"] = df_pos_train_split.index // 25_000\n    df_pos_train_split[\"folds\"] = np.concatenate([(np.arange(25_000) // 5000)] * 3)\n\n    JMiPOD_files = df_pos_train_split[\n        df_pos_train_split.alg_idx == alg_mapping[\"JMiPOD\"]\n    ].image.values\n    JUNIWARD_files = df_pos_train_split[\n        df_pos_train_split.alg_idx == alg_mapping[\"JUNIWARD\"]\n    ].image.values\n    UERD_files = df_pos_train_split[\n        df_pos_train_split.alg_idx == alg_mapping[\"UERD\"]\n    ].image.values\n\n    path1 = [(alaska.data_dir / \"JMiPOD\" / i).as_posix() for i in JMiPOD_files]\n    path2 = [(alaska.data_dir / \"JUNIWARD\" / i).as_posix() for i in JUNIWARD_files]\n    path3 = [(alaska.data_dir / \"UERD\" / i).as_posix() for i in UERD_files]\n\n    df_pos_train_split[\"img_path\"] = np.concatenate(np.array([path1, path2, path3]))\n\n    df_pos = df_pos_train_split.assign(\n        img_path=np.concatenate(np.array([path1, path2, path3])), label=1\n    )\n\n    df_neg = df_pos[[\"image\", \"folds\"]].copy()\n    df_neg = df_neg.assign(\n        img_path=[(alaska.data_dir / \"Cover\" / i).as_posix() for i in df_neg.image],\n        label=0,\n    )\n\n    cols = [\"image\", \"folds\", \"img_path\", \"label\"]\n    df_train = pd.concat([df_pos[cols], df_neg[cols]])\n    df_train = df_train.sample(frac=1.0, random_state=101).reset_index(drop=True)\n    del df_pos_train_split, df_pos, df_neg\n    _ = gc.collect()\n\n    return df_train","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms() -> Callable:\n    \"\"\"\n    source: https://www.kaggle.com/shonenkov/train-inference-gpu-baseline\n    \"\"\"\n    return A.Compose(\n        [\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.Resize(height=512, width=512, p=1.0),\n            ToTensorV2(),\n        ],\n        p=1.0,\n    )\n\ndef get_valid_transforms() -> Callable:\n    \"\"\"\n    source: https://www.kaggle.com/shonenkov/train-inference-gpu-baseline\n    \"\"\"\n    return A.Compose([A.Resize(height=512, width=512, p=1.0), ToTensorV2(),], p=1.0)\n\n\ntransform_strategy = {\"train\": get_train_transforms(), \"val\": get_valid_transforms()}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class AlaskaClassifierDataset(Dataset):\n    \"\"\"\n    source: https://www.kaggle.com/shonenkov/train-inference-gpu-baseline\n    \"\"\"\n\n    def __init__(\n        self, data: pd.DataFrame, transform: dict, val_fold: int = 0, train: bool = True\n    ):\n\n        super().__init__()\n        self.train = train\n        self.val_fold = val_fold\n        if self.train:\n            self.data = data[data[\"folds\"] != self.val_fold].reset_index(drop=True)\n        else:\n            self.data = data[data[\"folds\"] == self.val_fold].reset_index(drop=True)\n\n        self.transform = transform[\"train\"] if self.train else transform[\"val\"]\n\n    def __len__(self) -> int:\n        return len(self.data)\n\n    def __getitem__(self, idx: int) -> Dict[torch.Tensor, torch.Tensor]:\n        image = cv2.imread(self.data.loc[idx, \"img_path\"], cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)\n        image /= 255.0\n        if self.transform:\n            sample = {\"image\": image}\n            sample = self.transform(**sample)\n            image = sample[\"image\"]\n\n        target = self.data.loc[idx, \"label\"]\n\n        return {\"image\": image, \"target\": target}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_train = train_val_split()\ntrain_dataset = AlaskaClassifierDataset(df_train, transform_strategy)\nval_dataset = AlaskaClassifierDataset(df_train, transform_strategy, train=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def build_model(pretrained_model: str = \"efficientnet-b3\") -> nn.Module:\n    model = EfficientNet.from_name(pretrained_model)\n    n_input_feats = model._fc.in_features\n    model._fc = nn.Linear(n_input_feats, 1)\n    return model\n\n\nmodel = build_model()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!git clone https://github.com/openai/spinningup.git -q\n%cd spinningup\n!pip install -q -e . > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"try: \n    from spinup.utils.logx import EpochLogger\nexcept:\n    print('ones more..')\n    from spinup.utils.logx import EpochLogger","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def to_numpy(tensor: Union[Tensor, Image.Image, np.array]) -> np.ndarray:\n    \"\"\"\n    source: https://www.kaggle.com/sermakarevich/complete-handcrafted-pipeline-in-pytorch-resnet9\n    \"\"\"\n    if type(tensor) == np.array or type(tensor) == np.ndarray:\n        return np.array(tensor)\n    elif type(tensor) == Image.Image:\n        return np.array(tensor)\n    elif type(tensor) == Tensor:\n        return tensor.cpu().detach().numpy()\n    else:\n        raise ValueError(msg)\n\n\ndef copy_data_to_device(data: Any, device: str) -> None:\n    if torch.is_tensor(data):\n        return data.to(device, dtype=torch.float)\n    elif isinstance(data, (list, tuple)):\n        return [copy_data_to_device(elem, device) for elem in data]\n    raise ValueError(\"Invalid data type {}\".format(type(data)))\n\n\ndef get_auc(y_true: np.array, y_hat: np.array) -> float:\n    try:\n        auc = roc_auc_score(np.vstack(y_true), np.vstack(y_hat))\n    except:\n        auc = -1\n    return auc","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"LOSS_FN = nn.BCEWithLogitsLoss()\nLR_SCHEDULER = lambda optim: torch.optim.lr_scheduler.ReduceLROnPlateau(\n    optim, patience=3, factor=0.5, verbose=True\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_eval_loop(\n    model=model,\n    train_dataset=train_dataset,\n    val_dataset=val_dataset,\n    criterion=LOSS_FN,\n    lr=1e-4,\n    epoch_n=3,\n    batch_size=8,\n    device=None,\n    early_stopping_patience=10,\n    l2_reg_alpha=0.0,\n    max_batches_per_epoch_train=500,\n    max_batches_per_epoch_val=500,\n    data_loader_ctor=DataLoader,\n    optimizer_ctor=None,\n    lr_scheduler_ctor=LR_SCHEDULER,\n    shuffle_train=False,\n    dataloader_workers_n=1,\n):\n\n    logger_kwargs = {\n        \"output_dir\": (alaska.checkpoint_dir).as_posix(),\n        \"output_fname\": \"alaska_progress.txt\",\n        \"exp_name\": \"val_fold_0\",\n    }\n\n    logger = EpochLogger(**logger_kwargs)\n    logger.save_config(locals())\n    start_time = time.time()\n\n    if device is None:\n        device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n        print(\"device: \", device)\n    device = torch.device(device)\n\n    model.to(device)\n\n    if optimizer_ctor is None:\n        optimizer = torch.optim.Adam(\n            model.parameters(), lr=lr, weight_decay=l2_reg_alpha\n        )\n    else:\n        optimizer = optimizer_ctor(model.parameters(), lr=lr)\n\n    if lr_scheduler_ctor is not None:\n        lr_scheduler = lr_scheduler_ctor(optimizer)\n    else:\n        lr_scheduler = None\n\n    train_dataloader = data_loader_ctor(\n        train_dataset,\n        batch_size=batch_size,\n        shuffle=shuffle_train,\n        num_workers=dataloader_workers_n,\n    )\n    val_dataloader = data_loader_ctor(\n        val_dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=dataloader_workers_n,\n    )\n\n    best_val_loss = float(\"inf\")\n    auc_valid = 0\n    best_epoch_i = 0\n    best_model = copy.deepcopy(model)\n\n    for epoch_i in range(epoch_n):\n        try:\n            epoch_start = datetime.now()\n            print(\"epoch {}\".format(epoch_i + 1))\n\n            # *************** training part ***********************\n            mean_train_loss = 0\n            train_batches_n = 0\n            y_true_train, y_pred_train = [], []\n            model.train()\n            pbar = tqdm(\n                enumerate(train_dataloader),\n                total=max_batches_per_epoch_train,\n                desc=\"Epoch {}\".format(epoch_i),\n                ncols=0,\n            )\n\n            for batch_i, d in pbar:\n                if batch_i > max_batches_per_epoch_train:\n                    break\n\n                batch_x = copy_data_to_device(d[\"image\"], device)\n                batch_y = copy_data_to_device(d[\"target\"].view(-1, 1), device)\n                pred = model(batch_x)\n\n                y_true_train.append(to_numpy(batch_y))\n                y_pred_train.append(to_numpy(pred))\n\n                loss = criterion(pred, batch_y)\n\n                model.zero_grad()\n                loss.backward()\n\n                optimizer.step()\n\n                mean_train_loss += float(loss)\n                train_batches_n += 1\n\n            mean_train_loss /= train_batches_n\n            score = get_auc(y_true_train, y_pred_train)\n            logger.store(TrainLoss=mean_train_loss, TrainAUC=score)\n\n            # logging training info\n            Mode_train = \"*******\"\n            logger.log_tabular(\"Mode_train\", Mode_train)\n            logger.log_tabular(\"Epoch\", epoch_i + 1)\n            logger.log_tabular(\"TrainLoss\", average_only=True)\n            logger.log_tabular(\"TrainAUC\", average_only=True)\n            logger.log_tabular(\n                \"TotalGradientSteps\", (epoch_i + 1) * max_batches_per_epoch_train\n            )\n\n            # *************** eval part ***********************\n            model.eval()\n            mean_val_loss = 0\n            val_batches_n = 0\n\n            total_samples, correct = 0, 0\n            y_true_valid, y_pred_valid = [], []\n\n            with torch.no_grad():\n                for batch_i, d in enumerate(val_dataloader):\n                    if batch_i > max_batches_per_epoch_val:\n                        break\n\n                    batch_x = copy_data_to_device(d[\"image\"], device)\n                    batch_y = copy_data_to_device(d[\"target\"].view(-1, 1), device)\n                    pred = model(batch_x)\n\n                    y_true_valid.append(to_numpy(batch_y))\n                    y_pred_valid.append(to_numpy(pred))\n\n                    loss = criterion(pred, batch_y)\n                    mean_val_loss += float(loss)\n\n                    val_batches_n += 1\n\n                    correct += pred.eq(batch_y.view_as(pred)).sum().item()\n                    total_samples += batch_x.size()[0]\n\n                score = get_auc(y_true_valid, y_pred_valid)\n                accuracy = 100.0 * correct / total_samples\n                mean_val_loss /= val_batches_n\n                logger.store(Loss=mean_val_loss, AUC=score, Accuracy=accuracy)\n\n            if mean_val_loss < best_val_loss:\n                best_epoch_i = epoch_i\n                best_val_loss = mean_val_loss\n                best_model = copy.deepcopy(model)\n                logger.setup_pytorch_saver(best_model)  # Setup model saving\n                print(\"New best model!\")\n            elif epoch_i - best_epoch_i > early_stopping_patience:\n                print(\n                    \"The model has not improved over the past {} epochs, stop training\".format(\n                        early_stopping_patience\n                    )\n                )\n                break\n\n            if lr_scheduler is not None:\n                lr_scheduler.step(mean_val_loss)\n\n            # logging valid info\n            Mode_valid = \"*******\"\n            logger.log_tabular(\"Mode_valid\", Mode_valid)\n            logger.log_tabular(\"Loss\", average_only=True)\n            logger.log_tabular(\"AUC\", average_only=True)\n            logger.log_tabular(\"Accuracy\", average_only=True)\n            logger.log_tabular(\n                \"Total time\",\n                datetime.utcfromtimestamp(time.time() - start_time).strftime(\"%M:%S\"),\n            )\n            logger.dump_tabular()\n\n            # exception handling\n            print()\n        except KeyboardInterrupt:\n            print(\"Stopped by user\")\n            break\n        except Exception as ex:\n            print(\"Training error: {}\\n{}\".format(ex, traceback.format_exc()))\n            break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_eval_loop()","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}