{"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":12038896,"sourceType":"datasetVersion","datasetId":7377931}],"dockerImageVersionId":31041,"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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-06-03T23:39:02.807855Z","iopub.execute_input":"2025-06-03T23:39:02.80833Z","iopub.status.idle":"2025-06-03T23:39:11.158372Z","shell.execute_reply.started":"2025-06-03T23:39:02.808299Z","shell.execute_reply":"2025-06-03T23:39:11.157596Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Caformer Improved Notebook\n\nThis is the final notebook I will release in this competition. I recommend reading the previous notebooks, as many components are reused.\n\n1. [HGNet-V2 - Starter](https://www.kaggle.com/code/brendanartley/hgnet-v2-starter)\n2. [ConvNeXt - Full Resolution Baseline](https://www.kaggle.com/code/brendanartley/convnext-full-resolution-baseline)\n3. [CAFormer - Full Resolution Improved](https://www.kaggle.com/code/brendanartley/caformer-full-resolution-improved)\n\n\nThis notebook features the strongest model I have found so far. A modified CAFormer encoder combined with an improved decoder (pixel shuffle, intermediate convolutions, and SCSE blocks).\n\nLike previous notebooks, I provide a pre-trained model checkpoint (trained for 150 epochs) which achieved a validation MAE of ~24.","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport tensorflow as tf\nfrom tensorflow.keras.layers import AveragePooling2D\n\nfeature_map = np.array([\n    [1, 1, 1, 2,2,2],\n    [1, 1, 1, 2,2,2],\n    [4, 4, 4, 3,3,3],\n    [4, 4, 4, 3,3,3]\n], dtype=np.float32).reshape(1, 4, 6, 1)  # Convert to float32\n\nprint(feature_map)\n\n# Applying average pooling\navg_pool = AveragePooling2D(pool_size=(2,3), strides=(2,3))\noutput = avg_pool(feature_map)\nprint(output.numpy().reshape(2,2))","metadata":{"trusted":true,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile _cfg.py\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\n# cfg.backbone = \"caformer_b36.sail_in22k_ft_in1k\"\ncfg.backbone = \"eva02_base_patch16_clip_224.merged2b_s8b_b131k\"\n#cfg.batch_size_val = 105\ncfg.ema = True\ncfg.ema_decay = 0.99\n\ncfg.precision = \"bfloat16\"  # or \"float32\" or bfloat16\ncfg.use_amp = True\n\ncfg.epochs = 60\ncfg.batch_size = 30\ncfg.numworkerstrain = 16\ncfg.batch_size_val = 30\ncfg.numworkersval = 16\n\n#cfg.batch_size = 85\n#cfg.batch_size_val = 85\n\ncfg.early_stopping = {\"patience\": 10, \"streak\": 5}\ncfg.logging_steps = 5","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T23:38:54.003396Z","iopub.execute_input":"2025-06-03T23:38:54.003742Z","iopub.status.idle":"2025-06-03T23:38:54.011811Z","shell.execute_reply.started":"2025-06-03T23:38:54.003723Z","shell.execute_reply":"2025-06-03T23:38:54.011032Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Scheduler\n\nI have removed the training loop from this notebook, though it is the same as previous notebooks. \n\nThe only difference was the use of a custom learning rate scheduler. The scheduler uses a constant learning rate followed by a cosine annealing learning rate. It seems that a learning rate of 1e-4 works well at the beggining, but a lower learning rate is required to achieve lower training and validation MAE.","metadata":{}},{"cell_type":"code","source":"!export MASTER_PORT=25565\n!export MASTER_ADDR=localhost\n!export RANK=1\n!export WORLD_SIZE=2","metadata":{"trusted":true,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%env MASTER_PORT = 6677","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:33:51.519275Z","iopub.execute_input":"2025-06-03T01:33:51.519903Z","iopub.status.idle":"2025-06-03T01:33:51.524295Z","shell.execute_reply.started":"2025-06-03T01:33:51.519872Z","shell.execute_reply":"2025-06-03T01:33:51.523507Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%env MASTER_ADDR=localhost","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:28:30.418372Z","iopub.execute_input":"2025-06-03T01:28:30.418668Z","iopub.status.idle":"2025-06-03T01:28:30.423133Z","shell.execute_reply.started":"2025-06-03T01:28:30.418649Z","shell.execute_reply":"2025-06-03T01:28:30.422258Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%env MASTER_ADDR=localhost\n%env MASTER_PORT = 29500\n\nimport monai\nimport os\nimport time \nimport random\nimport numpy as np\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nfrom torch.amp import autocast, GradScaler\n\nimport torch.optim.lr_scheduler\nimport torch.distributed as dist\nfrom torch.utils.data import DistributedSampler\nfrom torch.nn.parallel import DistributedDataParallel\n\nfrom _cfg import cfg\nfrom _dataset import CustomDataset\nfrom _model import ModelEMA, Net\nfrom _utils import format_time\nimport matplotlib.pyplot as plt\n\nimport math\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\ndef set_seed(seed=1234):\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = False\n    torch.backends.cudnn.benchmark = True\n\ndef get_dtype(precision):\n    \"\"\"Get torch dtype based on precision string\"\"\"\n    dtype_map = {\n        \"float32\": torch.float32,\n        \"bfloat16\": torch.bfloat16,\n        \"float16\": torch.float16,\n    }\n    return dtype_map.get(precision, torch.float32)\n\ndef setup(rank, world_size):\n    torch.cuda.set_device(rank)\n    dist.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n    return\n\ndef cleanup():\n    dist.barrier()\n    dist.destroy_process_group()\n    return\n\ndef main(cfg):\n    # ========== Setup Precision ==========\n    dtype = get_dtype(cfg.precision)\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # Check if bfloat16 is supported\n    if cfg.precision == \"bfloat16\" and not torch.cuda.is_bf16_supported():\n        print(\"Warning: BFloat16 not supported on this GPU. Falling back to float32.\")\n        cfg.precision = \"float32\"\n        dtype = torch.float32\n        cfg.use_amp = False\n    \n    print(f\"Training with precision: {cfg.precision}\")\n    print(f\"Automatic Mixed Precision: {'Enabled' if cfg.use_amp else 'Disabled'}\")\n\n    # ========== Datasets / Dataloaders ==========\n    if cfg.local_rank == 0:\n        print(\"=\"*25)\n        print(\"Loading data..\")\n    train_ds = CustomDataset(cfg=cfg, mode=\"train\")\n    sampler= DistributedSampler(\n        train_ds, \n        num_replicas=cfg.world_size, \n        rank=cfg.local_rank,\n    )\n    train_dl = torch.utils.data.DataLoader(\n        train_ds, \n        sampler= sampler,\n        batch_size= cfg.batch_size, \n        num_workers= cfg.numworkerstrain,\n    )\n    \n    valid_ds = CustomDataset(cfg=cfg, mode=\"valid\")\n    sampler= DistributedSampler(\n        valid_ds, \n        num_replicas=cfg.world_size, \n        rank=cfg.local_rank,\n    )\n    valid_dl = torch.utils.data.DataLoader(\n        valid_ds, \n        sampler= sampler,\n        batch_size= cfg.batch_size_val, \n        num_workers= cfg.numworkersval,\n    )\n\n    # ========== Model / Optim ==========\n    model = Net(backbone=cfg.backbone)\n    model= model.to(cfg.local_rank)\n    # Convert model to appropriate dtype\n    if cfg.precision != \"float32\":\n        model = model.to(dtype)\n    \n    model = model.to(device)\n    if cfg.ema:\n        if cfg.local_rank == 0:\n            print(\"Initializing EMA model..\")\n        ema_model = ModelEMA(\n            model, \n            decay=cfg.ema_decay, \n            device=cfg.local_rank,\n        )\n    else:\n        ema_model = None\n    model= DistributedDataParallel(\n        model, \n        device_ids=[cfg.local_rank], \n        )\n    \n    criterion = nn.L1Loss()\n    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) #1e-3\n    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.epochs)\n    scheduler = ConstantCosineLR(optimizer, total_steps=len(train_dl)+1, pct_cosine=0.5)\n\n\n    scaler = GradScaler()\n    if cfg.use_amp and cfg.precision == \"float16\":\n        scaler = GradScaler()\n    else:\n        scaler = None\n\n    # ========== Training ==========\n    if cfg.local_rank == 0:\n        print(\"=\"*25)\n        print(\"Give me warp {}, Mr. Sulu.\".format(cfg.world_size))\n        print(\"=\"*25)\n    \n    best_loss= 1_000_000\n    val_loss= 1_000_000\n    val_arr = []\n    epoch_arr = 0\n    for epoch in range(0, cfg.epochs+1):\n        if epoch != 0:\n            tstart= time.time()\n            train_dl.sampler.set_epoch(epoch)\n    \n            # Train loop\n            model.train()\n            total_loss = []\n            for i, (x, y) in enumerate(train_dl):\n                x = x.to(cfg.local_rank)\n                y = y.to(cfg.local_rank)\n                # Convert to appropriate dtype if needed\n                if cfg.precision != \"float32\":\n                    x = x.to(dtype)\n                    y = y.to(dtype)\n        \n                # Training step with precision handling\n                if cfg.use_amp and cfg.precision in [\"float16\", \"bfloat16\"]:\n                    with autocast(device_type='cuda', dtype=dtype):\n                        logits = model(x)\n                        loss = criterion(logits, y)\n                    \n                    if scaler is not None:  # float16 with GradScaler\n                        scaler.scale(loss).backward()\n                        scaler.unscale_(optimizer)\n                        torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)\n                        scaler.step(optimizer)\n                        scaler.update()\n                    else:  # bfloat16 without GradScaler\n                        loss.backward()\n                        torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)\n                        optimizer.step()\n                else:  # float32\n                    logits = model(x)\n                    loss = criterion(logits, y)\n                    loss.backward()\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)\n                    optimizer.step()\n                \n                optimizer.zero_grad()\n                total_loss.append(loss.item())\n                \n                if ema_model is not None:\n                    ema_model.update(model)\n                    \n                if cfg.local_rank == 0 and (len(total_loss) >= cfg.logging_steps or i == 0):\n                    train_loss = np.mean(total_loss)\n                    total_loss = []\n                    print(\"Epoch {}:     Train MAE: {:.2f}     Val MAE: {:.2f}     Time: {}     Step: {}/{}\".format(\n                        epoch, \n                        train_loss,\n                        val_loss,\n                        format_time(time.time() - tstart),\n                        i+1, \n                        len(train_dl)+1, \n                    ))\n    \n        # ========== Valid ==========\n        model.eval()\n        val_logits = []\n        val_targets = []\n        with torch.no_grad():\n            for x, y in tqdm(valid_dl, disable=cfg.local_rank != 0):\n                x = x.to(cfg.local_rank)\n                y = y.to(cfg.local_rank)\n    \n                # Convert to appropriate dtype if needed\n                if cfg.precision != \"float32\":\n                    x = x.to(dtype)\n                    y = y.to(dtype)\n    \n                if cfg.use_amp and cfg.precision in [\"float16\", \"bfloat16\"]:\n                    with autocast(device_type='cuda', dtype=dtype):\n                        if ema_model is not None:\n                            out = ema_model.module(x)\n                        else:\n                            out = model(x)\n                        batch_loss = criterion(out, y)\n                else:\n                    if ema_model is not None:\n                        out = ema_model.module(x)\n                    else:\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            loss = criterion(val_logits, val_targets).item()\n\n        # Gather loss\n        v = torch.tensor([loss], device=cfg.local_rank)\n        torch.distributed.all_reduce(v, op=dist.ReduceOp.SUM)\n        val_loss = (v[0] / cfg.world_size).item()\n        \n        # val_arr.append(val_loss)\n        # epoch_arr = range(1,epoch+1)\n        # plt.plot(epoch_arr,val_arr)\n        # plt.xlabel(\"Epoch\")\n        # plt.ylabel(\"VAL MAE\")\n        # plt.title(\"ConstantCosineLR\")\n        # plt.savefig('EPOCHVSVALMAE.png')\n        \n        \n        # ========== Weights / Early stopping ==========\n        stop_train = torch.tensor([0], device=cfg.local_rank)\n        if cfg.local_rank == 0:\n            es= cfg.early_stopping\n            if val_loss < best_loss:\n                print(\"New best: {:.2f} -> {:.2f}\".format(best_loss, val_loss))\n                print(\"Saved weights..\")\n                best_loss = val_loss\n                if ema_model is not None:\n                    torch.save(ema_model.module.state_dict(), f'best_model_{cfg.seed}.pt')\n                else:\n                    torch.save(model.state_dict(), f'best_model_{cfg.seed}.pt')\n        \n                es[\"streak\"] = 0\n            else:\n                es= cfg.early_stopping\n                es[\"streak\"] += 1\n                if es[\"streak\"] > es[\"patience\"]:\n                    print(\"Ending training (early_stopping).\")\n                    stop_train = torch.tensor([1], device=cfg.local_rank)\n        \n        # Exits training on all ranks\n        dist.broadcast(stop_train, src=0)\n        if stop_train.item() == 1:\n            return\n        \n        scheduler.step(epoch)  \n\n\n    return\n    \n\n\nif __name__ == \"__main__\":\n\n    # GPU Specs\n    rank = 1\n    world_size = 2\n    _, total = torch.cuda.mem_get_info(device=rank)\n\n    # Init\n    setup(rank, world_size)\n    time.sleep(rank)\n    print(f\"Rank: {rank}, World size: {world_size}, GPU memory: {total / 1024**3:.2f}GB\", flush=True)\n    time.sleep(world_size - rank)\n\n    # Seed\n    set_seed(cfg.seed+rank)\n\n    # Run\n    cfg.local_rank= rank\n    cfg.world_size= world_size\n    main(cfg)\n    cleanup()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:38:33.232621Z","iopub.execute_input":"2025-06-03T01:38:33.233354Z","iopub.status.idle":"2025-06-03T01:59:33.368378Z","shell.execute_reply.started":"2025-06-03T01:38:33.233321Z","shell.execute_reply":"2025-06-03T01:59:33.36699Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile _utils.py\nimport datetime\n\ndef format_time(elapsed):\n    elapsed_rounded = int(round((elapsed)))\n    return str(datetime.timedelta(seconds=elapsed_rounded))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-03T01:36:44.768523Z","iopub.execute_input":"2025-06-03T01:36:44.768833Z","iopub.status.idle":"2025-06-03T01:36:44.773569Z","shell.execute_reply.started":"2025-06-03T01:36:44.7688Z","shell.execute_reply":"2025-06-03T01:36:44.773031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile .env\nMASTER_PORT=25565\nMASTER_ADDR=localhost","metadata":{"trusted":true,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"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","metadata":{"trusted":true,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import 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()\n","metadata":{"trusted":true,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Dataset\n\nSame as previous notebook. ","metadata":{}},{"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,"execution":{"iopub.status.busy":"2025-06-03T01:36:50.940957Z","iopub.execute_input":"2025-06-03T01:36:50.941497Z","iopub.status.idle":"2025-06-03T01:36:50.947551Z","shell.execute_reply.started":"2025-06-03T01:36:50.941472Z","shell.execute_reply":"2025-06-03T01:36:50.946673Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model\n\nThis time we use the `CAFormer` backbone from timm. See more info on this backbone [here](https://huggingface.co/timm/caformer_b36.sail_in22k_ft_in1k) and the original paper [here](https://arxiv.org/abs/2210.13452).\n\n\n### Encoder\n\nLike with Convnext, we modify the encoder so that the feature maps are aligned with the target output shape. I think there is room for improvement at the `nn.ReflectionPad2d` step. Currently, the model uses lots of padding here and I am afraid the detail in the shallowest feature map is lacking.\n\n### Decoder\n\nThe biggest changes in this notebook are to the decoder. \n\nFirst, we use PixelShuffle for upsampling. Pixelshuffle typically works well when fine detail is important, though it is more computatially expensive. Second, we add SCSE blocks. These are commonly used to increase decoder capacity with a minimal increase in parameter count and runtime. Finally, we add intermediate convolutions between the encoder output and decoder blocks. I beleive this trick was first introduced on Kaggle in the 3rd place solution of the Contrails Competition [here](https://www.kaggle.com/competitions/google-research-identify-contrails-redu), and also increases decoder capacity.","metadata":{}},{"cell_type":"code","source":"%%writefile _model.py\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,2) #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,"execution":{"iopub.status.busy":"2025-06-03T23:51:11.081778Z","iopub.execute_input":"2025-06-03T23:51:11.082333Z","iopub.status.idle":"2025-06-03T23:51:11.098778Z","shell.execute_reply.started":"2025-06-03T23:51:11.082308Z","shell.execute_reply":"2025-06-03T23:51:11.09807Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Utils\n\nSame as previous notebook. ","metadata":{}},{"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,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Pretrained Model\n\nNext, we load in the pretrained model. This model was trained with a batch size of 16 for 200 epochs.","metadata":{}},{"cell_type":"code","source":"import 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,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Valid\n\nNext, we score the ensemble on the validation set.","metadata":{}},{"cell_type":"code","source":"from 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,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Test\n\nFinally, we make predictions on the test data.","metadata":{}},{"cell_type":"code","source":"import 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,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import 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,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"We can also view a few samples to make sure things look reasonable.","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt \n\nif RUN_TEST:\n    # Plot a few samples\n    fig, axes = plt.subplots(3, 5, figsize=(10, 6))\n    axes= axes.flatten()\n\n    n = min(len(outputs), len(axes))\n    \n    for i in range(n):\n        img= outputs[0, 0, ...].cpu().numpy()\n        img = outputs[i, 0].cpu().numpy()\n        idx= oids_test[i]\n    \n        # Plot\n        axes[i].imshow(img, cmap='gray')\n        axes[i].set_title(idx)\n        axes[i].axis('off')\n\n    for i in range(n, len(axes)):\n        axes[i].axis('off')\n    \n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"execution_failed":"2025-06-03T01:27:49.648Z"}},"outputs":[],"execution_count":null}]}