{"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":11367935,"sourceType":"datasetVersion","datasetId":7116013},{"sourceId":11368083,"sourceType":"datasetVersion","datasetId":7116134},{"sourceId":11368169,"sourceType":"datasetVersion","datasetId":7116196},{"sourceId":11368268,"sourceType":"datasetVersion","datasetId":7116272},{"sourceId":11368499,"sourceType":"datasetVersion","datasetId":7116445},{"sourceId":11368545,"sourceType":"datasetVersion","datasetId":7116479},{"sourceId":11368547,"sourceType":"datasetVersion","datasetId":7116481},{"sourceId":11376433,"sourceType":"datasetVersion","datasetId":7122462},{"sourceId":11376448,"sourceType":"datasetVersion","datasetId":7122476},{"sourceId":11376464,"sourceType":"datasetVersion","datasetId":7122489},{"sourceId":11376742,"sourceType":"datasetVersion","datasetId":7122712},{"sourceId":11376868,"sourceType":"datasetVersion","datasetId":7122812},{"sourceId":11376871,"sourceType":"datasetVersion","datasetId":7122814},{"sourceId":11376872,"sourceType":"datasetVersion","datasetId":7122815},{"sourceId":11376935,"sourceType":"datasetVersion","datasetId":7122866},{"sourceId":11377083,"sourceType":"datasetVersion","datasetId":7122981},{"sourceId":11377231,"sourceType":"datasetVersion","datasetId":7123090},{"sourceId":11377291,"sourceType":"datasetVersion","datasetId":7123138},{"sourceId":11377325,"sourceType":"datasetVersion","datasetId":7123163},{"sourceId":11377334,"sourceType":"datasetVersion","datasetId":7123172},{"sourceId":11377594,"sourceType":"datasetVersion","datasetId":7123380},{"sourceId":11377614,"sourceType":"datasetVersion","datasetId":7123394},{"sourceId":11377741,"sourceType":"datasetVersion","datasetId":7123490},{"sourceId":11377752,"sourceType":"datasetVersion","datasetId":7123499},{"sourceId":11377756,"sourceType":"datasetVersion","datasetId":7123503},{"sourceId":11377935,"sourceType":"datasetVersion","datasetId":7123649},{"sourceId":11377970,"sourceType":"datasetVersion","datasetId":7123675},{"sourceId":11378141,"sourceType":"datasetVersion","datasetId":7123813},{"sourceId":11378162,"sourceType":"datasetVersion","datasetId":7123831},{"sourceId":11378178,"sourceType":"datasetVersion","datasetId":7123842}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# [GWI] Improved UNet pipepline with larger dataset\n- The aim is to improve the best score, that is currently publicly available, using Kaggle's computational resources.\n- We are going to use UNet model as introduced in [5 depth U net with residual](https://www.kaggle.com/code/adhok93/5-depth-u-net-with-residual) Notebook.\n- We are going to use part of Full OpenWFI dataset, which was introduced in [OpenFWI InversionNet Train with 670G Datasets](https://www.kaggle.com/code/seshurajup/openfwi-inversionnet-train-with-670g-datasets) Notebook.","metadata":{"execution":{"iopub.status.busy":"2025-04-25T14:53:18.823205Z","iopub.execute_input":"2025-04-25T14:53:18.823481Z","iopub.status.idle":"2025-04-25T14:53:23.506381Z","shell.execute_reply.started":"2025-04-25T14:53:18.823463Z","shell.execute_reply":"2025-04-25T14:53:23.505626Z"}}},{"cell_type":"code","source":"%%writefile config.yaml\n\ndata_path: /kaggle/input/waveform-inversion\nmodel: \n    name: UNet\n    unet_params:\n        init_features: 32\n        depth: 5\nread_weights: null\nbatch_size: 64\nprint_freq: 500\nmax_epochs: 20\nes_epochs: 4\nseed: 42\nvalid_frac: 36\ntrain_frac: 5\noptimizer:\n    lr: 0.0001\n    weight_decay: 0.001\nscheduler:\n    params:\n        factor: 0.316227766\n        patience: 1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T16:01:28.757532Z","iopub.execute_input":"2025-04-25T16:01:28.757786Z","iopub.status.idle":"2025-04-25T16:01:28.766263Z","shell.execute_reply.started":"2025-04-25T16:01:28.757764Z","shell.execute_reply":"2025-04-25T16:01:28.765449Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Data\n\nimport os\nimport numpy as np\nfrom pathlib import Path\nfrom torch.utils.data import Dataset, DataLoader\n\ndef inputs_files_to_output_files(input_files):\n    return [\n        Path(str(f).replace('seis', 'vel').replace('data', 'model'))\n        for f in input_files\n    ]\n\n\ndef get_train_files(data_path):\n\n    all_inputs = [\n        f\n        for f in\n        Path(data_path).rglob('*.npy')\n        if ('seis' in f.stem) or ('data' in f.stem)\n    ]\n\n    all_outputs = inputs_files_to_output_files(all_inputs)\n\n    assert all(f.exists() for f in all_outputs)\n\n    return all_inputs, all_outputs\n\n\nclass SeismicDataset(Dataset):\n    def __init__(self, inputs_files, output_files, n_examples_per_file=500):\n        assert len(inputs_files) == len(output_files)\n        self.inputs_files = inputs_files\n        self.output_files = output_files\n        self.n_examples_per_file = n_examples_per_file\n\n    def __len__(self):\n        return len(self.inputs_files) * self.n_examples_per_file\n\n    def __getitem__(self, idx):\n        # Calculate file offset and sample offset within file\n        file_idx = idx // self.n_examples_per_file\n        sample_idx = idx % self.n_examples_per_file\n\n        X = np.load(self.inputs_files[file_idx], mmap_mode='r')\n        y = np.load(self.output_files[file_idx], mmap_mode='r')\n\n        try:\n            return X[sample_idx].copy(), y[sample_idx].copy()\n        finally:\n            del X, y\n\n\nclass TestDataset(Dataset):\n    def __init__(self, test_files):\n        self.test_files = test_files\n\n\n    def __len__(self):\n        return len(self.test_files)\n\n\n    def __getitem__(self, i):\n        test_file = self.test_files[i]\n\n        return np.load(test_file), test_file.stem","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T16:01:28.767697Z","iopub.execute_input":"2025-04-25T16:01:28.768091Z","iopub.status.idle":"2025-04-25T16:01:33.324282Z","shell.execute_reply.started":"2025-04-25T16:01:28.768062Z","shell.execute_reply":"2025-04-25T16:01:33.323480Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Model\n\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch\n\nclass ResidualDoubleConv(nn.Module):\n    \"\"\"(Convolution => [BN] => ReLU) * 2 + Residual Connection\"\"\"\n\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n\n        # First convolution layer\n        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(mid_channels)\n        self.relu = nn.ReLU(inplace=True)\n\n        # Second convolution layer\n        self.conv2 = nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(out_channels)\n\n        # Shortcut connection to handle potential channel mismatch\n        if in_channels == out_channels:\n            self.shortcut = nn.Identity()\n        else:\n            # Projection shortcut: 1x1 conv + BN to match output channels\n            self.shortcut = nn.Sequential(\n                nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),\n                nn.BatchNorm2d(out_channels)\n            )\n\n    def forward(self, x):\n        identity = x  # Store the input for the residual connection\n\n        # First conv block\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        # Second conv block (without final ReLU yet)\n        out = self.conv2(out)\n        out = self.bn2(out)\n\n        # Apply shortcut to the identity path\n        identity_mapped = self.shortcut(identity)\n\n        # Add the residual connection\n        out += identity_mapped\n\n        # Apply final ReLU\n        out = self.relu(out)\n        return out\n\n\nclass Up(nn.Module):\n    \"\"\"Upscaling then ResidualDoubleConv\"\"\"\n\n    def __init__(self, in_channels, out_channels, bilinear=True):\n        super().__init__()\n        self.bilinear = bilinear\n\n        if bilinear:\n            self.up = nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False)\n            # Input to ResidualDoubleConv = channels from upsampled layer below + channels from skip connection\n            # Output of ResidualDoubleConv = desired output channels for this decoder stage\n            self.conv = ResidualDoubleConv(in_channels + out_channels, out_channels) # Use ResidualDoubleConv\n\n        else: # Using ConvTranspose2d\n            # ConvTranspose halves the channels: in_channels -> in_channels // 2\n            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)\n            # Input channels to ResidualDoubleConv\n            conv_in_channels = in_channels // 2 # Channels after ConvTranspose\n            skip_channels = out_channels       # Channels from skip connection\n            total_in_channels = conv_in_channels + skip_channels\n            self.conv = ResidualDoubleConv(total_in_channels, out_channels) # Use ResidualDoubleConv\n\n    def forward(self, x1, x2):\n        # x1 is the feature map from the layer below (needs upsampling)\n        # x2 is the skip connection from the corresponding encoder layer\n        x1 = self.up(x1)\n\n        # Pad x1 if its dimensions don't match x2 after upsampling\n        # Input is CHW\n        diffY = x2.size(2) - x1.size(2)\n        diffX = x2.size(3) - x1.size(3)\n\n        # Pad format: (padding_left, padding_right, padding_top, padding_bottom)\n        x1 = F.pad(\n            x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]\n        )\n\n        # Concatenate along the channel dimension\n        x = torch.cat([x2, x1], dim=1)\n        return self.conv(x)\n\nclass OutConv(nn.Module):\n    \"\"\"1x1 Convolution for the output layer\"\"\"\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass UNet(nn.Module):\n    \"\"\"U-Net architecture implementation with Residual Blocks\"\"\"\n\n    def __init__(\n        self,\n        n_channels=5,\n        n_classes=1,\n        init_features=32,\n        depth=5, # number of pooling layers\n        bilinear=True,\n    ):\n        super().__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n        self.bilinear = bilinear\n        self.depth = depth\n\n        self.initial_pool = nn.AvgPool2d(kernel_size=(14, 1), stride=(14, 1))\n\n        # --- Encoder ---\n        self.encoder_convs = nn.ModuleList() # Store conv blocks\n        self.encoder_pools = nn.ModuleList() # Store pool layers\n\n        # Initial conv block (no pooling before it)\n        # Use ResidualDoubleConv for the initial convolution block\n        self.inc = ResidualDoubleConv(n_channels, init_features)\n        self.encoder_convs.append(self.inc)\n\n        current_features = init_features\n        for _ in range(depth):\n            # Define convolution block for this stage\n            conv = ResidualDoubleConv(current_features, current_features * 2)\n            # Define pooling layer for this stage\n            pool = nn.MaxPool2d(2)\n            self.encoder_convs.append(conv)\n            self.encoder_pools.append(pool)\n            current_features *= 2\n\n        # --- Bottleneck ---\n        # Use ResidualDoubleConv for the bottleneck\n        self.bottleneck = ResidualDoubleConv(current_features, current_features)\n\n        # --- Decoder ---\n        self.decoder_blocks = nn.ModuleList()\n        # Input features start from bottleneck output features\n        # Output features at each stage are halved\n        for _ in range(depth):\n            # Up block uses ResidualDoubleConv internally and handles channels\n            up_block = Up(current_features, current_features // 2, bilinear)\n            self.decoder_blocks.append(up_block)\n            current_features //= 2 # Halve features for next Up block input\n\n        # --- Output Layer ---\n        # Input features are the output features of the last Up block\n        self.outc = OutConv(current_features, n_classes)\n\n    def _pad_or_crop(self, x, target_h=70, target_w=70):\n        \"\"\"Pads or crops input tensor x to target height and width.\"\"\"\n        _, _, h, w = x.shape\n        # Pad Height if needed\n        if h < target_h:\n            pad_top = (target_h - h) // 2\n            pad_bottom = target_h - h - pad_top\n            x = F.pad(x, (0, 0, pad_top, pad_bottom))  # Pad height only\n            h = target_h\n        # Pad Width if needed\n        if w < target_w:\n            pad_left = (target_w - w) // 2\n            pad_right = target_w - w - pad_left\n            x = F.pad(x, (pad_left, pad_right, 0, 0))  # Pad width only\n            w = target_w\n        # Crop Height if needed\n        if h > target_h:\n            crop_top = (h - target_h) // 2\n            # Use slicing to crop\n            x = x[:, :, crop_top : crop_top + target_h, :]\n            h = target_h\n        # Crop Width if needed\n        if w > target_w:\n            crop_left = (w - target_w) // 2\n            x = x[:, :, :, crop_left : crop_left + target_w]\n            w = target_w\n        return x\n\n    def forward(self, x):\n        # Initial pooling and resizing\n        x_pooled = self.initial_pool(x)\n        x_resized = self._pad_or_crop(x_pooled, target_h=70, target_w=70)\n\n        # --- Encoder Path ---\n        skip_connections = []\n        xi = x_resized\n\n        # Apply initial conv (inc)\n        xi = self.encoder_convs[0](xi)\n        skip_connections.append(xi) # Store output of inc\n\n        # Apply subsequent encoder convs and pools\n        # self.depth is the number of pooling layers\n        for i in range(self.depth):\n            # Apply conv block for this stage\n            xi = self.encoder_convs[i+1](xi)\n            # Store skip connection *before* pooling\n            skip_connections.append(xi)\n            # Apply pooling layer for this stage\n            xi = self.encoder_pools[i](xi)\n\n        # Apply bottleneck conv\n        xi = self.bottleneck(xi)\n\n        # --- Decoder Path ---\n        xu = xi # Start with bottleneck output\n        # Iterate through decoder blocks and corresponding skip connections in reverse\n        for i, block in enumerate(self.decoder_blocks):\n            # Determine the correct skip connection index from the end\n            # Example: depth=5. Skips stored: [inc, enc1, enc2, enc3, enc4] (indices 0-4)\n            # Decoder 0 (Up(1024, 512)) needs skip 4 (enc4)\n            # Decoder 1 (Up(512, 256)) needs skip 3 (enc3) ...\n            # Decoder 4 (Up(64, 32)) needs skip 0 (inc)\n            skip_index = self.depth - 1 - i\n            skip = skip_connections[skip_index]\n            xu = block(xu, skip) # Up block combines xu (from below) and skip\n\n        # --- Final Output ---\n        logits = self.outc(xu)\n        # Apply scaling and offset specific to the problem's target range\n        output = logits * 1000.0 + 1500.0\n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T16:01:33.325487Z","iopub.execute_input":"2025-04-25T16:01:33.325773Z","iopub.status.idle":"2025-04-25T16:01:33.345274Z","shell.execute_reply.started":"2025-04-25T16:01:33.325755Z","shell.execute_reply":"2025-04-25T16:01:33.344378Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Utils\n\nimport datetime\nimport random\nimport torch\nimport numpy as np\n\ndef format_time(elapsed):\n    \"\"\"Take a time in seconds and return a string hh:mm:ss.\"\"\"\n    elapsed_rounded = int(round((elapsed)))\n    return str(datetime.timedelta(seconds=elapsed_rounded))\n\ndef seed_everything(\n    seed_value: int\n) -> None:\n    \"\"\"\n    Controlling a unified seed value for Python, NumPy, and PyTorch (CPU, GPU).\n\n    Parameters:\n    ----------\n    seed_value : int\n        The unified random seed value.\n    \"\"\"\n    random.seed(seed_value) # Python\n    np.random.seed(seed_value) # cpu vars\n    torch.manual_seed(seed_value) # cpu  vars    \n    if torch.cuda.is_available(): \n        torch.cuda.manual_seed(seed_value)\n        torch.cuda.manual_seed_all(seed_value) # gpu vars\n    if torch.backends.cudnn.is_available:\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T16:01:33.346202Z","iopub.execute_input":"2025-04-25T16:01:33.346529Z","iopub.status.idle":"2025-04-25T16:01:33.372034Z","shell.execute_reply.started":"2025-04-25T16:01:33.346501Z","shell.execute_reply":"2025-04-25T16:01:33.371187Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Preparation and training\n\nimport os\nimport sys\nimport yaml\nfrom torch.utils.data import DataLoader\nfrom pprint import pprint\nimport torch\nimport torch.nn as nn\nimport numpy as np\nimport time\n\nprint(f\"GPU: {torch.cuda.get_device_name(0)}\")\n_, total = torch.cuda.mem_get_info(device=0)\nprint(f\"GPU memory: {total / 1024**3:.2f}GB\")\n\nwith open(\"config.yaml\", \"r\") as file_obj:\n    config = yaml.safe_load(file_obj)\nprint()\npprint(config)\nif config[\"data_path\"] is None:\n    config[\"data_path\"] = os.environ[\"TMPDIR\"]\n    print(\"data_path:\", config[\"data_path\"])\nprint()\n\nseed_everything(config[\"seed\"])\n\nall_inputs, all_outputs = [], []\nfor i in range(1, 31):\n    all_inputs1, all_outputs1 = get_train_files(config[\"data_path\"]+f\"-{i}\")\n    all_inputs.extend(all_inputs1)\n    all_outputs.extend(all_outputs1)\nprint(\"Total number of input/output files:\", len(all_inputs))\n\nvalid_inputs = [all_inputs[i] for i in range(0, len(all_inputs), config[\"valid_frac\"])]\ntrain_inputs = [f for f in all_inputs if not f in valid_inputs]\nif config[\"train_frac\"] > 1:\n    train_inputs = [train_inputs[i] for i in range(0, len(train_inputs), config[\"train_frac\"])]\n\nprint(\"Number of train files:\", len(train_inputs))\nprint(\"Number of valid files:\", len(valid_inputs))\nprint()\n\ntrain_outputs = inputs_files_to_output_files(train_inputs)\nvalid_outputs = inputs_files_to_output_files(valid_inputs)\n\ndstrain = SeismicDataset(train_inputs, train_outputs)\ndltrain = DataLoader(\n    dstrain,\n    batch_size=config[\"batch_size\"],\n    shuffle=True,\n    pin_memory=False,\n    drop_last=True,\n    num_workers=4,\n    persistent_workers=True,\n)\n\ndsvalid = SeismicDataset(valid_inputs, valid_outputs)\ndlvalid = DataLoader(\n    dsvalid,\n    batch_size=4*config[\"batch_size\"],\n    shuffle=False,\n    pin_memory=False,\n    drop_last=False,\n    num_workers=4,\n    persistent_workers=True,\n)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = UNet(**config[\"model\"][\"unet_params\"]).to(device)\n\nif config[\"read_weights\"] is not None:\n    print(\"Reading weights from:\", config[\"read_weights\"])\n    model.load_state_dict(torch.load(config[\"read_weights\"], weights_only=True))\n\ncriterion = nn.L1Loss()\noptimizer = torch.optim.AdamW(model.parameters(), **config[\"optimizer\"])  # hparams\nscheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', **config[\"scheduler\"][\"params\"])\n\nbest_val_loss = 10000.0\nepochs_wo_improvement = 0\nt0 = time.time()  # Measure staring time\n\nfor epoch in range(1, config[\"max_epochs\"] + 1):\n\n    # Train\n    model.train()\n    train_losses = []\n    for step, (inputs, targets) in enumerate(dltrain):\n\n        inputs = inputs.to(device)\n        targets = targets.to(device)\n        optimizer.zero_grad()\n\n        with torch.autocast(device_type=\"cuda\"):\n            outputs = model(inputs)\n            loss = criterion(outputs, targets)\n\n        loss.backward()\n        optimizer.step()\n        train_losses.append(loss.item())\n\n        if step % config[\"print_freq\"] == config[\"print_freq\"] - 1 or step == len(dltrain) - 1:\n            trn_loss = np.mean(train_losses)\n            t1 = format_time(time.time() - t0)\n            free, total = torch.cuda.mem_get_info(device=0)\n            mem_used = (total - free) / 1024**3\n            lr = optimizer.param_groups[-1]['lr']\n            print(\n                f\"Epoch: {epoch:02d}  Step {step+1}/{len(dltrain)}  Trn Loss: {trn_loss:.2f}  LR: {lr:.2e}  GPU Usage: {mem_used:.2f}GB  Elapsed Time: {t1}\",\n                flush=True,\n            )\n\n    # Valid\n    model.eval()\n    valid_losses = []\n    for inputs, targets in dlvalid:\n        inputs = inputs.to(device)\n        targets = targets.to(device)\n\n        with torch.inference_mode():\n            with torch.autocast(device_type=\"cuda\"):\n                outputs = model(inputs)\n\n        loss = criterion(outputs, targets)\n\n        valid_losses.append(loss.item())\n\n    t1 = format_time(time.time() - t0)\n    trn_loss = np.mean(train_losses)\n    val_loss = np.mean(valid_losses)\n\n    free, total = torch.cuda.mem_get_info(device=0)\n    mem_used = (total - free) / 1024**3\n\n    print(\n        f\"\\nEpoch: {epoch:02d}  Trn Loss: {trn_loss:.2f}  Val Loss: {val_loss:.2f}  GPU Usage: {mem_used:.2f}GB  Elapsed Time: {t1}\",\n        flush=True,\n    )\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        epochs_wo_improvement = 0\n        torch.save(model.state_dict(), \"best_model.pth\")\n        print(f\"\\nNew best val_loss: {val_loss:.2f}\\n\", flush=True)\n    else:\n        epochs_wo_improvement += 1\n        print(f\"\\nEpochs without improvement: {epochs_wo_improvement}\\n\", flush=True)\n\n    if epochs_wo_improvement == config[\"es_epochs\"]:\n        break\n\n    scheduler.step(val_loss)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T16:01:33.373648Z","iopub.execute_input":"2025-04-25T16:01:33.373947Z","execution_failed":"2025-04-25T16:01:47.961Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Inference\n\nimport csv\nfrom pathlib import Path\nimport os\n\nt0 = time.time()\n\ntest_files = list(Path(os.path.join(config[\"data_path\"], \"test\")).glob(\"*.npy\"))\nx_cols = [f\"x_{i}\" for i in range(1, 70, 2)]\nfieldnames = [\"oid_ypos\"] + x_cols\nds = TestDataset(test_files)\ndl = DataLoader(ds, batch_size=4*config[\"batch_size\"], num_workers=4, pin_memory=False)\n\nmodel.load_state_dict(torch.load(\"best_model.pth\", weights_only=True))\n\nmodel.eval()\nwith open(\"submission.csv\", \"wt\", newline=\"\") as csvfile:\n    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n    writer.writeheader()\n\n    for inputs, oids_test in dl:\n        inputs = inputs.to(device)\n        with torch.inference_mode():\n            with torch.autocast(device_type=\"cuda\"):\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\nt1 = format_time(time.time() - t0)\nprint(f\"Inference Time: {t1}\")","metadata":{"trusted":true,"execution":{"execution_failed":"2025-04-25T16:01:47.963Z"}},"outputs":[],"execution_count":null}]}