{"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":12345893,"sourceType":"datasetVersion","datasetId":7783028}],"dockerImageVersionId":31041,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"from pathlib import Path\n\nimport timm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport matplotlib.pyplot as plt","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:42:42.392802Z","iopub.execute_input":"2025-07-02T10:42:42.393088Z","iopub.status.idle":"2025-07-02T10:42:56.054834Z","shell.execute_reply.started":"2025-07-02T10:42:42.393040Z","shell.execute_reply":"2025-07-02T10:42:56.054099Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"markdown","source":"## Note\n\nMy submission used `erase_last_row=False`, which is not consistent with my training process.\nFixing the issue by `erase_last_row=True` gives (Public 0.796, Private 0.799).","metadata":{}},{"cell_type":"code","source":"class TestDataset(torch.utils.data.Dataset):\n    def __init__(self, data_root, erase_last_row=False):\n        self.data_root = Path(data_root)\n        self.seis_files = sorted(self.data_root.glob(\"*.npy\"))\n        self.erase_last_row = erase_last_row\n\n    def __len__(self):\n        return len(self.seis_files)\n\n    def __getitem__(self, idx):\n        seis_file = self.seis_files[idx]\n        seis = np.load(seis_file)\n        seis = torch.from_numpy(seis)\n        if self.erase_last_row:\n            seis[:, -1, :] = 0\n        oid = Path(seis_file).stem\n        return {\"seis\": seis, \"oid\": oid}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:42:56.057180Z","iopub.execute_input":"2025-07-02T10:42:56.057543Z","iopub.status.idle":"2025-07-02T10:42:56.063277Z","shell.execute_reply.started":"2025-07-02T10:42:56.057522Z","shell.execute_reply":"2025-07-02T10:42:56.062451Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_ds = TestDataset(\"/kaggle/input/waveform-inversion/test\", erase_last_row=True)\ntest_loader = torch.utils.data.DataLoader(test_ds, batch_size=32, shuffle=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:42:56.064224Z","iopub.execute_input":"2025-07-02T10:42:56.064500Z","iopub.status.idle":"2025-07-02T10:42:58.802484Z","shell.execute_reply.started":"2025-07-02T10:42:56.064472Z","shell.execute_reply":"2025-07-02T10:42:58.801891Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"def process_input(x):\n    # x: (N, 5, 1000, 70)\n    N = x.size(0)\n    x = x.reshape(N, 5, 250, 4, 70)\n    x = x.permute(0, 1, 2, 4, 3)\n    x = x.reshape(N, 5, 250, 280)\n    x = F.interpolate(x, size=(490, 490), mode=\"bilinear\", align_corners=False)\n    return x\n\n\nclass FWIModel(nn.Module):\n    def __init__(self, pretrained=True, split_at=4):\n        super().__init__()\n        backbone = timm.create_model(\n            \"eva02_large_patch14_clip_224\",\n            pretrained=pretrained,\n            dynamic_img_size=True,\n            in_chans=1,\n        )\n        self.backbone = backbone\n        self.head = nn.Linear(1024, 4)\n        self.pixel_shuffle = nn.PixelShuffle(2)\n        self.split_at = split_at\n\n    def forward(self, x):\n        x = process_input(x)\n        N, C, H, W = x.size()\n        # x = self.backbone.forward_features(x)\n        x = self.backbone.patch_embed(x.reshape(N * C, 1, H, W))\n        x, rot_pos_embed = self.backbone._pos_embed(x)\n        for blk in self.backbone.blocks[: self.split_at]:\n            x = blk(x, rope=rot_pos_embed)\n        x = x.reshape(N, 5, *x.size()[1:])\n        x = x.mean(1)\n        for blk in self.backbone.blocks[self.split_at :]:\n            x = blk(x, rope=rot_pos_embed)\n        x = self.backbone.norm(x)\n\n        x = x[:, 1:]\n        x = self.head(x)\n        x = x.permute(0, 2, 1).reshape(N, 4, 35, 35)\n        x = self.pixel_shuffle(x)\n        # x = x[:, :, 1:-1, 1:-1]\n\n        x = F.sigmoid(x.float()) * 6000\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:42:58.803209Z","iopub.execute_input":"2025-07-02T10:42:58.803425Z","iopub.status.idle":"2025-07-02T10:42:58.811966Z","shell.execute_reply.started":"2025-07-02T10:42:58.803408Z","shell.execute_reply":"2025-07-02T10:42:58.811239Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = FWIModel(pretrained=False, split_at=8)\nmodel.load_state_dict(torch.load(\"/kaggle/input/fwi-checkpoints/release03_epoch1.pth\", \"cpu\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:42:58.812791Z","iopub.execute_input":"2025-07-02T10:42:58.813856Z","iopub.status.idle":"2025-07-02T10:43:17.488308Z","shell.execute_reply.started":"2025-07-02T10:42:58.813835Z","shell.execute_reply":"2025-07-02T10:43:17.487573Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = model.cuda().eval()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:43:17.488949Z","iopub.execute_input":"2025-07-02T10:43:17.489188Z","iopub.status.idle":"2025-07-02T10:43:18.228602Z","shell.execute_reply.started":"2025-07-02T10:43:17.489170Z","shell.execute_reply":"2025-07-02T10:43:18.227958Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference","metadata":{"execution":{"iopub.status.busy":"2025-07-02T02:54:03.724284Z","iopub.execute_input":"2025-07-02T02:54:03.724781Z","iopub.status.idle":"2025-07-02T02:54:03.804111Z","shell.execute_reply.started":"2025-07-02T02:54:03.724755Z","shell.execute_reply":"2025-07-02T02:54:03.803407Z"}}},{"cell_type":"code","source":"data = next(iter(test_loader))\nwith torch.no_grad(), torch.amp.autocast(device_type=\"cuda\", dtype=torch.float16):\n    seis = data[\"seis\"].cuda()\n    pred_vel = model(seis)\npred_vel = pred_vel.cpu().float()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:43:18.230180Z","iopub.execute_input":"2025-07-02T10:43:18.230392Z","iopub.status.idle":"2025-07-02T10:43:27.030288Z","shell.execute_reply.started":"2025-07-02T10:43:18.230375Z","shell.execute_reply":"2025-07-02T10:43:27.029726Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"idx = 20\nprint(data[\"oid\"][idx])\nplt.imshow(pred_vel[idx, 0].numpy())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-02T10:43:27.030889Z","iopub.execute_input":"2025-07-02T10:43:27.031134Z","iopub.status.idle":"2025-07-02T10:43:27.263559Z","shell.execute_reply.started":"2025-07-02T10:43:27.031113Z","shell.execute_reply":"2025-07-02T10:43:27.262931Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}