{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Pretraining the CNN with regressing the number of atoms of each element in the structure","metadata":{}},{"cell_type":"markdown","source":"![](https://i.ibb.co/WPDk5d7/Presentation1.jpg)","metadata":{}},{"cell_type":"markdown","source":"As discussed [here](https://www.kaggle.com/c/bms-molecular-translation/discussion/224257#1258149) and mentioned by [@hengck23](https://www.kaggle.com/hengck23), I'm going to first pre-train the CNN part of the final model on a task **similar** to this competition's. Then we can freeze the CNN and only train the RNN part on the real task of the competition. In the final step, both the CNN and RNN will be trained with a lower learning_rate end-to-end for the main task. So, here is what will be done:\n\n1. Pretraining the CNN on a similar task to competition's [This notebook]\n2. Freeze CNN and train the RNN part on the main task\n3. Train both CNN and RNN with lower learning_rate on the main task\n\nIn this notebook, **we are going to regress the number of atoms of each element** present in the biochemical structure. For doing this, the model needs to learn a lot of useful features of the image which I hope will help the CNN for the final task of this competition which is giving us the whole InChI of the molecule.\n\n\n","metadata":{}},{"cell_type":"markdown","source":"My **preprocessing** notebook related to this one is [here](https://www.kaggle.com/moeinshariatnia/cnn-rnn-cnn-pretraining-w-regression-preprocess). \n\n\nI want to thank [Y.Nakama](https://www.kaggle.com/yasufuminakama) who shared really great starter notebooks.","metadata":{}},{"cell_type":"markdown","source":"## Importing libraries","metadata":{}},{"cell_type":"code","source":"!pip install timm","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport gc\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom tqdm.autonotebook import tqdm\n\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import mean_absolute_error\nfrom sklearn.preprocessing import StandardScaler\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\n\nimport timm\nimport albumentations as A\n\nnp.set_printoptions(suppress=True)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Config","metadata":{}},{"cell_type":"markdown","source":"Set the **debug to False** to train the model on the whole data","metadata":{}},{"cell_type":"code","source":"class CFG:\n    debug = True # set to False to run on the whole data\n    epochs = 4\n    batch_size = 128\n    num_workers = 4\n    size = 224\n    classes = [\"C\", \"Cl\", \"F\", \"H\", \"N\", \"O\", \"S\"]\n    num_classes = len([\"C\", \"Cl\", \"F\", \"H\", \"N\", \"O\", \"S\"])\n    model_name = 'resnet34'\n    fc_hidden_size = 128\n    dropout = 0.2 # set to 0 or None to remove it\n    seed = 42\n    n_fold = 5\n    learning_rate = 3e-4\n    scheduler = \"ReduceLROnPlateau\"\n    patience = 2\n    factor = 0.5\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Taking a look at the dataframe","metadata":{}},{"cell_type":"code","source":"train = pd.read_pickle(\"../input/cnn-rnn-cnn-pretraining-w-regression-preprocess/cnn_pretrain.pkl\")\ntrain.head(3)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"elements = [\"C\", \"Cl\", \"F\", \"H\", \"N\", \"O\", \"S\"]\nweights = []\nfor element in elements:\n    zeros = train[train[element] == 0].shape[0]\n    non_zeros = train[train[element] != 0].shape[0]\n    weight = non_zeros / len(train)\n    print(f\"Element: {element} | Percent non zero: {weight * 100} | weight: {round(weight, 2)}\")\n    weights.append(weight)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We will ignore elements B, Br, and Si because they are really rare (less than 0.01 percent non zero)","metadata":{}},{"cell_type":"markdown","source":"## Calculating the weights","metadata":{}},{"cell_type":"markdown","source":"Using the non-zero ratios as the weights associated to each of the elements. Used in the final loss function.","metadata":{}},{"cell_type":"code","source":"weights = np.array(weights)\nweights_normalized = (weights / sum(weights)).round(3)\nweights_normalized","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CFG.weights = torch.tensor(weights_normalized).float().to(CFG.device)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Set CFG.debug to False to train on the whole data","metadata":{}},{"cell_type":"code","source":"if CFG.debug:\n    train = train.sample(10_000).reset_index(drop=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataset","metadata":{}},{"cell_type":"markdown","source":"We will scale the counts of atoms with **StandardScaler** from scikit-learn. It helps the model's learning if all the outputs are on roughly similar scales (Carbon has much more atoms compared to say Oxygen; so, scaling will help)\nKeep in mind that we will use train stats for the validation set; so, only using .fit function in train dataset","metadata":{}},{"cell_type":"code","source":"class BMSDataset(torch.utils.data.Dataset):\n    def __init__(self, img_paths, counts, scaler=None, transforms=None):\n        self.img_paths = img_paths\n        if scaler is None:\n            self.scaler = StandardScaler()\n            self.counts = self.scaler.fit_transform(counts)\n        else:\n            self.scaler = scaler\n            self.counts = scaler.transform(counts)\n        \n        self.transforms = transforms\n    \n    def __getitem__(self, idx):\n        img = cv2.imread(self.img_paths[idx])[..., ::-1]\n        if self.transforms is not None:\n            img = self.transforms(image=img)['image']\n        counts = self.counts[idx]\n        img = torch.tensor(img).permute(2, 0, 1).float()\n        counts = torch.tensor(counts).float()\n        return img, counts\n    \n    def inverse_transform(self, x):\n        \"\"\"\n        x type is np array\n        \"\"\"\n        return self.scaler.inverse_transform(x)\n        \n    def __len__(self):\n        return len(self.img_paths)\n    \n\ndef get_transforms(mode=\"train\"):\n    if mode == \"train\":\n        return A.Compose([\n            A.Resize(CFG.size, CFG.size),\n            A.Normalize()\n        ])\n    else:\n        return A.Compose([\n            A.Resize(CFG.size, CFG.size),\n            A.Normalize()\n        ])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self,\n                 model_name=CFG.model_name, \n                 num_classes=CFG.num_classes, \n                 dropout=CFG.dropout,\n                 fc_hidden_size=CFG.fc_hidden_size,\n                 pretrained=True):\n        \n        super().__init__()\n        self.cnn = timm.create_model(\n            model_name, pretrained=pretrained, num_classes=0, global_pool=\"\"\n        )\n        num_features = self.cnn.num_features\n        self.dropout = nn.Dropout(dropout) if dropout is not None else nn.Identity()\n        self.fc = nn.Linear(num_features, fc_hidden_size)\n        self.output = nn.Linear(fc_hidden_size, num_classes)\n        \n    def get_features(self, x):\n        return self.cnn(x)\n    \n    def forward(self, x):\n        batch_size = x.size(0)\n        x = self.get_features(x)\n        channels = x.size(1)\n        x = F.adaptive_avg_pool2d(x, 1).reshape(batch_size, channels)\n        x = self.dropout(x)\n        x = F.relu(self.fc(x))\n        x = self.dropout(x)\n        out = self.output(x)\n\n        return out","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utils","metadata":{}},{"cell_type":"code","source":"class AvgMeter:\n    def __init__(self, name=\"Metric\"):\n        self.name = name\n        self.reset()\n    \n    def reset(self):\n        self.avg, self.sum, self.count = [0]*3\n    \n    def update(self, val, count=1):\n        self.count += count\n        self.sum += val * count\n        self.avg = self.sum / self.count\n    \n    def __repr__(self):\n        text = f\"{self.name}: {self.avg:.4f}\"\n        return text\n    \ndef get_lr(optimizer):\n    for param_group in optimizer.param_groups:\n        return param_group[\"lr\"]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def one_epoch(model, \n              criterion, \n              loader, \n              optimizer=None, \n              lr_scheduler=None,\n              mode=\"train\", \n              step=\"batch\"):\n    \n    loss_meter = AvgMeter()\n    mae_orig_meter = AvgMeter()\n    \n    tqdm_object = tqdm(loader, total=len(loader))\n    for images, targets in tqdm_object:\n        images, targets = images.to(CFG.device), targets.to(CFG.device)\n        preds = model(images)\n        loss = criterion(preds, targets)\n        loss *= CFG.weights.unsqueeze(0)\n        loss = loss.mean()\n        \n        if mode == \"train\":\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            if step == \"batch\":\n                lr_scheduler.step()\n                \n        count = images.size(0)\n        loss_meter.update(loss.item(), count)\n        \n        targets_orig = loader.dataset.inverse_transform(targets.cpu().numpy())\n        preds_orig = loader.dataset.inverse_transform(preds.detach().cpu().numpy())\n        mae_orig = mean_absolute_error(targets_orig, preds_orig)\n        mae_orig_meter.update(mae_orig, count)\n\n        if mode == \"train\":\n            tqdm_object.set_postfix(train_loss=loss_meter.avg, train_mae=mae_orig_meter.avg, lr=get_lr(optimizer))\n        else:\n            tqdm_object.set_postfix(valid_loss=loss_meter.avg, valid_mae=mae_orig_meter.avg)\n    if mode != \"train\":\n        print(f\"Targets: \\n\"\n              f\"{np.round(targets_orig[:5], 2)}\")\n        \n        print(f\"Preds: \\n\"\n              f\"{np.round(preds_orig[:5], 2)}\")\n        \n    return loss_meter, mae_orig_meter","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_eval(model, \n               train_loader, \n               valid_loader, \n               criterion, \n               optimizer, \n               lr_scheduler=None, \n               step=None,\n               fold=0):\n    \n    best_loss = float('inf')\n    \n    for epoch in range(CFG.epochs):\n        print(\"*\" * 30)\n        print(f\"Epoch {epoch + 1}\")\n        current_lr = get_lr(optimizer)\n        \n        model.train()\n        train_loss, train_mae = one_epoch(model, \n                                          criterion, \n                                          train_loader, \n                                          optimizer=optimizer,\n                                          lr_scheduler=lr_scheduler,\n                                          mode=\"train\",\n                                          step=step)                     \n        model.eval()\n        with torch.no_grad():\n            valid_loss, valid_mae = one_epoch(model, \n                                              criterion, \n                                              valid_loader, \n                                              optimizer=None,\n                                              lr_scheduler=None,\n                                              mode=\"valid\")\n        \n        if valid_loss.avg < best_loss:\n            best_loss = valid_loss.avg\n            torch.save(model.state_dict(), f'best_fold_{fold}.pt')\n            print(\"Saved best model!\")\n        \n        # or you could do: if step == \"epoch\":\n        if isinstance(lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):\n            lr_scheduler.step(valid_loss.avg)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def make_folds(dataframe):\n    folds = dataframe.copy()\n    Fold = StratifiedKFold(n_splits=CFG.n_fold, shuffle=True, random_state=CFG.seed)\n    for n, (train_index, val_index) in enumerate(Fold.split(folds, folds['InChI_length'])):\n        folds.loc[val_index, 'fold'] = int(n)\n    folds['fold'] = folds['fold'].astype(int)\n    print(folds.groupby(['fold']).size())\n    return folds","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def make_loader(dataframe, mode=\"train\", scaler=None):\n    transforms = get_transforms(mode=mode)\n    dataset = BMSDataset(dataframe['file_path'].values, \n                         dataframe.loc[:, CFG.classes].values, \n                         scaler=scaler, \n                         transforms=transforms)\n    \n    dataloader = torch.utils.data.DataLoader(dataset, \n                                             batch_size=CFG.batch_size, \n                                             shuffle=True if mode == \"train\" else False,\n                                             num_workers=CFG.num_workers)\n    return dataloader","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def one_fold(folds, fold):  \n    print(f\"Training Fold: {fold}\")\n    \n    \n    train_dataframe = folds[folds['fold'] != fold].reset_index(drop=True)\n    valid_dataframe = folds[folds['fold'] == fold].reset_index(drop=True)\n\n    train_loader = make_loader(train_dataframe, \"train\", None) # Setting scaler to None, fits a new scaler on train data\n    valid_loader = make_loader(valid_dataframe, \"valid\", train_loader.dataset.scaler) # Using train scaler for valid data\n\n    model = Model().to(CFG.device)\n    optimizer = torch.optim.Adam(model.parameters(), lr=CFG.learning_rate)\n    if CFG.scheduler == \"ReduceLROnPlateau\":\n        lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, \n                                                                  mode=\"min\", \n                                                                  factor=CFG.factor, \n                                                                  patience=CFG.patience)\n        step = \"epoch\"\n    \n    criterion = nn.MSELoss(reduction='none')\n    train_eval(model, \n               train_loader, \n               valid_loader,\n               criterion, \n               optimizer, \n               lr_scheduler=lr_scheduler,\n               step=step,\n               fold=fold)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"folds = make_folds(train)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"If debug == True, only one fold will be trained.","metadata":{}},{"cell_type":"code","source":"if CFG.debug:\n    one_fold(folds, 0)\nelse:\n    for i in range(CFG.n_fold):\n        one_fold(folds, i)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}