{"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":"# G2Net Multi-Model PyTorch Training Script with W&B Experiment Tracking 🚀\n\nTraining Script with Multiple model architecture support along with W&B experiment tracking.\n\nI have tried to write as generalised training script as possible and I will be pushing major updates to it in the future so keep checking!","metadata":{}},{"cell_type":"markdown","source":"<center>\n<img src=\"https://i.imgur.com/gb6B4ig.png\" width=300px height=100px>\n</center>\n<br>\nWandb is a developer tool for companies turn deep learning research projects into deployed software by helping teams track their models, visualize model performance and easily automate training and improving models. We will use their tools to log hyperparameters and output metrics from your runs, then visualize and compare results and quickly share findings with your colleagues.\n\nYou can check more about them here: [wandb.ai](https://www.wandb.ai)","metadata":{}},{"cell_type":"markdown","source":"*Note: Huge Credits to Y.Nakama's Datasets ([train](https://www.kaggle.com/yasufuminakama/g2net-n-mels-128-train-images) and [test](https://www.kaggle.com/yasufuminakama/g2net-n-mels-128-test-images)) that I am using in this notebook along with the loading functions!*","metadata":{}},{"cell_type":"markdown","source":"**If you found this notebook useful you can leave an upvote! If there's scope for improvement or any mistakes you found, please comment down below!**","metadata":{}},{"cell_type":"code","source":"%%sh\npip install -q timm\npip install -q wandb --upgrade","metadata":{"execution":{"iopub.status.busy":"2021-07-20T06:53:16.186885Z","iopub.execute_input":"2021-07-20T06:53:16.187204Z","iopub.status.idle":"2021-07-20T06:53:33.715331Z","shell.execute_reply.started":"2021-07-20T06:53:16.187172Z","shell.execute_reply":"2021-07-20T06:53:33.714369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1. Imports and Utility functions\n\nSome imports, utility functions and W&B login.","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\nimport gc\nimport platform\nimport random\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom tqdm.notebook import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\n\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.model_selection import StratifiedKFold\n\nimport timm\nimport cv2\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform\n\nimport wandb\n\nimport warnings\nwarnings.simplefilter('ignore')","metadata":{"execution":{"iopub.status.busy":"2021-07-20T07:21:31.802486Z","iopub.execute_input":"2021-07-20T07:21:31.80296Z","iopub.status.idle":"2021-07-20T07:21:31.818732Z","shell.execute_reply.started":"2021-07-20T07:21:31.802918Z","shell.execute_reply":"2021-07-20T07:21:31.817435Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# W&B Login\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n\nwandb.login(key=wb_key)\n\nCONFIG = dict(\n    lr=1e-5,\n    autocast = True,\n    resize=(224, 224),\n    model_name = 'tf_efficientnet_b4',\n    pretrained = True,\n    epochs = 3,\n    scheduler = 'CosineAnnealingLR',\n    n_splits = 5,\n    split = 0.97,\n    folds = [1, 2, 3, 4, 5],\n    workers = 4,\n    train_bs = 64,\n    valid_bs = 64,\n    seed = 42,\n    num_labels = 1,\n    grad_acc_steps = 1,\n    max_gnorm = 1000,\n    architecture = \"CNN\",\n    infra = \"Kaggle\",\n    competition = 'g2net',\n    _wandb_kernel = 'tanaym'\n)\n\nrun = wandb.init(project='g2net', \n                 config=CONFIG,\n                 group='effnet',\n                 job_type='train'\n                )","metadata":{"execution":{"iopub.status.busy":"2021-07-20T06:53:35.299119Z","iopub.execute_input":"2021-07-20T06:53:35.29947Z","iopub.status.idle":"2021-07-20T06:53:35.306576Z","shell.execute_reply.started":"2021-07-20T06:53:35.299431Z","shell.execute_reply":"2021-07-20T06:53:35.305453Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Config:\n    lr=1e-5\n    autocast = False\n    resize=(224, 224)\n    model_name = 'tf_efficientnet_b4'\n    pretrained = True\n    epochs = 3\n    scheduler = 'CosineAnnealingLR'\n    n_splits = 5\n    split = 0.97\n    folds = [1, 2, 3, 4, 5]\n    workers = 4\n    train_bs = 64\n    valid_bs = 64\n    seed = 42\n    scaler = GradScaler()\n    num_labels = 1\n    target_name = 'target'\n    grad_acc_steps = 1\n    wandb = True\n    max_gnorm = 1000\n    TRAIN_PATH = \"../input/g2net-gravitational-wave-detection/train\"\n    TEST_PATH = \"../input/g2net-gravitational-wave-detection/test\"\n    train_file = \"../input/g2net-gravitational-wave-detection-file-paths/training_labels_with_paths.csv\"\n    sample_sub = \"../input/g2net-gravitational-wave-detection/sample_submission.csv\"","metadata":{"execution":{"iopub.status.busy":"2021-07-20T07:05:38.814592Z","iopub.execute_input":"2021-07-20T07:05:38.814916Z","iopub.status.idle":"2021-07-20T07:05:38.823733Z","shell.execute_reply.started":"2021-07-20T07:05:38.814889Z","shell.execute_reply":"2021-07-20T07:05:38.820745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def wandb_log(**kwargs):\n    \"\"\"\n    Logs a key-value pair to W&B\n    \"\"\"\n    for k, v in kwargs.items():\n        wandb.log({k: v})\n\ndef get_train_file_path(image_id):\n    \"\"\"\n    Taken from Y.Nakama's notebook\n    \"\"\"\n    return \"../input/g2net-n-mels-128-train-images/{}.npy\".format(image_id)\n\ndef get_test_file_path(image_id):\n    \"\"\"\n    Taken from Y.Nakama's notebook\n    \"\"\"\n    return \"../input/g2net-n-mels-128-test-images/{}.npy\".format(image_id)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2021-07-20T06:57:12.114341Z","iopub.execute_input":"2021-07-20T06:57:12.114701Z","iopub.status.idle":"2021-07-20T06:57:12.119913Z","shell.execute_reply.started":"2021-07-20T06:57:12.114668Z","shell.execute_reply":"2021-07-20T06:57:12.118991Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Model Architectures\n\nList of all different model that you can use.","metadata":{}},{"cell_type":"code","source":"# Models\nclass VITModel(nn.Module):\n    def __init__(self):\n        super(VITModel, self).__init__()\n        self.backbone = timm.create_model(Config.model_name, pretrained=Config.pretrained, in_chans=1)\n        self.n_f = self.backbone.head.in_features\n        self.backbone.head = nn.Linear(self.n_f, Config.num_labels)\n    def forward(self, x):\n        return self.backbone(x)\n\nclass EffNetModel(nn.Module):\n    def __init__(self):\n        super(EffNetModel, self).__init__()\n        self.backbone = timm.create_model(Config.model_name, pretrained=Config.pretrained, in_chans=1)\n        self.n_f = self.backbone.classifier.in_features\n        self.backbone.classifier = nn.Linear(self.n_f, Config.num_labels)\n    def forward(self, x):\n        return self.backbone(x)\n    \nclass ResNextModel(nn.Module):\n    def __init__(self):\n        super(ResNextModel, self).__init__()\n        self.backbone = timm.create_model(Config.model_name, pretrained=Config.pretrained, in_chans=1)\n        self.n_f = self.backbone.fc.in_features\n        self.backbone.fc = nn.Linear(self.n_f, Config.num_labels)\n    def forward(self, x):\n        return self.backbone(x)","metadata":{"execution":{"iopub.status.busy":"2021-07-20T06:57:17.696049Z","iopub.execute_input":"2021-07-20T06:57:17.696419Z","iopub.status.idle":"2021-07-20T06:57:17.70693Z","shell.execute_reply.started":"2021-07-20T06:57:17.696387Z","shell.execute_reply":"2021-07-20T06:57:17.705603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. Custom Dataset Class\n\nNow we define a custom dataset class that will load our data when training the model.\n\nThe `_get_fname()` function essentially takes an id and gets the corresponding `.npy` file from the folder structure.","metadata":{}},{"cell_type":"code","source":"class G2NetData(Dataset):\n    def __init__(self, data, is_test=False):\n        self.data = data\n        self.is_test = is_test        \n        self.file_names = self.data['file_path'].values\n        self.labels = self.data['target'].values\n        \n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        file_path = self.file_names[idx]\n        image = np.load(file_path)\n        image = image[np.newaxis, :, :]\n        image = torch.from_numpy(image).float()\n        label = torch.tensor(self.labels[idx]).float()\n        \n        return image, label","metadata":{"execution":{"iopub.status.busy":"2021-07-20T06:57:39.735182Z","iopub.execute_input":"2021-07-20T06:57:39.735518Z","iopub.status.idle":"2021-07-20T06:57:39.742068Z","shell.execute_reply.started":"2021-07-20T06:57:39.735472Z","shell.execute_reply":"2021-07-20T06:57:39.741223Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Augments\n\nSome basic Augments using Albumentations library.\n\nI am currently not using them in the code but planning to do so in the future.","metadata":{}},{"cell_type":"code","source":"class Augments:\n    \"\"\"\n    Contains Train, Validation Augments\n    \"\"\"\n    train_augments = A.Compose([\n        ToTensorV2(p=1.0),\n    ],p=1.)\n    \n    valid_augments = A.Compose([\n        ToTensorV2(p=1.0),\n    ], p=1.)","metadata":{"execution":{"iopub.status.busy":"2021-07-20T06:57:41.44394Z","iopub.execute_input":"2021-07-20T06:57:41.444265Z","iopub.status.idle":"2021-07-20T06:57:41.449071Z","shell.execute_reply.started":"2021-07-20T06:57:41.444234Z","shell.execute_reply":"2021-07-20T06:57:41.44796Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5. Trainer Class\n\nThis is the main trainer class. It has training and validation function for one-one epoch each.","metadata":{}},{"cell_type":"code","source":"class Trainer:\n    def __init__(self, model, optimizer, scheduler, train_dataloader, valid_dataloader, device):\n        self.model = model\n        self.optimizer = optimizer\n        self.scheduler = scheduler\n        self.train_data = train_dataloader\n        self.valid_data = valid_dataloader\n        self.loss_fn = self.yield_loss\n        self.val_loss_fn = self.yield_loss\n        self.device = device\n        \n    def yield_loss(self, outputs, targets):\n        \"\"\"\n        Returns the loss function\n        \"\"\"\n        return nn.BCEWithLogitsLoss()(outputs, targets)\n    \n    def train_one_epoch(self):\n        \"\"\"\n        Trains the model for one epoch\n        \"\"\"\n        prog_bar = tqdm(enumerate(self.train_data), total=len(self.train_data))\n        self.model.train()\n        avg_loss = 0\n        for idx, inputs in prog_bar:\n            image = inputs[0].to(self.device, dtype=torch.float)\n            targets = inputs[1].to(self.device, dtype=torch.float)\n\n            if Config.autocast:\n                with autocast():\n                    outputs = self.model(image).view(-1)\n                    loss = self.loss_fn(outputs, targets)\n                Config.scaler.scale(loss).backward()\n                Config.scaler.step(self.optimizer)\n                Config.scaler.update()\n                \n            else:\n                outputs = self.model(image).view(-1)\n                loss = self.loss_fn(outputs, targets)\n                loss.backward()\n                self.optimizer.step()\n            \n            self.optimizer.zero_grad()\n            prog_bar.set_description('loss: {:.2f}'.format(loss.item()))\n\n            avg_loss += loss.item()\n\n        return avg_loss / len(self.train_data)\n    \n    def valid_one_epoch(self):\n        \"\"\"\n        Validates the model over all batches (1 epoch)\n        \"\"\"\n        prog_bar = tqdm(enumerate(self.valid_data), total=len(self.valid_data))\n        self.model.eval()\n        all_targets = []\n        all_predictions = []\n        avg_loss = 0\n        with torch.no_grad():\n            for idx, inputs in prog_bar:\n                image = inputs[0].to(self.device, dtype=torch.float)\n                targets = inputs[1].to(self.device, dtype=torch.float)\n                \n                outputs = self.model(image).view(-1)\n                \n                val_loss = self.val_loss_fn(outputs, targets)\n                prog_bar.set_description('val_loss: {:.2f}'.format(val_loss.item()))\n                \n                all_targets.extend(targets.cpu().detach().numpy().tolist())\n                all_predictions.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())\n                \n                avg_loss += val_loss.item()\n        val_roc_auc = roc_auc_score(all_targets, all_predictions)\n        return val_roc_auc, avg_loss / len(self.valid_data)\n    \n    def get_model(self):\n        return self.model","metadata":{"execution":{"iopub.status.busy":"2021-07-20T06:58:01.553736Z","iopub.execute_input":"2021-07-20T06:58:01.554066Z","iopub.status.idle":"2021-07-20T06:58:01.569725Z","shell.execute_reply.started":"2021-07-20T06:58:01.554036Z","shell.execute_reply":"2021-07-20T06:58:01.568905Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6. Main Training Code\n\nBelow code combines everything into one big training and validation procedure.","metadata":{}},{"cell_type":"code","source":"# Training Code\nif __name__ == '__main__':\n    if torch.cuda.is_available():\n        print(\"[INFO] Using GPU: {}\\n\".format(torch.cuda.get_device_name()))\n        DEVICE = torch.device('cuda')\n    else:\n        print(\"\\n[INFO] GPU not found. Using CPU: {}\\n\".format(platform.processor()))\n        DEVICE = torch.device('cpu')\n    \n    # Data loading and stuff\n    data = pd.read_csv(Config.train_file)\n    data['file_path'] = data['id'].apply(get_train_file_path)\n    data = data.sample(frac=1).reset_index(drop=True)\n    \n    kf = StratifiedKFold(n_splits=Config.n_splits, shuffle=True)\n    \n    print(f\"Training Model: {Config.model_name}\")\n    for fold_, (train_idx, valid_idx) in enumerate(kf.split(data, data[Config.target_name])):\n        print(f\"{'='*20} Fold: {fold_} {'='*20}\")\n        \n        train_data = data.loc[train_idx].reset_index(drop=True)\n        valid_data = data.loc[valid_idx].reset_index(drop=True)\n    \n        print(f\"Training on {train_data.shape[0]} samples and validating on {valid_data.shape[0]} samples\")\n\n        # Make Training and Validation Datasets\n        training_set = G2NetData(\n            data=train_data\n        )\n\n        validation_set = G2NetData(\n            data=valid_data\n        )\n\n        train = DataLoader(\n            training_set,\n            batch_size=Config.train_bs,\n            shuffle=True,\n            num_workers=Config.workers,\n            pin_memory=True\n        )\n\n        valid = DataLoader(\n            validation_set,\n            batch_size=Config.valid_bs,\n            shuffle=False,\n            num_workers=Config.workers\n        )\n\n        # Declare model and initialize other things\n        model = EffNetModel().to(DEVICE)\n        nb_train_steps = int(len(train_data) / Config.train_bs * Config.epochs)\n        optimizer = torch.optim.AdamW(model.parameters(), lr=Config.lr, weight_decay=1e-6)\n\n        trainer = Trainer(model, optimizer, None, train, valid, DEVICE)\n\n        # Do the training and validation\n        for epoch in range(1, Config.epochs+1):\n            print(f\"\\n{'--'*5} EPOCH: {epoch} {'--'*5}\\n\")\n\n            # Train for 1 epoch\n            tr_lss = trainer.train_one_epoch()\n\n            # Validate for 1 epoch\n            current_roc, vl_lss = trainer.valid_one_epoch()\n            if Config.wandb:\n                wandb_log(\n                    epoch_train_loss=tr_lss,\n                    epoch_valid_loss=vl_lss,\n                    roc_auc_score=current_roc,\n                )\n            print(f\"Validation ROC-AUC: {current_roc:.4f}\")\n\n            torch.save(trainer.get_model().state_dict(), f\"fold_{fold_}_{Config.model_name}.pt\")\n        \n        del train_data, valid_data, training_set, validation_set, train, valid, model, optimizer, trainer\n        gc.collect()\n        torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2021-07-20T07:21:54.281169Z","iopub.execute_input":"2021-07-20T07:21:54.281481Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}