{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.16","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","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":31013,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Caformar TPU training","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":"import torch\n#if not torch.cuda.is_available():\n#    raise RuntimeError(\"Requires GPUs with CUDA enabled.\")\ntry: \n    import monai\nexcept: \n    !pip install --no-deps monai -q\n\n!pip uninstall -y tensorflow\n!pip install tensorflow-cpu","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-09T10:34:41.641813Z","iopub.execute_input":"2025-06-09T10:34:41.641995Z","iopub.status.idle":"2025-06-09T10:36:21.680816Z","shell.execute_reply.started":"2025-06-09T10:34:41.641974Z","shell.execute_reply":"2025-06-09T10:36:21.675806Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile _cfg.py\nfrom types import SimpleNamespace\nimport torch\n\ncfg= SimpleNamespace()\ncfg.seed = 3\ncfg.subsample = None\ncfg.ema = True\ncfg.ema_decay = 0.99\ncfg.backbone = \"caformer_b36.sail_in22k_ft_in1k\"\ncfg.epochs = 5\ncfg.batch_size = 32\ncfg.RUN_VALID = True\ncfg.RUN_TEST  = True\ncfg.RUN_TRAIN = True\ncfg.RUN_TRAIN_ALL = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-09T10:36:21.683216Z","iopub.execute_input":"2025-06-09T10:36:21.683728Z","iopub.status.idle":"2025-06-09T10:36:21.694727Z","shell.execute_reply.started":"2025-06-09T10:36:21.683692Z","shell.execute_reply":"2025-06-09T10:36:21.690375Z"}},"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":"%%writefile _scheduler.py\nimport 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]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-09T10:36:21.696585Z","iopub.execute_input":"2025-06-09T10:36:21.696789Z","iopub.status.idle":"2025-06-09T10:36:22.042906Z","shell.execute_reply.started":"2025-06-09T10:36:21.696769Z","shell.execute_reply":"2025-06-09T10:36:22.038097Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport matplotlib.pyplot as plt\nfrom _scheduler import ConstantCosineLR\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":{"iopub.status.busy":"2025-06-09T10:36:22.044733Z","iopub.execute_input":"2025-06-09T10:36:22.044975Z","iopub.status.idle":"2025-06-09T10:36:30.391951Z","shell.execute_reply.started":"2025-06-09T10:36:22.044952Z","shell.execute_reply":"2025-06-09T10:36:30.386452Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset","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)):\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\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-09T10:36:30.394117Z","iopub.execute_input":"2025-06-09T10:36:30.394458Z","iopub.status.idle":"2025-06-09T10:36:30.405096Z","shell.execute_reply.started":"2025-06-09T10:36:30.394431Z","shell.execute_reply":"2025-06-09T10:36:30.400556Z"}},"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 = []\n\n        for m in self.models:\n            logits = m(x)\n\n            output.append(logits)\n\n        output = torch.stack(output)\n        output = torch.quantile(output, 0.5, dim=0)\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        ) # Output [B, C, 1, 1]\n        self.sSE = nn.Sequential(\n            nn.Conv2d(in_channels, 1, 1), \n            nn.Sigmoid(),\n            ) # Output [B, 1, H, W]\n\n    def forward(self, x):\n        return x * self.cSE(x) + x * self.sSE(x)\n\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\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,"execution":{"iopub.status.busy":"2025-06-09T10:36:30.407544Z","iopub.execute_input":"2025-06-09T10:36:30.407792Z","iopub.status.idle":"2025-06-09T10:36:30.426487Z","shell.execute_reply.started":"2025-06-09T10:36:30.407768Z","shell.execute_reply":"2025-06-09T10:36:30.422957Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Utils\n\nSame as previous notebook. ","metadata":{}},{"cell_type":"code","source":"%%writefile _utils.py\n\nimport datetime\n\ndef format_time(elapsed):\n    elapsed_rounded = int(round((elapsed)))\n    return str(datetime.timedelta(seconds=elapsed_rounded))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-09T10:36:30.429076Z","iopub.execute_input":"2025-06-09T10:36:30.429350Z","iopub.status.idle":"2025-06-09T10:36:30.440091Z","shell.execute_reply.started":"2025-06-09T10:36:30.429303Z","shell.execute_reply":"2025-06-09T10:36:30.434759Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train and Valid","metadata":{}},{"cell_type":"code","source":"%%writefile train_fwi_xla_8tpu.py\nimport sys\nimport os\nfrom tqdm import tqdm\nimport numpy as np\nimport time \nimport gc\nimport random\n\nimport torch\nimport torch.nn as nn\nfrom torch.amp import autocast\nfrom torch.utils.data import DistributedSampler\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.utils as xu\nfrom torch_xla.amp import syncfree\n\n\nfrom _cfg import cfg\nfrom _dataset import CustomDataset\nfrom _model import ModelEMA, Net\nfrom _utils import format_time\nfrom _scheduler import ConstantCosineLR\n\n\ndef _mp_fn(index, cfg):\n    # Setup\n    torch.manual_seed(cfg.seed)\n    world_size = xm.xrt_world_size()\n    local_rank = xm.get_ordinal()\n    device = xm.xla_device()\n    xm.master_print(f\"Process {local_rank} initialized on device: {device}\")\n\n    # Prepare Dataset\n    train_ds = CustomDataset(cfg=cfg, mode=\"train\")\n    valid_ds = CustomDataset(cfg=cfg, mode=\"valid\")\n    xm.rendezvous('Dataset preparing complete')\n    # Train Data\n    sampler_tr = DistributedSampler(\n        train_ds,\n        num_replicas=world_size,\n        rank=local_rank,\n        shuffle=True\n        )\n    train_pytorch_dl = torch.utils.data.DataLoader(\n        train_ds,\n        batch_size=cfg.batch_size,\n        sampler=sampler_tr,\n        num_workers=1,\n    )\n    train_dl = pl.MpDeviceLoader(train_pytorch_dl, device)\n\n    # Valid Data\n    sampler_vl = DistributedSampler(\n        valid_ds,\n        num_replicas=world_size,\n        rank=local_rank,\n        shuffle=True\n        )\n    valid_pytorch_dl = torch.utils.data.DataLoader(\n        valid_ds,\n        batch_size=cfg.batch_size,\n        sampler=sampler_vl,\n        num_workers=1,\n    )\n    valid_dl = pl.MpDeviceLoader(valid_pytorch_dl, device)\n\n    # Define Models\n    model = Net(\n        backbone=cfg.backbone,\n        pretrained=False,\n        )\n    state_dict= torch.load(\"/kaggle/input/openfwi-preprocessed-72x72/models_1000x70/unet2d_caformer_seed3_epochbest.pt\", map_location=torch.device('cpu'), weights_only=True)\n    state_dict= {k.removeprefix(\"_orig_mod.\"):v for k,v in state_dict.items()} # Remove torch.compile() prefix\n    \n    model.load_state_dict(state_dict)\n    model = model.to(device)\n    \n    if cfg.ema:\n        print(\"Initializing EMA model..\")\n        ema_model = ModelEMA(\n            model,\n            decay=cfg.ema_decay,\n            device=device            )\n    else:\n        ema_model = None\n\n\n    # Larning Parameters\n    criterion = nn.L1Loss()\n    optimizer = syncfree.SGD(model.parameters(), lr=1e-6)\n    # Custom Scheduler TODO: Fill the value of total_steps\n    scheduler = ConstantCosineLR(optimizer, total_steps=len(train_dl)*cfg.epochs, pct_cosine=0.7)\n\n\n    # ================  Train and Valid ===================\n    val_best_loss = 1e5\n    for epoch in range(cfg.epochs):\n        if epoch == 0:\n            tstart = time.time()\n            print(\"Start Training\")        \n        # =============== Train Loop ===============\n        tracker = xm.RateTracker()\n        model.train()\n        tr_total_loss = []\n        sampler_tr.set_epoch(epoch)\n        for i, (inputs, targets) in enumerate(train_dl):\n            optimizer.zero_grad(set_to_none=True)\n            with autocast('xla'):   \n                logits = model(inputs)\n                loss = criterion(logits, targets)\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 2.0)\n            tr_total_loss.append(loss.item())\n            if ema_model is not None:\n                ema_model.update(model)\n\n            xm.optimizer_step(optimizer, barrier=True)\n            scheduler.step()\n\n            tracker.add(cfg.batch_size)\n            if (i+1) % 100 == 0 or i == 0:\n                train_loss = np.mean(tr_total_loss)\n                tr_total_loss = []\n                if xm.is_master_ordinal():\n                    xm.master_print(f\"Epoch: {epoch}, Step: {i+1}/{len(train_dl)}, Loss: {train_loss}, Rate: {tracker.rate()} samples/sec\")\n        # For insurance\n        if xm.is_master_ordinal():\n            xm.save(model.state_dict(), f'epoch_{epoch}_train.pt')\n\n        # =============== Valid Loop ===============\n        model.eval()\n        val_logits = []\n        val_targets = []\n        with torch.inference_mode():\n            for i, (inputs, targets) in enumerate(valid_dl):\n                with autocast('xla'):\n                    if ema_model is not None:\n                        out = ema_model.module(inputs)\n                    else:\n                        out = model(inputs)\n                    \n                val_logits.append(out)\n                val_targets.append(targets)\n                \n            val_logits = torch.cat(val_logits, dim=0)\n            val_targets = torch.cat(val_targets, dim=0)\n\n            gathered_logits = xm.all_gather(val_logits)\n            gathered_targets = xm.all_gather(val_targets)\n\n            if xm.is_master_ordinal():\n                loss = criterion(gathered_logits, gathered_targets).item()\n                xm.master_print(f'Epoch {epoch} Val Loss={loss} Time={time.asctime()}')\n                # Save model\n                xm.save(model.state_dict(), f'epoch_{epoch}.pt')\n                xm.master_print(f'Model checkpoint saved!')\n                \n        xm.master_print(f\"Finish epoch {epoch}\")\n\nif __name__ == '__main__':\n    import os\n    os.environ.pop(\"TPU_PROCESS_ADDRESSES\")\n    xmp.spawn(fn=_mp_fn, args=(cfg,), start_method='fork')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-09T10:36:30.441586Z","iopub.execute_input":"2025-06-09T10:36:30.441786Z","iopub.status.idle":"2025-06-09T10:36:30.452848Z","shell.execute_reply.started":"2025-06-09T10:36:30.441766Z","shell.execute_reply":"2025-06-09T10:36:30.448621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python train_fwi_xla_8tpu.py","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-09T10:36:30.454942Z","iopub.execute_input":"2025-06-09T10:36:30.455137Z","iopub.status.idle":"2025-06-09T11:04:40.112904Z","shell.execute_reply.started":"2025-06-09T10:36:30.455117Z","shell.execute_reply":"2025-06-09T11:04:40.107093Z"}},"outputs":[],"execution_count":null}]}