{"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":"# Pytorch Lightning:aux targets+weighted loss+thresholds","metadata":{}},{"cell_type":"markdown","source":"In this notebook, I try to convert notebook from [vslaykovsky](https://www.kaggle.com/code/vslaykovsky/train-pytorch-aux-targets-weighted-loss-thres) to PytorchLightning, which allow us easily switch between single and multiple GPUs training.","metadata":{}},{"cell_type":"markdown","source":"## 1. Imports, Constants, Dependencies","metadata":{}},{"cell_type":"code","source":"! pip install pytorch-lightning --quiet","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:30.03925Z","iopub.execute_input":"2023-01-04T13:22:30.040018Z","iopub.status.idle":"2023-01-04T13:22:39.81199Z","shell.execute_reply.started":"2023-01-04T13:22:30.039921Z","shell.execute_reply":"2023-01-04T13:22:39.81068Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\nimport os\n\n# import cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport sklearn\nimport torch\nfrom PIL import Image\nfrom sklearn.preprocessing import LabelEncoder\nimport sys\nsys.path.append('../input/timm-pytorch-image-models/pytorch-image-models-master')\nfrom timm import create_model, list_models\nfrom timm.data import create_transform\nfrom torch.cuda.amp import GradScaler, autocast\nfrom torch.utils.data import DataLoader, Dataset\nimport pytorch_lightning as pl\nfrom pytorch_lightning.loggers import WandbLogger\nfrom pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor\n\nfrom tqdm import tqdm\n\nimport wandb\n\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_key = user_secrets.get_secret(\"WANDB_API_KEY\")\n\nwandb.login(key=wb_key)","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:39.814139Z","iopub.execute_input":"2023-01-04T13:22:39.817641Z","iopub.status.idle":"2023-01-04T13:22:45.219408Z","shell.execute_reply.started":"2023-01-04T13:22:39.817597Z","shell.execute_reply":"2023-01-04T13:22:45.218413Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# DEBUG MODE\nDEBUG=False\n\n\n# GLOBAL VARIABLES (Unchanged variables)\nTRAIN_IMAGES_PATH = f'/kaggle/input/rsna-cut-off-empty-space-from-images'\nTARGET = 'cancer'\nCATEGORY_AUX_TARGETS = ['site_id', 'laterality', 'view', 'implant', 'biopsy', 'invasive', 'BIRADS', 'density', 'difficult_negative_case', 'machine_id', 'age']\nRSNA_2022_PATH = '../input/rsna-breast-cancer-detection'\nTRAIN_IMAGES_PATH = f'/kaggle/input/rsna-cut-off-empty-space-from-images'\nMAX_TRAIN_BATCHES = 40000\nMAX_EVAL_BATCHES = 400\nMODELS_PATH = '/kaggle/input/wandb-models/models'\nNUM_WORKERS = 2\nPREDICT_MAX_BATCHES = 1e9\nN_FOLDS = 5\nFOLDS = np.array(os.environ.get('FOLDS', '0,1,2,3,4').split(',')).astype(int)\nWANDB_SWEEP_PROJECT = 'rsna-breast-cancer-sweeps'\n\nclass CFG:\n    ONE_CYCLE = True\n    ONE_CYCLE_PCT_START = 0.1\n    ADAMW = False\n    ADAMW_DECAY = 0.024\n    # ONE_CYCLE_MAX_LR = float(os.environ.get('LR', '0.0008'))\n    ONE_CYCLE_MAX_LR = float(os.environ.get('LR', '0.0004'))\n    EPOCHS = int(os.environ.get('EPOCHS', 3))\n    MODEL_TYPE = os.environ.get('MODEL', 'seresnext50_32x4d')\n    DROPOUT = float(os.environ.get('DROPOUT', 0.0))\n    AUG = os.environ.get('AUG', 'true').lower() == 'true'\n    AUX_LOSS_WEIGHT = 94\n    POSITIVE_TARGET_WEIGHT=20\n#     BATCH_SIZE = 32\n    BATCH_SIZE = 16\n    AUTO_AUG_M = 10\n    AUTO_AUG_N = 2\n    TTA = False\n    CHECKPOINT_PATH=\"./checkpoints\"\n    \nWANDB_RUN_NAME = f'plot_lr_{CFG.MODEL_TYPE}_lr{CFG.ONE_CYCLE_MAX_LR}_ep{CFG.EPOCHS}_bs{CFG.BATCH_SIZE}_pw{CFG.POSITIVE_TARGET_WEIGHT}_aux{CFG.AUX_LOSS_WEIGHT}_{\"adamw\" if CFG.ADAMW else \"adam\"}_{\"aug\" if CFG.AUG else \"noaug\"}_drop{CFG.DROPOUT}'\nWANDB_PROJECT = 'RSNA-breast-cancer-v1'\nprint('run', WANDB_RUN_NAME, 'folds', FOLDS)\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\npl.seed_everything(seed=42)\n\n\n# DEBUG=True\n# if DEBUG:\n#     FOLDS = np.array(os.environ.get('FOLDS', '0,1,2,3,4').split(',')).astype(int)\n","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:45.221055Z","iopub.execute_input":"2023-01-04T13:22:45.221717Z","iopub.status.idle":"2023-01-04T13:22:45.266188Z","shell.execute_reply.started":"2023-01-04T13:22:45.221678Z","shell.execute_reply":"2023-01-04T13:22:45.26504Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # W&B Logger\n# wandb_logger = WandbLogger(\n#     project='RSNA-Lightning-torch-converted', \n#     job_type='train', \n# #     config=(CFG.__dict__),\n# )\n    ","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:45.270017Z","iopub.execute_input":"2023-01-04T13:22:45.270439Z","iopub.status.idle":"2023-01-04T13:22:45.276276Z","shell.execute_reply.started":"2023-01-04T13:22:45.270396Z","shell.execute_reply":"2023-01-04T13:22:45.274264Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2.   Loading train/eval/test DF","metadata":{}},{"cell_type":"code","source":"original_df = pd.read_csv(\"/kaggle/input/rsna-breast-cancer-detection/train.csv\")\ndf = original_df\n# df = original_df.head(10000)\n# ========== stratifiedGroupKFold ===========\nfrom sklearn.model_selection import StratifiedGroupKFold\n\nsplit = StratifiedGroupKFold(N_FOLDS)\nfor k, (_, test_idx) in enumerate(split.split(df, df.cancer, groups=df.patient_id)):\n    df.loc[test_idx, 'split'] = k\ndf.split = df.split.astype(int)\ndf.groupby('split').cancer.mean()","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:45.278301Z","iopub.execute_input":"2023-01-04T13:22:45.278814Z","iopub.status.idle":"2023-01-04T13:22:49.809465Z","shell.execute_reply.started":"2023-01-04T13:22:45.278777Z","shell.execute_reply":"2023-01-04T13:22:49.808249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df.age.fillna(df.age.mean(), inplace=True)\ndf['age'] = pd.qcut(df.age, 10, labels=range(10), retbins=False).astype(int)\ndf","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.811223Z","iopub.execute_input":"2023-01-04T13:22:49.811955Z","iopub.status.idle":"2023-01-04T13:22:49.843671Z","shell.execute_reply.started":"2023-01-04T13:22:49.811912Z","shell.execute_reply":"2023-01-04T13:22:49.842556Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df[CATEGORY_AUX_TARGETS] = df[CATEGORY_AUX_TARGETS].apply(LabelEncoder().fit_transform)","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.845605Z","iopub.execute_input":"2023-01-04T13:22:49.846004Z","iopub.status.idle":"2023-01-04T13:22:49.912952Z","shell.execute_reply.started":"2023-01-04T13:22:49.845966Z","shell.execute_reply":"2023-01-04T13:22:49.911831Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. Dataset class & transform function","metadata":{}},{"cell_type":"code","source":"import torchvision\n\ndef get_transforms(aug=False):\n    \"\"\"\n    # old transforms\n    create_transform(\n        (1024, 512), \n        mean=0.53, #(0.53, 0.53, 0.53),\n        std=0.23, #(0.23, 0.23, 0.23),\n        is_training=is_training, \n        auto_augment=f'rand-m{config.AUTO_AUG_M}-n{config.AUTO_AUG_N}'\n    )\n    \"\"\"\n    def transforms(img):\n        img = img.convert('RGB')#.resize((512, 512))\n        if aug:\n            tfm = [\n                torchvision.transforms.RandomHorizontalFlip(0.5),\n                torchvision.transforms.RandomRotation(degrees=(-5, 5)), \n                torchvision.transforms.RandomResizedCrop((1024, 512), scale=(0.8, 1), ratio=(0.45, 0.55)) \n            ]\n        else:\n            tfm = [\n                torchvision.transforms.RandomHorizontalFlip(0.5),\n                torchvision.transforms.Resize((1024, 512))\n            ]\n        img = torchvision.transforms.Compose(tfm + [            \n            torchvision.transforms.ToTensor(), # chanel, height, width.\n            torchvision.transforms.Normalize(mean=0.2179, std=0.0529),\n            \n        ])(img)\n        return img\n\n    return lambda img: transforms(img)\n\n# if DEBUG:\n#     tfm = get_transforms(aug=True)\n#     img = Image.open(f\"{TRAIN_IMAGES_PATH}/10006/1459541791.png\")\n#     print(img.size)\n#     plt.imshow(np.array(img), cmap='gray')\n#     plt.show()\n\n#     plt.figure(figsize=(20, 20))\n#     for i in range(8):\n#         v = tfm(img).permute(1, 2, 0)\n#         v -= v.min()\n#         v /= v.max()\n#         # plt.imshow(v)\n#         # break\n#         plt.subplot(2, 4, i + 1).imshow(v)\n#     plt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.914729Z","iopub.execute_input":"2023-01-04T13:22:49.915215Z","iopub.status.idle":"2023-01-04T13:22:49.924447Z","shell.execute_reply.started":"2023-01-04T13:22:49.915178Z","shell.execute_reply":"2023-01-04T13:22:49.923384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BreastCancerDataSet(torch.utils.data.Dataset):\n    def __init__(self, df, path, transforms=None):\n        super().__init__()\n        self.df = df\n        self.path = path\n        self.transforms = transforms\n\n    def __getitem__(self, i):\n\n        path = f'{self.path}/{self.df.iloc[i].patient_id}/{self.df.iloc[i].image_id}.png'\n        try:\n            img = Image.open(path).convert('RGB')\n        except Exception as ex:\n            print(path, ex)\n            return None\n\n        if self.transforms is not None:\n            img = self.transforms(img)\n\n        # If not in test.\n        if TARGET in self.df.columns:\n            cancer_target = torch.as_tensor(self.df.iloc[i].cancer)\n            cat_aux_targets = torch.as_tensor(self.df.iloc[i][CATEGORY_AUX_TARGETS])\n            return img, cancer_target, cat_aux_targets\n\n        return img\n\n    def __len__(self):\n        return len(self.df)","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.926037Z","iopub.execute_input":"2023-01-04T13:22:49.92656Z","iopub.status.idle":"2023-01-04T13:22:49.938494Z","shell.execute_reply.started":"2023-01-04T13:22:49.926521Z","shell.execute_reply":"2023-01-04T13:22:49.937385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BreastCancerDataModule(pl.LightningDataModule):\n    def __init__(self, df_train, df_valid):\n        super().__init__()\n        self.df_train = df_train\n        self.df_valid = df_valid\n\n    def setup(self, stage=None):\n        self.train_dataset = BreastCancerDataSet(\n            self.df_train,\n            path=TRAIN_IMAGES_PATH,\n            transforms= get_transforms(aug=True),\n        )\n\n        self.valid_dataset = BreastCancerDataSet(\n            self.df_valid,\n            path=TRAIN_IMAGES_PATH,\n            transforms= get_transforms(aug=True),\n        )\n\n    def train_dataloader(self):\n        return DataLoader(\n            self.train_dataset,\n            batch_size=CFG.BATCH_SIZE,\n            num_workers=NUM_WORKERS,\n            shuffle=True,\n            pin_memory=True,\n        )\n\n    def val_dataloader(self):\n        return DataLoader(\n            self.valid_dataset,\n            batch_size=CFG.BATCH_SIZE,\n            num_workers=NUM_WORKERS,\n            shuffle=False,\n            pin_memory=True,\n        )\nif DEBUG:\n    df_train_demo = df.iloc[:5]\n    df_valid_demo = df.iloc[5:10]\n    demo_data_module = BreastCancerDataModule(df_train_demo, df_valid_demo)\n    demo_data_module.setup()","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.942867Z","iopub.execute_input":"2023-01-04T13:22:49.943201Z","iopub.status.idle":"2023-01-04T13:22:49.953756Z","shell.execute_reply.started":"2023-01-04T13:22:49.943172Z","shell.execute_reply":"2023-01-04T13:22:49.95282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Model","metadata":{}},{"cell_type":"code","source":"class BSDRSNAModule(pl.LightningModule):\n    def __init__(self, aux_classes,len_dl_train ,model_type=CFG.MODEL_TYPE, dropout=0.):\n        super().__init__()\n        self.len_dl_train = len_dl_train\n\n        # back bone\n        self.model = create_model(model_type, pretrained=True, num_classes=0, drop_rate=dropout)\n        self.backbone_dim = self.model(torch.randn(1, 3, 512, 512)).shape[-1]\n\n        # head for cancer:\n        self.nn_cancer = torch.nn.Sequential(\n            torch.nn.Linear(self.backbone_dim, 1),\n        )\n\n        # heads for aux categories:\n        self.nn_aux = torch.nn.ModuleList([\n            torch.nn.Linear(self.backbone_dim, n) for n in aux_classes\n        ])\n\n        self.automatic_optimization = False\n\n    def forward(self, x):\n        x = self.model(x)\n        cancer = self.nn_cancer(x).squeeze() # [batch_size]\n        aux = []\n        for nn in self.nn_aux:\n            aux.append(nn(x).squeeze()) # aux is list of [batch_size, n_classes] for n_classes is number of classes for each category.\n        return cancer, aux\n\n    def training_step(self, batch, batch_idx):\n        y_cancer_pred, aux_pred = self(batch[0])\n        cancer_loss = torch.nn.functional.binary_cross_entropy_with_logits(\n                        y_cancer_pred,\n                        batch[1].to(float),# check\n                        pos_weight=torch.tensor([CFG.POSITIVE_TARGET_WEIGHT]).to(DEVICE)\n                    ).item()\n        aux_loss = torch.mean(torch.stack([torch.nn.functional.cross_entropy(aux_pred[i], batch[2][:, i]) for i in range(batch[2].shape[-1])]))\n        loss = cancer_loss + CFG.AUX_LOSS_WEIGHT * aux_loss\n        \n        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)\n        self.log('cancer_loss', cancer_loss, on_epoch=True)\n        self.log('aux_loss', aux_loss, on_epoch=True)\n        return loss\n    \n    def validation_step(self, batch, batch_idx):\n        y_cancer, y_aux = batch[1], batch[2]\n        y_cancer_pred, aux_pred = self(batch[0])\n\n        cancer_loss = torch.nn.functional.binary_cross_entropy_with_logits(\n                        y_cancer_pred, \n                        y_cancer.to(float),\n                        # pos_weight=torch.tensor([CFG.POSITIVE_TARGET_WEIGHT])\n                        pos_weight=torch.tensor([CFG.POSITIVE_TARGET_WEIGHT]).to(DEVICE)\n                    ).item()\n        aux_loss = torch.mean(torch.stack([torch.nn.functional.cross_entropy(aux_pred[i], y_aux[:, i]) for i in range(y_aux.shape[-1])])).item()\n        val_loss = cancer_loss + CFG.AUX_LOSS_WEIGHT * aux_loss\n        self.log('val_loss', val_loss, on_step=True, on_epoch=True, prog_bar=True)\n        self.log('val_cancer_loss', cancer_loss, on_epoch=True)\n        self.log('val_aux_loss', aux_loss, on_epoch=True)\n        return val_loss\n\n    def predict_step(self, batch, batch_idx):\n        cancer, aux = self.forward(batch[0])\n        sigaux = []\n        for a in aux:\n            sigaux.append(torch.softmax(a, dim=-1))\n        return torch.sigmoid(cancer), sigaux # check\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.parameters())\n        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=CFG.ONE_CYCLE_MAX_LR, epochs=CFG.EPOCHS,\n                                                        steps_per_epoch=self.len_dl_train,\n                                                        pct_start=CFG.ONE_CYCLE_PCT_START)\n#         lr = scheduler.get_last_lr()[0] if scheduler else CFG.ONE_CYCLE_MAX_LR\n#         self.log('lr', lr)\n        return [optimizer], [scheduler]\n#         return optimizer\n\n#     def on_after_backward(self):\n#         global_step = self.global_step\n#         if (global_step % 25 == 0) and (global_step > 0):\n#             for name, param in self.named_parameters():\n#                 self.logger.experiment.add_histogram(name, param, global_step)\n#         wandb.watch(self.model, log = 'gradients')\n#         lr = scheduler.get_last_lr()[0] if scheduler else config.ONE_CYCLE_MAX_LR\n#         self.log('lr', lr)\n\nAUX_TARGET_NCLASSES = df[CATEGORY_AUX_TARGETS].max() + 1\nDEBUG = False\nif DEBUG:\n    modelModule = BSDRSNAModule(model_type=CFG.MODEL_TYPE, aux_classes=AUX_TARGET_NCLASSES, dropout=0., len_dl_train=1)\n    df_train_demo = df.iloc[:5]\n    df_valid_demo = df.iloc[5:10]\n    train_dataset = BreastCancerDataSet(df_train_demo, path=TRAIN_IMAGES_PATH, transforms= get_transforms(aug=True))\n    dataloader = DataLoader(train_dataset, batch_size=CFG.BATCH_SIZE, num_workers=2, shuffle=True, pin_memory=True)\n    trainer = pl.Trainer()\n    predictions = trainer.predict(modelModule, dataloader)\n","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.955356Z","iopub.execute_input":"2023-01-04T13:22:49.956078Z","iopub.status.idle":"2023-01-04T13:22:49.988104Z","shell.execute_reply.started":"2023-01-04T13:22:49.95604Z","shell.execute_reply":"2023-01-04T13:22:49.98716Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# predictions","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.989437Z","iopub.execute_input":"2023-01-04T13:22:49.990326Z","iopub.status.idle":"2023-01-04T13:22:49.994919Z","shell.execute_reply.started":"2023-01-04T13:22:49.990286Z","shell.execute_reply":"2023-01-04T13:22:49.993944Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5. Train","metadata":{}},{"cell_type":"code","source":"TRAIN = True\nif TRAIN:\n    for fold in FOLDS:\n        name = f\"{WANDB_RUN_NAME}-{fold}\"\n        \n        # each name for each fold\n        with wandb.init(project=WANDB_PROJECT, name=name, group=WANDB_RUN_NAME) as run:\n            gc.collect()\n            wandb_logger = WandbLogger()\n            train_df = df.query('split != @fold')\n            valid_df = df.query('split == @fold')\n            \n\n            data_module = BreastCancerDataModule(train_df, valid_df)\n\n            len_dl_train = len(train_df) // CFG.BATCH_SIZE + (1 if len(train_df) % CFG.BATCH_SIZE != 0 else 0)\n            # early_stopping_callback = EarlyStopping(monitor='val_loss', patience=2)\n            checkpoint_callback = ModelCheckpoint(\n                dirpath=CFG.CHECKPOINT_PATH,\n                filename= f\"fold_{fold}_{CFG.MODEL_TYPE}\",\n                save_top_k=1,\n                verbose=True,\n                monitor=\"val_loss\",\n                mode=\"min\"\n            )\n            lr_monitor = LearningRateMonitor(logging_interval='step')\n            \n            modelModule = BSDRSNAModule(model_type=CFG.MODEL_TYPE, aux_classes=AUX_TARGET_NCLASSES, dropout=0., len_dl_train=len_dl_train)\n#             wandb_logger.watch(modelModule, log='gradients')\n            trainer = pl.Trainer(\n                logger=wandb_logger,\n                callbacks=[checkpoint_callback, lr_monitor],\n                # callbacks=[checkpoint_callback, early_stopping_callback],\n                max_epochs=CFG.EPOCHS,\n                accelerator=\"gpu\", \n                devices=1, \n#                 strategy='dp',\n                log_every_n_steps=1,            \n                precision=16,\n                )\n\n            trainer.fit(modelModule, data_module)\n","metadata":{"execution":{"iopub.status.busy":"2023-01-04T13:22:49.996229Z","iopub.execute_input":"2023-01-04T13:22:49.997224Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6. WandB sweep","metadata":{}},{"cell_type":"markdown","source":"## 7. Cross Validation","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}