{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"},{"sourceId":11568812,"sourceType":"datasetVersion","datasetId":7253205},{"sourceId":11569667,"sourceType":"datasetVersion","datasetId":7253605},{"sourceId":11569755,"sourceType":"datasetVersion","datasetId":7253661},{"sourceId":11910577,"sourceType":"datasetVersion","datasetId":7487776},{"sourceId":12038896,"sourceType":"datasetVersion","datasetId":7377931}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"RUN_VALID = True\nRUN_TEST  = True\n\nimport torch\nif not torch.cuda.is_available() or torch.cuda.device_count() < 2:\n    raise RuntimeError(\"Requires >= 2 GPUs with CUDA enabled.\")\n\ntry: \n    import monai\nexcept: \n    !pip install --no-deps monai -q","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile _cfg.py\n\nfrom types import SimpleNamespace\nimport torch\n\ncfg= SimpleNamespace()\ncfg.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ncfg.local_rank = 0\ncfg.seed = 123\ncfg.subsample = None\n\ncfg.backbone = \"caformer_b36.sail_in22k_ft_in1k\"\ncfg.batch_size_val = 16","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import math\n\nfrom torch.optim.lr_scheduler import _LRScheduler\n\nclass ConstantCosineLR(_LRScheduler):\n    \"\"\"\n    Constant learning rate followed by CosineAnnealing.\n    \"\"\"\n    def __init__(\n        self, \n        optimizer,\n        total_steps, \n        pct_cosine, \n        last_epoch=-1,\n        ):\n        self.total_steps = total_steps\n        self.milestone = int(total_steps * (1 - pct_cosine))\n        self.cosine_steps = max(total_steps - self.milestone, 1)\n        self.min_lr = 0\n        super().__init__(optimizer, last_epoch)\n\n    def get_lr(self):\n        step = self.last_epoch + 1\n        if step <= self.milestone:\n            factor = 1.0\n        else:\n            s = step - self.milestone\n            factor = 0.5 * (1 + math.cos(math.pi * s / self.cosine_steps))\n        return [lr * factor for lr in self.base_lrs]\n\nimport torch\nimport matplotlib.pyplot as plt\n\n# Dummy model\nn_steps = 10_000\nmodel = torch.nn.Linear(1, 1)\noptimizer = torch.optim.SGD(model.parameters(), lr=1.0)\n\n# Scheduler\nscheduler = ConstantCosineLR(optimizer, total_steps=n_steps, pct_cosine=0.5)\n\n# Get LRs\narr = []\nfor _ in range(n_steps):\n    scheduler.step()\n    arr.append(optimizer.param_groups[0]['lr'])\n\nplt.plot(arr)\nplt.xlabel(\"Step\")\nplt.ylabel(\"LR\")\nplt.title(\"ConstantCosineLR\")\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile _dataset.py\n\nimport os\nimport glob\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nimport torch\n\nclass CustomDataset(torch.utils.data.Dataset):\n    def __init__(\n        self, \n        cfg,\n        mode = \"train\", \n    ):\n        self.cfg = cfg\n        self.mode = mode\n        \n        self.data, self.labels, self.records = self.load_metadata()\n\n    def load_metadata(self, ):\n\n        # Select rows\n        df= pd.read_csv(\"/kaggle/input/openfwi-preprocessed-72x72/folds.csv\")\n        if self.cfg.subsample is not None:\n            df= df.groupby([\"dataset\", \"fold\"]).head(self.cfg.subsample)\n\n        if self.mode == \"train\":\n            df= df[df[\"fold\"] != 0]\n        else:\n            df= df[df[\"fold\"] == 0]\n\n        \n        data = []\n        labels = []\n        records = []\n        mmap_mode = \"r\"\n\n        for idx, row in tqdm(df.iterrows(), total=len(df), disable=self.cfg.local_rank != 0):\n            row= row.to_dict()\n\n            # Hacky way to get exact file name\n            p1 = os.path.join(\"/kaggle/input/open-wfi-1/openfwi_float16_1/\", row[\"data_fpath\"])\n            p2 = os.path.join(\"/kaggle/input/open-wfi-1/openfwi_float16_1/\", row[\"data_fpath\"].split(\"/\")[0], \"*\", row[\"data_fpath\"].split(\"/\")[-1])\n            p3 = os.path.join(\"/kaggle/input/open-wfi-2/openfwi_float16_2/\", row[\"data_fpath\"])\n            p4 = os.path.join(\"/kaggle/input/open-wfi-2/openfwi_float16_2/\", row[\"data_fpath\"].split(\"/\")[0], \"*\", row[\"data_fpath\"].split(\"/\")[-1])\n            farr= glob.glob(p1) + glob.glob(p2) + glob.glob(p3) + glob.glob(p4)\n        \n            # Map to lbl fpath\n            farr= farr[0]\n            flbl= farr.replace('seis', 'vel').replace('data', 'model')\n            \n            # Load\n            arr= np.load(farr, mmap_mode=mmap_mode)\n            lbl= np.load(flbl, mmap_mode=mmap_mode)\n\n            # Append\n            data.append(arr)\n            labels.append(lbl)\n            records.append(row[\"dataset\"])\n\n        return data, labels, records\n\n    def __getitem__(self, idx):\n        row_idx= idx // 500\n        col_idx= idx % 500\n\n        d= self.records[row_idx]\n        x= self.data[row_idx][col_idx, ...]\n        y= self.labels[row_idx][col_idx, ...]\n\n        # Augs \n        if self.mode == \"train\":\n            \n            # Temporal flip\n            if np.random.random() < 0.5:\n                x= x[::-1, :, ::-1]\n                y= y[..., ::-1]\n\n        x= x.copy()\n        y= y.copy()\n        \n        return x, y\n\n    def __len__(self, ):\n        return len(self.records) * 500","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Model\n\nfrom copy import deepcopy\nfrom types import MethodType\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nimport timm\n\nfrom monai.networks.blocks import UpSample, SubpixelUpsample\n\n####################\n## EMA + Ensemble ##\n####################\n\nclass ModelEMA(nn.Module):\n    def __init__(self, model, decay=0.99, device=None):\n        super().__init__()\n        self.module = deepcopy(model)\n        self.module.eval()\n        self.decay = decay\n        self.device = device\n        if self.device is not None:\n            self.module.to(device=device)\n\n    def _update(self, model, update_fn):\n        with torch.no_grad():\n            for ema_v, model_v in zip(self.module.state_dict().values(), model.state_dict().values()):\n                if self.device is not None:\n                    model_v = model_v.to(device=self.device)\n                ema_v.copy_(update_fn(ema_v, model_v))\n\n    def update(self, model):\n        self._update(model, update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m)\n\n    def set(self, model):\n        self._update(model, update_fn=lambda e, m: m)\n\n\nclass EnsembleModel(nn.Module):\n    def __init__(self, models):\n        super().__init__()\n        self.models = nn.ModuleList(models).eval()\n\n    def forward(self, x):\n        output = None\n        \n        for m in self.models:\n            logits= m(x)\n            \n            if output is None:\n                output = logits\n            else:\n                output += logits\n                \n        output /= len(self.models)\n        return output\n        \n\n#############\n## Decoder ##\n#############\n\nclass ConvBnAct2d(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size,\n        padding: int = 0,\n        stride: int = 1,\n        norm_layer: nn.Module = nn.Identity,\n        act_layer: nn.Module = nn.ReLU,\n    ):\n        super().__init__()\n\n        self.conv= nn.Conv2d(\n            in_channels, \n            out_channels,\n            kernel_size,\n            stride=stride, \n            padding=padding, \n            bias=False,\n        )\n        self.norm = norm_layer(out_channels) if norm_layer != nn.Identity else nn.Identity()\n        self.act= act_layer(inplace=True)\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.norm(x)\n        x = self.act(x)\n        return x\n\n\nclass SCSEModule2d(nn.Module):\n    def __init__(self, in_channels, reduction=16):\n        super().__init__()\n        self.cSE = nn.Sequential(\n            nn.AdaptiveAvgPool2d(1),\n            nn.Conv2d(in_channels, in_channels // reduction, 1),\n            nn.Tanh(),\n            nn.Conv2d(in_channels // reduction, in_channels, 1),\n            nn.Sigmoid(),\n        )\n        self.sSE = nn.Sequential(\n            nn.Conv2d(in_channels, 1, 1), \n            nn.Sigmoid(),\n            )\n\n    def forward(self, x):\n        return x * self.cSE(x) + x * self.sSE(x)\n\nclass Attention2d(nn.Module):\n    def __init__(self, name, **params):\n        super().__init__()\n        if name is None:\n            self.attention = nn.Identity(**params)\n        elif name == \"scse\":\n            self.attention = SCSEModule2d(**params)\n        else:\n            raise ValueError(\"Attention {} is not implemented\".format(name))\n\n    def forward(self, x):\n        return self.attention(x)\n\nclass DecoderBlock2d(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        skip_channels,\n        out_channels,\n        norm_layer: nn.Module = nn.Identity,\n        attention_type: str = None,\n        intermediate_conv: bool = False,\n        upsample_mode: str = \"deconv\",\n        scale_factor: int = 2,\n    ):\n        super().__init__()\n\n        # Upsample block\n        if upsample_mode == \"pixelshuffle\":\n            self.upsample= SubpixelUpsample(\n                spatial_dims= 2,\n                in_channels= in_channels,\n                scale_factor= scale_factor,\n            )\n        else:\n            self.upsample = UpSample(\n                spatial_dims= 2,\n                in_channels= in_channels,\n                out_channels= in_channels,\n                scale_factor= scale_factor,\n                mode= upsample_mode,\n            )\n\n        if intermediate_conv:\n            k= 3\n            c= skip_channels if skip_channels != 0 else in_channels\n            self.intermediate_conv = nn.Sequential(\n                ConvBnAct2d(c, c, k, k//2),\n                ConvBnAct2d(c, c, k, k//2),\n                )\n        else:\n            self.intermediate_conv= None\n\n        self.attention1 = Attention2d(\n            name= attention_type, \n            in_channels= in_channels + skip_channels,\n            )\n\n        self.conv1 = ConvBnAct2d(\n            in_channels + skip_channels,\n            out_channels,\n            kernel_size= 3,\n            padding= 1,\n            norm_layer= norm_layer,\n        )\n\n        self.conv2 = ConvBnAct2d(\n            out_channels,\n            out_channels,\n            kernel_size= 3,\n            padding= 1,\n            norm_layer= norm_layer,\n        )\n        self.attention2 = Attention2d(\n            name= attention_type, \n            in_channels= out_channels,\n            )\n\n    def forward(self, x, skip=None):\n        x = self.upsample(x)\n\n        if self.intermediate_conv is not None:\n            if skip is not None:\n                skip = self.intermediate_conv(skip)\n            else:\n                x = self.intermediate_conv(x)\n\n        if skip is not None:\n            # print(x.shape, skip.shape)\n            x = torch.cat([x, skip], dim=1)\n            x = self.attention1(x)\n\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x = self.attention2(x)\n        return x\n\n\nclass UnetDecoder2d(nn.Module):\n    \"\"\"\n    Unet decoder.\n    Source: https://arxiv.org/abs/1505.04597\n    \"\"\"\n    def __init__(\n        self,\n        encoder_channels: tuple[int],\n        skip_channels: tuple[int] = None,\n        decoder_channels: tuple = (256, 128, 64, 32),\n        scale_factors: tuple = (2,2,2,2),\n        norm_layer: nn.Module = nn.Identity,\n        attention_type: str = \"scse\",\n        intermediate_conv: bool = True,\n        upsample_mode: str = \"pixelshuffle\",\n    ):\n        super().__init__()\n        \n        if len(encoder_channels) == 4:\n            decoder_channels= decoder_channels[1:]\n        self.decoder_channels= decoder_channels\n        \n        if skip_channels is None:\n            skip_channels= list(encoder_channels[1:]) + [0]\n\n        # Build decoder blocks\n        in_channels= [encoder_channels[0]] + list(decoder_channels[:-1])\n        self.blocks = nn.ModuleList()\n\n        for i, (ic, sc, dc) in enumerate(zip(in_channels, skip_channels, decoder_channels)):\n            # print(i, ic, sc, dc)\n            self.blocks.append(\n                DecoderBlock2d(\n                    ic, sc, dc, \n                    norm_layer= norm_layer,\n                    attention_type= attention_type,\n                    intermediate_conv= intermediate_conv,\n                    upsample_mode= upsample_mode,\n                    scale_factor= scale_factors[i],\n                    )\n            )\n\n    def forward(self, feats: list[torch.Tensor]):\n        res= [feats[0]]\n        feats= feats[1:]\n\n        # Decoder blocks\n        for i, b in enumerate(self.blocks):\n            skip= feats[i] if i < len(feats) else None\n            res.append(\n                b(res[-1], skip=skip),\n                )\n            \n        return res\n\nclass SegmentationHead2d(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        scale_factor: tuple[int] = (2,2),\n        kernel_size: int = 3,\n        mode: str = \"nontrainable\",\n    ):\n        super().__init__()\n        self.conv= nn.Conv2d(\n            in_channels, out_channels, kernel_size= kernel_size,\n            padding= kernel_size//2\n        )\n        self.upsample = UpSample(\n            spatial_dims= 2,\n            in_channels= out_channels,\n            out_channels= out_channels,\n            scale_factor= scale_factor,\n            mode= mode,\n        )\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.upsample(x)\n        return x\n        \n\n#############\n## Encoder ##\n#############\n\nclass Net(nn.Module):\n    def __init__(\n        self,\n        backbone: str,\n        pretrained: bool = True,\n    ):\n        super().__init__()\n        \n        # Encoder\n        self.backbone= timm.create_model(\n            backbone,\n            in_chans= 5,\n            pretrained= pretrained,\n            features_only= True,\n            drop_path_rate=0.0,\n            )\n        ecs= [_[\"num_chs\"] for _ in self.backbone.feature_info][::-1]\n\n        # Decoder\n        self.decoder= UnetDecoder2d(\n            encoder_channels= ecs,\n        )\n\n        self.seg_head= SegmentationHead2d(\n            in_channels= self.decoder.decoder_channels[-1],\n            out_channels= 1,\n            scale_factor= 1,\n        )\n        \n        self._update_stem(backbone)\n\n    def _update_stem(self, backbone):\n        m = self.backbone\n\n        m.stem.conv.stride=(4,1)\n        m.stem.conv.padding=(0,4)\n        m.stages_0.downsample = nn.AvgPool2d(kernel_size=(4,1), stride=(4,1))\n        m.stem= nn.Sequential(\n            nn.ReflectionPad2d((0,0,78,78)),\n            m.stem,\n        )\n\n        pass\n\n        \n    def proc_flip(self, x_in):\n        x_in= torch.flip(x_in, dims=[-3, -1])\n        x= self.backbone(x_in)\n        x= x[::-1]\n\n        # Decoder\n        x= self.decoder(x)\n        x_seg= self.seg_head(x[-1])\n        x_seg= x_seg[..., 1:-1, 1:-1]\n        x_seg= torch.flip(x_seg, dims=[-1])\n        x_seg= x_seg * 1500 + 3000\n        return x_seg\n\n    def forward(self, batch):\n        x= batch\n\n        # Encoder\n        x_in = x\n        x= self.backbone(x)\n        # print([_.shape for _ in x])\n        x= x[::-1]\n\n        # Decoder\n        x= self.decoder(x)\n        # print([_.shape for _ in x])\n        x_seg= self.seg_head(x[-1])\n        x_seg= x_seg[..., 1:-1, 1:-1]\n        x_seg= x_seg * 1500 + 3000\n    \n        if self.training:\n            return x_seg\n        else:\n            p1 = self.proc_flip(x_in)\n            x_seg = torch.mean(torch.stack([x_seg, p1]), dim=0)\n            return x_seg","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import datetime\n\ndef format_time(elapsed):\n    elapsed_rounded = int(round((elapsed)))\n    return str(datetime.timedelta(seconds=elapsed_rounded))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Pretrained model\n\nimport glob\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom _cfg import cfg\n\nif RUN_VALID or RUN_TEST:\n\n    # Load pretrained models\n    models = []\n    for f in sorted(glob.glob(\"/kaggle/input/openfwi-preprocessed-72x72/models_1000x70/unet2d_caformer*.pt\")):\n        print(\"Loading: \", f)\n        m = Net(\n            backbone=cfg.backbone,\n            pretrained=False,\n        )\n        state_dict= torch.load(f, map_location=cfg.device, weights_only=True)\n        state_dict= {k.removeprefix(\"_orig_mod.\"):v for k,v in state_dict.items()} # Remove torch.compile() prefix\n\n        m.load_state_dict(state_dict)\n        models.append(m)\n    \n    # Combine\n    model = EnsembleModel(models)\n    model = model.to(cfg.device)\n    model = model.eval()\n    print(\"n_models: {:_}\".format(len(models)))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Validation\n\nfrom tqdm import tqdm\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nfrom torch.amp import autocast\n\nfrom _dataset import CustomDataset\n\n\nif RUN_VALID:\n\n    # Dataset / Dataloader\n    valid_ds = CustomDataset(cfg=cfg, mode=\"valid\")\n    sampler = torch.utils.data.SequentialSampler(valid_ds)\n    valid_dl = torch.utils.data.DataLoader(\n        valid_ds, \n        sampler= sampler,\n        batch_size= cfg.batch_size_val, \n        num_workers= 4,\n    )\n\n    # Valid loop\n    criterion = nn.L1Loss()\n    val_logits = []\n    val_targets = []\n    \n    with torch.no_grad():\n        for x, y in tqdm(valid_dl):\n            x = x.to(cfg.device)\n            y = y.to(cfg.device)\n    \n            with autocast(cfg.device.type):\n                out = model(x)\n    \n            val_logits.append(out.cpu())\n            val_targets.append(y.cpu())\n    \n        val_logits= torch.cat(val_logits, dim=0)\n        val_targets= torch.cat(val_targets, dim=0)\n    \n        total_loss= criterion(val_logits, val_targets).item()\n    \n    # Dataset Scores\n    ds_idxs= np.array([valid_ds.records])\n    ds_idxs= np.repeat(ds_idxs, repeats=500)\n    \n    print(\"=\"*25)\n    with torch.no_grad():    \n        for idx in sorted(np.unique(ds_idxs)):\n    \n            # Mask\n            mask = ds_idxs == idx\n            logits_ds = val_logits[mask]\n            targets_ds = val_targets[mask]\n    \n            # Score predictions\n            loss = criterion(val_logits[mask], val_targets[mask]).item()\n            print(\"{:15} {:.2f}\".format(idx, loss))\n    print(\"=\"*25)\n    print(\"Val MAE: {:.2f}\".format(total_loss))\n    print(\"=\"*25)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Predictions\n\nimport torch\n\nclass TestDataset(torch.utils.data.Dataset):\n    def __init__(self, test_files):\n        self.test_files = test_files\n\n    def __len__(self):\n        return len(self.test_files)\n\n    def __getitem__(self, i):\n        test_file = self.test_files[i]\n        test_stem = test_file.split(\"/\")[-1].split(\".\")[0]\n        return np.load(test_file), test_stem","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Submission file\n\nimport csv\nimport time\nimport glob\nfrom tqdm import tqdm\nimport numpy as np\nimport pandas as pd\n\n\nif RUN_TEST:\n\n    ss= pd.read_csv(\"/kaggle/input/waveform-inversion/sample_submission.csv\")    \n    row_count = 0\n    t0 = time.time()\n    \n    test_files = sorted(glob.glob(\"/kaggle/input/open-wfi-test/test/*.npy\"))\n    x_cols = [f\"x_{i}\" for i in range(1, 70, 2)]\n    fieldnames = [\"oid_ypos\"] + x_cols\n    \n    test_ds = TestDataset(test_files)\n    test_dl = torch.utils.data.DataLoader(\n        test_ds, \n        sampler=torch.utils.data.SequentialSampler(test_ds),\n        batch_size=cfg.batch_size_val, \n        num_workers=4,\n    )\n    \n    with open(\"submission.csv\", \"wt\", newline=\"\") as csvfile:\n        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n        writer.writeheader()\n\n        with torch.inference_mode():\n            with torch.autocast(cfg.device.type):\n                for inputs, oids_test in tqdm(test_dl, total=len(test_dl)):\n                    inputs = inputs.to(cfg.device)\n            \n                    outputs = model(inputs)\n                            \n                    y_preds = outputs[:, 0].cpu().numpy()\n                    \n                    for y_pred, oid_test in zip(y_preds, oids_test):\n                        for y_pos in range(70):\n                            row = dict(zip(x_cols, [y_pred[y_pos, x_pos] for x_pos in range(1, 70, 2)]))\n                            row[\"oid_ypos\"] = f\"{oid_test}_y_{y_pos}\"\n            \n                            writer.writerow(row)\n                            row_count += 1\n\n                            # Clear buffer\n                            if row_count % 100_000 == 0:\n                                csvfile.flush()\n    \n    t1 = format_time(time.time() - t0)\n    print(f\"Inference Time: {t1}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}