{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Import modeuls"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install pretrainedmodels iterative-stratification > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport time\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport pretrainedmodels\nimport albumentations as albu\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import AdamW\nfrom torch.nn.parameter import Parameter\nfrom torch.utils.data import DataLoader, Dataset\nfrom catalyst.utils import get_device, set_global_seed\nfrom iterstrat.ml_stratifiers import MultilabelStratifiedKFold\nfrom sklearn.metrics import recall_score","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Utils for loading dataset and calculating metrics"},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_data(BASE_PATH):\n    print('Reading train.csv file....')\n    train = pd.read_csv(BASE_PATH + 'train.csv')\n    print('Training.csv file have {} rows and {} columns'.format(\n        train.shape[0], train.shape[1]))\n\n    print('Reading test.csv file....')\n    test = pd.read_csv(BASE_PATH + 'test.csv')\n    print('Test.csv file have {} rows and {} columns'.format(\n        test.shape[0], test.shape[1]))\n\n    print('Reading sample_submission.csv file....')\n    sample_submission = pd.read_csv(BASE_PATH + 'sample_submission.csv')\n    print('Sample_submission.csv file have {} rows and {} columns'.format(\n        sample_submission.shape[0], sample_submission.shape[1]))\n    return train, test, sample_submission\n\ndef prepare_image(datadir, data_type='train', submission=False):\n    assert data_type in ['train', 'test']\n    images = []\n    # only use 1/4 size because of CPU memory\n    for i in tqdm([0]):\n        if submission:\n            image_df_list = pd.read_parquet(datadir + f'{data_type}_image_data_{i}.parquet')\n        else:\n            image_df_list = pd.read_feather(datadir + f'{data_type}_image_data_{i}.feather')\n\n        HEIGHT = 137\n        WIDTH = 236\n        image = image_df_list.iloc[:, 1:].values.reshape(-1, HEIGHT, WIDTH)\n        images.append(image)\n        del image_df_list, image\n        gc.collect()\n\n    images = np.concatenate(images, axis=0)\n    print('Image shape : ', images.shape)\n    return images\n\ndef macro_recall(outputs, targets):\n    pred_labels = [np.argmax(out, axis=1) for out in outputs]\n    # target_col = ['grapheme_root', 'consonant_diacritic', 'vowel_diacritic']\n    recall_grapheme = recall_score(targets[:, 0], pred_labels[0], average='macro')\n    recall_consonant = recall_score(targets[:, 1], pred_labels[1], average='macro')\n    recall_vowel = recall_score(targets[:, 2], pred_labels[2], average='macro')\n    scores = [recall_grapheme, recall_consonant, recall_vowel]\n    final_score = np.average(scores, weights=[2, 1, 1])\n    return final_score, scores","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset, model and loss class"},{"metadata":{"trusted":true},"cell_type":"code","source":"HEIGHT = 137\nWIDTH = 236\nSIZE = 224\n\n\ndef bbox(img):\n    rows = np.any(img, axis=1)\n    cols = np.any(img, axis=0)\n    rmin, rmax = np.where(rows)[0][[0, -1]]\n    cmin, cmax = np.where(cols)[0][[0, -1]]\n    return rmin, rmax, cmin, cmax\n\n\n# https://www.kaggle.com/iafoss/image-preprocessing-128x128\ndef crop_resize(img, size=SIZE, pad=16):\n    # crop a box around pixels large than the threshold\n    # some images contain line at the sides\n    ymin, ymax, xmin, xmax = bbox(img[5:-5, 5:-5] > 80)\n    # cropping may cut too much, so we need to add it back\n    xmin = xmin - 13 if (xmin > 13) else 0\n    ymin = ymin - 10 if (ymin > 10) else 0\n    xmax = xmax + 13 if (xmax < WIDTH - 13) else WIDTH\n    ymax = ymax + 10 if (ymax < HEIGHT - 10) else HEIGHT\n    img = img[ymin:ymax, xmin:xmax]\n    # remove lo intensity pixels as noise\n    img[img < 28] = 0\n    lx, ly = xmax-xmin, ymax-ymin\n    l = max(lx, ly) + pad  # noqa\n    # make sure that the aspect ratio is kept in rescaling\n    img = np.pad(img, [((l-ly)//2,), ((l-lx)//2,)], mode='constant')\n    return cv2.resize(img, (size, size))\n\n\nclass BengaliAIDataset(Dataset):\n    def __init__(self, images=None, labels=None, size=None, transforms=None):\n        self.images = images\n        self.labels = labels\n        self.size = size\n        self.transforms = transforms\n\n        # set dummy labels\n        if self.labels is None:\n            self.labels = np.zeros(len(images))\n\n        # validation\n        if len(images) != len(labels):\n            raise ValueError('Do not match the data size between input and output')\n\n    def __getitem__(self, idx):\n        img = self.images[idx]\n        label = self.labels[idx]\n        if self.size is not None:\n            img = crop_resize(img, self.size)\n        if self.transforms is not None:\n            augmented = self.transforms(image=img)\n            img = augmented['image']\n        # 2dim to 3dim (N,N) -> (1,N,N)\n        img = img[None, :, :]\n        return torch.tensor(img, dtype=torch.float), torch.tensor(label, dtype=torch.long)\n\n    def __len__(self):\n        return len(self.images)\n    \nclass LinearBlock(nn.Module):\n    def __init__(self, in_features=None, out_features=None, bias=True,\n                 use_bn=True, activation=F.relu, dropout_ratio=-1, residual=False):\n        super(LinearBlock, self).__init__()\n        # validation\n        if in_features is None or out_features is None:\n            raise ValueError('You should set both in_features and out_features!!')\n        self.linear = nn.Linear(in_features, out_features, bias=bias)\n        self.activation = activation\n        self.use_bn = use_bn\n        self.dropout_ratio = dropout_ratio\n        self.residual = residual\n\n        if use_bn:\n            self.bn = nn.BatchNorm1d(out_features)\n        if dropout_ratio > 0.:\n            self.dropout = nn.Dropout(p=dropout_ratio)\n        else:\n            self.dropout = None\n\n    def __call__(self, x):\n        h = self.linear(x)\n        if self.use_bn:\n            h = self.bn(h)\n        if self.activation is not None:\n            h = self.activation(h)\n        if self.residual:\n            h = self._residual_add(h, x)\n        if self.dropout_ratio > 0:\n            h = self.dropout(h)\n        return h\n\n    def _residual_add(lhs, rhs):\n        lhs_ch, rhs_ch = lhs.shape[1], rhs.shape[1]\n        if lhs_ch < rhs_ch:\n            out = lhs + rhs[:, :lhs_ch]\n        elif lhs_ch > rhs_ch:\n            out = torch.cat([lhs[:, :rhs_ch] + rhs, lhs[:, rhs_ch:]], dim=1)\n        else:\n            out = lhs + rhs\n        return out\n\n\nclass GeM(nn.Module):\n    def __init__(self, p=3, eps=1e-6):\n        super(GeM, self).__init__()\n        self.p = Parameter(torch.ones(1)*p)\n        self.eps = eps\n\n    def forward(self, x):\n        return F.avg_pool2d(x.clamp(min=self.eps).pow(self.p),\n                            (x.size(-2), x.size(-1))).pow(1./self.p)\n\n\nclass BengaliBaselineClassifier(nn.Module):\n    def __init__(self, n_grapheme=168, n_vowel=11, n_consonant=7,\n                 pretrainedmodels=None, in_channels=1,\n                 hdim=512, use_bn=True, pretrained=None):\n        super(BengaliBaselineClassifier, self).__init__()\n        self.n_grapheme = n_grapheme\n        self.n_vowel = n_vowel\n        self.n_consonant = n_consonant\n        self.base_model = pretrainedmodels\n        self.conv0 = nn.Conv2d(in_channels, 3, kernel_size=3, stride=1, padding=1, bias=True)\n        inch = self.base_model.last_linear.in_features\n        self.gem_pool = GeM()\n        self.fc1 = LinearBlock(inch, hdim, use_bn=use_bn, activation=F.relu)\n        self.logits_for_grapheme = LinearBlock(hdim, n_grapheme, use_bn=False, activation=None)\n        self.logits_for_vowel = LinearBlock(hdim, n_vowel, use_bn=False, activation=None)\n        self.logits_for_consonant = LinearBlock(hdim, n_consonant, use_bn=False, activation=None)\n\n    def forward(self, x):\n        h = self.conv0(x)\n        h = self.base_model.features(h)\n        h = self.gem_pool(h)\n        h = h.view(h.size(0), -1)\n        h = self.fc1(h)\n        logits_for_grapheme = self.logits_for_grapheme(h)\n        logits_for_consonant = self.logits_for_consonant(h)\n        logits_for_vowel = self.logits_for_vowel(h)\n        # target_col = ['grapheme_root', 'consonant_diacritic', 'vowel_diacritic']\n        logits = (logits_for_grapheme, logits_for_consonant, logits_for_vowel)\n        return logits\n\n    \nclass BaselineLoss(nn.Module):\n    def __init__(self):\n        super(BaselineLoss, self).__init__()\n\n    def forward(self, pred, target):\n        criterion = nn.CrossEntropyLoss()\n        loss = criterion(pred[0], target[:, 0]) + \\\n            criterion(pred[1], target[:, 1]) + criterion(pred[2], target[:, 2])\n        return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Original pytorch runner"},{"metadata":{"trusted":true},"cell_type":"code","source":"class BengaliRunner:\n    def __init__(self, device='cpu'):\n        self.device = device\n\n    def train(self, model, criterion, optimizer, loaders, scheduler=None, logdir=None,\n              num_epochs=20, score_func=None):\n        # validation\n        for dict_val in [loaders]:\n            if 'train' in dict_val and 'valid' in dict_val:\n                pass\n            else:\n                raise ValueError('You should set train and valid key.')\n\n        # setup training\n        model = model.to(self.device)\n        train_loader = loaders['train']\n        valid_loader = loaders['valid']\n        train_criterion = criterion\n        valid_criterion = criterion\n        best_score = -1.0\n        best_avg_val_loss = 100\n        log_df = pd.DataFrame(\n            [], columns=['epoch', 'loss', 'valid_loss', 'score', 'recall_grapheme',\n                         'recall_consonant', 'recall_vowel', 'time'],\n            index=range(num_epochs)\n        )\n        for epoch in range(num_epochs):\n            start_time = time.time()\n            # release memory\n            torch.cuda.empty_cache()\n            gc.collect()\n            # train for one epoch\n            avg_loss = self._train_model(model, train_criterion, optimizer, train_loader, scheduler)\n            # evaluate on validation set\n            avg_val_loss, score, scores = self._validate_model(model, valid_criterion, valid_loader, score_func)\n\n            # log\n            elapsed_time = time.time() - start_time\n            log_df.iloc[epoch] = [epoch + 1, avg_loss, avg_val_loss, score, scores[0], scores[1], scores[2], elapsed_time]\n\n            # the position of this depends on the scheduler you use\n            if scheduler is not None:\n                scheduler.step()\n\n            # save best params\n            save_path = 'best_model.pth'\n            if logdir is not None:\n                save_path = os.path.join(logdir, save_path)\n\n            if score is None:\n                if best_avg_val_loss > avg_val_loss:\n                    best_avg_val_loss = avg_val_loss\n                    best_param_loss = model.state_dict()\n                    torch.save(best_param_loss, save_path)\n                    print('Save the best model on Epoch {}'.format(epoch + 1))\n            else:\n                if best_score < score:\n                    best_score = score\n                    best_param_score = model.state_dict()\n                    torch.save(best_param_score, save_path)\n                    print('Save the best model on Epoch {}'.format(epoch + 1))\n\n            # save log\n            log_df.to_csv(os.path.join(logdir, 'log.csv'))\n\n        return True\n\n    def predict_loader(self, model, loader, resume='best_model.pth'):\n        # set up models\n        model = model.to(self.device)\n        model.load_state_dict(torch.load(resume))\n        model.eval()\n\n        # prediction\n        grapheme_preds = []\n        consonant_preds = []\n        vowel_preds = []\n        with torch.no_grad():\n            for idx, batch in tqdm(enumerate(loader), total=len(loader)):\n                images, labels = batch\n                images = images.to(self.device)\n                labels = labels.to(self.device)\n\n                # output\n                output_valid = model(images)\n                logits = [out.detach().cpu().numpy() for out in output_valid]\n                # target_col = ['grapheme_root', 'consonant_diacritic', 'vowel_diacritic']\n                grapheme_preds.extend(logits[0])\n                consonant_preds.extend(logits[1])\n                vowel_preds.extend(logits[2])\n\n        return grapheme_preds, consonant_preds, vowel_preds\n\n    def _train_model(self, model, criterion, optimizer, train_loader, scheduler=None):\n        # switch to train mode\n        model.train()\n        avg_loss = 0.0\n        for idx, batch in tqdm(enumerate(train_loader), total=len(train_loader)):\n            images, labels = batch\n            images = images.to(self.device)\n            labels = labels.to(self.device)\n\n            # training\n            output_train = model(images)\n            loss = criterion(output_train, labels)\n\n            # update params\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n\n            # calc loss\n            avg_loss += loss.item() / len(train_loader)\n\n        return avg_loss\n\n    def _validate_model(self, model, criterion, valid_loader, score_func=None):\n        # switch to eval mode\n        model.eval()\n        avg_val_loss = 0.\n        valid_grapheme = []\n        valid_consonant = []\n        valid_vowel = []\n        targets = []\n        with torch.no_grad():\n            for idx, batch in tqdm(enumerate(valid_loader), total=len(valid_loader)):\n                images, labels = batch\n                images = images.to(self.device)\n                labels = labels.to(self.device)\n\n                # output\n                output_valid = model(images)\n                avg_val_loss += criterion(output_valid, labels).item() / len(valid_loader)\n\n                # calc score\n                logits = [out.detach().cpu().numpy() for out in output_valid]\n                labels = labels.detach().cpu().numpy()\n                valid_grapheme.extend(logits[0])\n                valid_consonant.extend(logits[1])\n                valid_vowel.extend(logits[2])\n                targets.extend(labels)\n\n            score = None\n            if score_func is not None:\n                # TODO : you should write valid score calculation\n                # In this case, we pass sigmoid function\n                targets = np.array(targets)\n                valid_preds = [valid_grapheme, valid_consonant, valid_vowel]\n                score, scores = score_func(valid_preds, targets)\n\n        return avg_val_loss, score, scores\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"# set your params\nBASE_PATH = '../input/bengaliai-cv19/'\nDATA_PATH = '../input/bengaliaicv19feather/'\nBASE_LOGDIR = './logs'\nNUM_FOLDS = 5\nBATCH_SIZE = 64\nEPOCHS = 10\nSEED = 1234\nSIZE = 96\nLR = 0.001\nHOLD_OUT = True\n\n# fix seed\nset_global_seed(SEED)\n\n# read dataset\ntrain, _, _ = read_data(BASE_PATH)\ntrain_all_images = prepare_image(DATA_PATH, data_type='train', submission=False)\ntrain = train.iloc[0:len(train_all_images)]\n\n# init\ntarget_col = ['grapheme_root', 'consonant_diacritic', 'vowel_diacritic']\ndevice = get_device()\ntrain_data_transforms = albu.Compose([\n    albu.ShiftScaleRotate(rotate_limit=10, scale_limit=.1),\n    albu.Cutout(p=0.5),\n])\ntest_data_transforms = None\n\n# cross validation\nkf = MultilabelStratifiedKFold(n_splits=NUM_FOLDS, random_state=SEED)\nids = kf.split(X=train_all_images, y=train[target_col].values)\nfor fold, (train_idx, valid_idx) in enumerate(ids):\n    print(\"Current Fold: \", fold + 1)\n    logdir = os.path.join(BASE_LOGDIR, 'fold_{}'.format(fold + 1))\n    os.makedirs(logdir, exist_ok=True)\n\n    train_df, valid_df = train.iloc[train_idx], train.iloc[valid_idx]\n    print(\"Train and Valid Shapes are\", train_df.shape, valid_df.shape)\n\n    print(\"Preparing train datasets....\")\n    train_dataset = BengaliAIDataset(\n        images=train_all_images[train_idx], labels=train_df[target_col].values,\n        size=SIZE, transforms=train_data_transforms\n    )\n\n    print(\"Preparing valid datasets....\")\n    valid_dataset = BengaliAIDataset(\n        images=train_all_images[valid_idx], labels=valid_df[target_col].values,\n        size=SIZE, transforms=test_data_transforms\n    )\n\n    print(\"Preparing dataloaders datasets....\")\n    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n    valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False)\n    loaders = {'train': train_loader, 'valid': valid_loader}\n\n    # release memory\n    del train_df, valid_df, train_dataset, valid_dataset\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    # init models\n    resnet34 = pretrainedmodels.__dict__[\"resnet34\"](pretrained=\"imagenet\")\n    model = BengaliBaselineClassifier(pretrainedmodels=resnet34, hdim=512)\n    model = model.to(device)\n    criterion = BaselineLoss()\n    optimizer = AdamW(model.parameters(), lr=LR)\n    scheduler = None\n\n    # catalyst trainer\n    runner = BengaliRunner(device=device)\n    # model training\n    runner.train(model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler,\n                     loaders=loaders, logdir=logdir, num_epochs=EPOCHS, score_func=macro_recall)\n\n    # release memory\n    del model, runner, train_loader, valid_loader, loaders\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    if HOLD_OUT is True:\n        break","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}