{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Cassava Classification - PyTorch Starter (Train)"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\n\nimport albumentations as A\nimport cv2\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torch.optim as optim\n\nfrom tqdm.notebook import tqdm\nfrom torch.utils.data import Dataset, DataLoader\nfrom albumentations.pytorch import ToTensorV2\n\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import KFold, StratifiedKFold\n\nimport plotly.express as px\nimport plotly.graph_objects as go\nfrom plotly.subplots import make_subplots\n\nimport warnings  \nwarnings.filterwarnings('ignore')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = '/kaggle/input/cassava-leaf-disease-classification'\nDEBUG = False\n\nSEED = 42\nN_FOLDS = 2 if DEBUG else 5\nN_EPOCHS = 2 if DEBUG else 10\nBATCH_SIZE = 64\nSIZE = 256","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaDataset(Dataset):\n    \n    def __init__(self, df, dataset='train', transforms=None):\n    \n        self.df = df\n        self.transforms=transforms\n        self.dataset=dataset\n        \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self, idx):\n        image_src = f'{DIR_INPUT}/{self.dataset}_images/{self.df.loc[idx, \"image_id\"]}'\n        # print(image_src)\n        image = cv2.imread(image_src, cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        # label = self.df.loc[idx, 'label']\n        labels = self.df.loc[idx, ['cls0', 'cls1', 'cls2', 'cls3', 'cls4']].values\n        labels = torch.from_numpy(labels.astype(np.int8))\n        labels = labels.unsqueeze(-1)\n        \n        if self.transforms:\n            transformed = self.transforms(image=image)\n            image = transformed['image']\n\n        return image, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaModel(nn.Module):\n    \n    def __init__(self, num_classes=5):\n        super().__init__()\n        \n        self.backbone = torchvision.models.resnet18(pretrained=True)\n        \n        in_features = self.backbone.fc.in_features\n\n        self.logit = nn.Linear(in_features, num_classes)\n        \n    def forward(self, x):\n        batch_size, C, H, W = x.shape\n        \n        x = self.backbone.conv1(x)\n        x = self.backbone.bn1(x)\n        x = self.backbone.relu(x)\n        x = self.backbone.maxpool(x)\n\n        x = self.backbone.layer1(x)\n        x = self.backbone.layer2(x)\n        x = self.backbone.layer3(x)\n        x = self.backbone.layer4(x)\n        \n        x = F.adaptive_avg_pool2d(x,1).reshape(batch_size,-1)\n        x = F.dropout(x, 0.25, self.training)\n\n        x = self.logit(x)\n\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = A.Compose([\n    A.RandomResizedCrop(height=SIZE, width=SIZE, p=1.0),\n    A.Flip(),\n    A.ShiftScaleRotate(rotate_limit=1.0, p=0.8),\n\n    A.Normalize(p=1.0),\n    ToTensorV2(p=1.0),\n])\n\ntransforms_valid = A.Compose([\n    A.Resize(height=SIZE, width=SIZE, p=1.0),\n    A.Normalize(p=1.0),\n    ToTensorV2(p=1.0),\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df[['cls0', 'cls1', 'cls2', 'cls3', 'cls4']] = train_labels = pd.get_dummies(train_df.iloc[:, 1])\n\n# For debugging.\nif DEBUG:\n    train_df = train_df.sample(n=100)\n    train_df.reset_index(drop=True, inplace=True)\n\ntrain_labels = train_df.iloc[:, 1].values\n\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"folds = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)\noof_preds = np.zeros((train_df.shape[0],))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class DenseCrossEntropy(nn.Module):\n\n    def __init__(self):\n        super(DenseCrossEntropy, self).__init__()\n        \n        \n    def forward(self, logits, labels):\n        logits = logits.float()\n        labels = labels.float()\n        \n        logprobs = F.log_softmax(logits, dim=-1)\n        \n        loss = -labels * logprobs\n        loss = loss.sum(-1)\n\n        return loss.mean()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_one_fold(i_fold, model, criterion, optimizer, dataloader_train, dataloader_valid):\n    \n    train_fold_results = []\n\n    for epoch in range(N_EPOCHS):\n\n        # print('  Epoch {}/{}'.format(epoch + 1, N_EPOCHS))\n        # print('  ' + ('-' * 20))\n        os.system(f'echo \\\"  Epoch {epoch}\\\"')\n\n        model.train()\n        tr_loss = 0\n\n        for step, batch in enumerate(dataloader_train):\n\n            images = batch[0]\n            labels = batch[1]\n\n            images = images.to(device, dtype=torch.float)\n            labels = labels.to(device, dtype=torch.float)\n            \n            outputs = model(images)\n            \n            loss = criterion(outputs, labels.squeeze(-1))                \n            loss.backward()\n\n            tr_loss += loss.item()\n\n            optimizer.step()\n            optimizer.zero_grad()\n\n        # Validate\n        model.eval()\n        val_loss = 0\n        val_preds = None\n        val_labels = None\n\n        for step, batch in enumerate(dataloader_valid):\n\n            images = batch[0]\n            labels = batch[1]\n\n            if val_labels is None:\n                val_labels = labels.clone().squeeze(-1)\n            else:\n                val_labels = torch.cat((val_labels, labels.squeeze(-1)), dim=0)\n\n            images = images.to(device, dtype=torch.float)\n            labels = labels.to(device, dtype=torch.float)\n\n            with torch.no_grad():\n                outputs = model(images)\n\n                loss = criterion(outputs, labels.squeeze(-1))\n                val_loss += loss.item()\n\n                preds = torch.softmax(outputs, dim=1).data.cpu()\n\n                if val_preds is None:\n                    val_preds = preds\n                else:\n                    val_preds = torch.cat((val_preds, preds), dim=0)\n        \n        val_preds = torch.argmax(val_preds, dim=1)\n\n        train_fold_results.append({\n            'fold': i_fold,\n            'epoch': epoch,\n            'train_loss': tr_loss / len(dataloader_train),\n            'valid_loss': val_loss / len(dataloader_valid),\n            'valid_score': accuracy_score(torch.argmax(val_labels, dim=1), val_preds)\n        })\n\n    return val_preds, train_fold_results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"submissions = None\ntrain_results = []\n\nfor i_fold, (train_idx, valid_idx) in enumerate(folds.split(train_df, train_labels)):\n    print(\"Fold {}/{}\".format(i_fold + 1, N_FOLDS))\n\n    valid = train_df.iloc[valid_idx]\n    valid.reset_index(drop=True, inplace=True)\n\n    train = train_df.iloc[train_idx]\n    train.reset_index(drop=True, inplace=True)    \n\n    dataset_train = CassavaDataset(df=train, dataset='train', transforms=transforms_train)\n    dataset_valid = CassavaDataset(df=valid, dataset='train', transforms=transforms_valid)\n\n    dataloader_train = DataLoader(dataset_train, batch_size=BATCH_SIZE, num_workers=4, shuffle=True)\n    dataloader_valid = DataLoader(dataset_valid, batch_size=BATCH_SIZE, num_workers=4, shuffle=False)\n\n    device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device('cpu')\n\n    model = CassavaModel(num_classes=5)\n    model.to(device)\n\n    criterion = DenseCrossEntropy()\n    plist = [{'params': model.parameters(), 'lr': 5e-5}]\n    optimizer = optim.Adam(plist, lr=5e-5)\n    \n    val_preds, train_fold_results = train_one_fold(i_fold, model, criterion, optimizer, dataloader_train, dataloader_valid)\n    oof_preds[valid_idx] = val_preds.numpy()\n    \n    train_results = train_results + train_fold_results\n    \n    torch.save({\n        'fold': i_fold,\n        'lr': optimizer.state_dict()[\"param_groups\"][0]['lr'],\n        'model_state_dict': model.cpu().state_dict(),\n        'optimizer_state_dict': optimizer.state_dict(),\n        # 'scheduler_state_dict'\n        # 'scaler_state_dict'\n    }, f\"model_state_fold_{i_fold}.pth\")\n\nprint(\"{}-Folds CV score: {:.4f}\".format(N_FOLDS, accuracy_score(train_labels, oof_preds)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Show train history"},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"train_results = pd.DataFrame(train_results)\ntrain_results.head(10)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"fig = make_subplots(rows=2, cols=1)\n\ncolors = [\n    ('#d32f2f', '#ef5350'),\n    ('#303f9f', '#5c6bc0'),\n    ('#00796b', '#26a69a'),\n    ('#fbc02d', '#ffeb3b'),\n    ('#5d4037', '#8d6e63'),\n]\n\nfor i in range(N_FOLDS):\n    data = train_results[train_results['fold'] == i]\n\n    fig.add_trace(go.Scatter(x=data['epoch'].values,\n                             y=data['train_loss'].values,\n                             mode='lines',\n                             visible='legendonly' if i > 0 else True,\n                             line=dict(color=colors[i][0], width=2),\n                             name='Train loss - Fold #{}'.format(i)),\n                 row=1, col=1)\n\n    fig.add_trace(go.Scatter(x=data['epoch'],\n                             y=data['valid_loss'].values,\n                             mode='lines+markers',\n                             visible='legendonly' if i > 0 else True,\n                             line=dict(color=colors[i][1], width=2),\n                             name='Valid loss - Fold #{}'.format(i)),\n                 row=1, col=1)\n    \n    fig.add_trace(go.Scatter(x=data['epoch'].values,\n                             y=data['valid_score'].values,\n                             mode='lines+markers',\n                             line=dict(color=colors[i][0], width=2),\n                             name='Valid score - Fold #{}'.format(i),\n                             showlegend=False),\n                 row=2, col=1)\n\nfig.update_layout({\n  \"annotations\": [\n    {\n      \"x\": 0.225, \n      \"y\": 1.0, \n      \"font\": {\"size\": 16}, \n      \"text\": \"Train / valid losses\", \n      \"xref\": \"paper\", \n      \"yref\": \"paper\", \n      \"xanchor\": \"center\", \n      \"yanchor\": \"bottom\", \n      \"showarrow\": False\n    }, \n    {\n      \"x\": 0.775, \n      \"y\": 1.0, \n      \"font\": {\"size\": 16}, \n      \"text\": \"Validation scores\", \n      \"xref\": \"paper\", \n      \"yref\": \"paper\", \n      \"xanchor\": \"center\", \n      \"yanchor\": \"bottom\", \n      \"showarrow\": False\n    }, \n  ]\n})\n\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}