![\"/Users/osman/Desktop/Screenshot](\"https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F20259944%2F27dc43a53e40e5e9ff69b8ab1eff4495%2FScreenshot%202026-03-02%20at%2012.37.40.png?generation=1772444684940610&alt=media\")
I'll be honest. I didn't do as well as I had hoped. I was aiming much higher like top 10 But the result didn't reflect that effort, and I think I know whys.
\nI'll be real about where I was coming from: I'm relatively new to competitive machine learning, TPUs, and JAX all at once. There was a steep learning curve just to get anything running, let alone running well.
\nI won't pretend the motivation was purely academic either. Top 10 meant prize money, and I genuinely needed it. That drove me to put in a lot of hours β probably more than I should have at certain points. I tried hard. I just didn't have the experience yet to convert that effort into a top result.
\nAnd that's okay. I learned an enormous amount. I went from knowing nothing about distributed TPU training to having a working end-to-end 3D segmentation pipeline. That's real progress, even if the leaderboard didn't fully show it. Next time, I'll come in with that foundation already in place.
\nA big thank you to Innat β his notebook was what got me started with the medicai framework and TPU-compatible 3D training.
This training was done as part of the TRC (TPU Research Cloud) program. I'm genuinely grateful to Google for providing access to TPU v4-32 hardware through TRC training 160Β³ 3D models at this scale would have been completely infeasible otherwise. Thank you.
\nI built a 3D segmentation pipeline to detect ancient papyrus surfaces from CT scans of the Herculaneum scrolls. My final solution uses TransUNet [1] with a SEResNeXt [2] encoder, trained on 160Β³ voxel patches with a Tversky [3] + CenterlineDice [4] loss on TPU v4-32. Inference combines sliding window, 6-fold TTA, dual checkpoint ensemble, and morphological post-processing. My biggest personal takeaway was learning to train 3D medical segmentation models end-to-end on TPUs with JAX/Keras.
\n
End-to-end pipeline: 3D data prep β TransUNet training β dual checkpoint blending β post-processing β binary ink mask.
\nInput (160Β³Γ1) β SEResNeXt CNN Encoder β Transformer Bottleneck β UNet Decoder β Softmax (3 classes)\nWhy this architecture?
\nI experimented with two encoder sizes:
\n| Variant | \nTraining | \nInference | \n
|---|---|---|
seresnext50 | \nβ (train-tpu.py) | \nβ | \n
seresnext101 | \nβ (final model) | \nβ (inference-kaggle6.py) | \n
| Parameter | \nValue | \n
|---|---|
| Hardware | \nKaggle TPU v4-32 (32 cores) | \n
| Framework | \nKeras 3 + JAX backend | \n
| Batch size | \ntotal_devices // 2 (16 on TPU v4-32) | \n
| Epochs | \n800 (CosineDecay + EarlyStopping patience=60) | \n
| Base LR | \n2e-4 with sqrt batch scaling | \n
| LR Schedule | \nSingle-cycle CosineDecay, 5-epoch warmup, alpha=0.005 | \n
| Optimizer | \nAdamW (weight_decay=1e-5, clipnorm=1.0) | \n
| Loss | \nSparseTverskyLoss + 0.2 Γ SparseCenterlineDiceLoss | \n
| Val split | \n5 TFRecord shards (~30 samples), seed=42 | \n
# Tversky loss β recall-focused (alpha=0.3 FP, beta=0.7 FN)\ntversky = SparseTverskyLoss(alpha=0.3, beta=0.7, ignore_class_ids=2)\n\n# CenterlineDice β topology preservation via soft-skeletonisation\ncldice = SparseCenterlineDiceLoss(target_class_ids=1, ignore_class_ids=2, iters=20)\n\n# Combined\nloss = tversky + 0.2 * cldice\nRationale: Tversky with Ξ²=0.7 penalizes missed surfaces (FN) more than false positives [3], aligning with the competition metric's sensitivity to gaps. CenterlineDice [4] preserves topological continuity via soft-skeletonisation on predictions β directly targeting the TopoScore sub-metric. The ignore_class_ids=2 automatically masks out unlabeled regions.
Data is stored as TFRecords (780 samples). Training augmentation includes:
\nmin_valid_ratio=0.3 β rejects patches >70% ignore-class)nonzero=True) β adapts to each volume's actual intensity range rather than fixed [0,255]β[0,1] scalingEarlier experiments used CosineDecayRestarts(m_mul=0.9), but the learning rate collapsed too early (~epoch 300). Switching to a single-cycle CosineDecay over the full 800-epoch budget produced much more stable convergence.
SlidingWindowInference(\n roi_size=(160, 160, 160), # must match ViT positional embeddings (5Β³=125 patches)\n overlap=0.5,\n blend_mode=\"gaussian\", # smooth boundary artifacts\n)\nTTA_AXES = [(), (0,), (1,), (2,), (0,1), (1,2)]\n# Original + 3 single-axis flips + 2 pair flips β averaged\npred = 0.80 * pred_best_val_dice + 0.20 * pred_best_train_loss\nTwo checkpoints from the same training run β best_val_dice generalizes better, best_train_loss is more expressive. Blending consistently improved over either single checkpoint.
ink_prob -= 0.15 Γ ignore_prob β leverages class-2 predictions to reduce false positives in uncertain regionsThis competition was my first serious experience training on TPUs. It required several specific codebase adjustments to run successfully:
\n1. JAX backend over TensorFlow:\nJAX is significantly more stable and faster on TPUs than TF eager mode.
\nimport os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n2. Disabling Preallocation (Critical for Inference):\nBy default, JAX reserves ~75% of VRAM upfront. For a 160Β³ 3D volume on inference, this guarantees an OOM error before the model even spins up.
\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\nos.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n3. Distributed Training Strategy:\nScaling across 32 cores requires the DataParallel distribution. Tuning the batch size is critical; too large wastes HBM, too small underutilizes the v4-32 hardware.
import jax\nimport keras\n\njax.distributed.initialize()\ndevices = keras.distribution.list_devices()\ndata_parallel = keras.distribution.DataParallel(devices=devices)\nkeras.distribution.set_distribution(data_parallel)\ndata_parallel.auto_shard_dataset = False # Essential for custom tf.data pipelines\n4. Flash Attention Incompatibility:\nDistributed training over multiple TPU workers currently requires disabling Flash Attention to prevent XlaRuntimeError.
keras.config.disable_flash_attention()\n5. Custom HBM Management:\n160Β³ volumes plus TransUNet push even the v4-32's generous memory boundaries. I wrote a custom Keras callback using jax.local_devices()[i].memory_stats() to monitor actual HBM usage dynamically during training.
6. Preemption Resilience:\nKaggle TPU sessions get preempted frequently. An auto-resume script that saves checkpoint weights and CSV logs per-epoch is mandatory, not optional.
\n---### Competition Strategy
\n| Category | \nTool | \n
|---|---|
| Framework | \nKeras 3 + JAX | \n
| Hardware | \nKaggle TPU v4-32 | \n
| Model library | \nmedicai (TransUNet, losses, metrics, transforms) | \n
| Data pipeline | \nTFRecords + tf.data | \n
| 3D transforms | \nmedicai.transforms (MONAI-inspired) | \n
| Inference | \nSlidingWindowInference + TTA | \n
| Post-processing | \nscipy.ndimage | \n
I will definitely focus more on reviewing open-source projects first, and I will continue to give my best, even in difficult situations.\nLastly, I would really appreciate it if you could leave your advice in the comments.
", "messages": [ { "id": "3416203", "postDate": "03/02/2026 09:47:44", "content": "I'll be honest. I didn't do as well as I had hoped. I was aiming much higher like top 10 But the result didn't reflect that effort, and I think I know whys.
\nI'll be real about where I was coming from: I'm relatively new to competitive machine learning, TPUs, and JAX all at once. There was a steep learning curve just to get anything running, let alone running well.
\nI won't pretend the motivation was purely academic either. Top 10 meant prize money, and I genuinely needed it. That drove me to put in a lot of hours β probably more than I should have at certain points. I tried hard. I just didn't have the experience yet to convert that effort into a top result.
\nAnd that's okay. I learned an enormous amount. I went from knowing nothing about distributed TPU training to having a working end-to-end 3D segmentation pipeline. That's real progress, even if the leaderboard didn't fully show it. Next time, I'll come in with that foundation already in place.
\nA big thank you to Innat β his notebook was what got me started with the medicai framework and TPU-compatible 3D training.
This training was done as part of the TRC (TPU Research Cloud) program. I'm genuinely grateful to Google for providing access to TPU v4-32 hardware through TRC training 160Β³ 3D models at this scale would have been completely infeasible otherwise. Thank you.
\nI built a 3D segmentation pipeline to detect ancient papyrus surfaces from CT scans of the Herculaneum scrolls. My final solution uses TransUNet [1] with a SEResNeXt [2] encoder, trained on 160Β³ voxel patches with a Tversky [3] + CenterlineDice [4] loss on TPU v4-32. Inference combines sliding window, 6-fold TTA, dual checkpoint ensemble, and morphological post-processing. My biggest personal takeaway was learning to train 3D medical segmentation models end-to-end on TPUs with JAX/Keras.
\n
End-to-end pipeline: 3D data prep β TransUNet training β dual checkpoint blending β post-processing β binary ink mask.
\nInput (160Β³Γ1) β SEResNeXt CNN Encoder β Transformer Bottleneck β UNet Decoder β Softmax (3 classes)\nWhy this architecture?
\nI experimented with two encoder sizes:
\n| Variant | \nTraining | \nInference | \n
|---|---|---|
seresnext50 | \nβ (train-tpu.py) | \nβ | \n
seresnext101 | \nβ (final model) | \nβ (inference-kaggle6.py) | \n
| Parameter | \nValue | \n
|---|---|
| Hardware | \nKaggle TPU v4-32 (32 cores) | \n
| Framework | \nKeras 3 + JAX backend | \n
| Batch size | \ntotal_devices // 2 (16 on TPU v4-32) | \n
| Epochs | \n800 (CosineDecay + EarlyStopping patience=60) | \n
| Base LR | \n2e-4 with sqrt batch scaling | \n
| LR Schedule | \nSingle-cycle CosineDecay, 5-epoch warmup, alpha=0.005 | \n
| Optimizer | \nAdamW (weight_decay=1e-5, clipnorm=1.0) | \n
| Loss | \nSparseTverskyLoss + 0.2 Γ SparseCenterlineDiceLoss | \n
| Val split | \n5 TFRecord shards (~30 samples), seed=42 | \n
# Tversky loss β recall-focused (alpha=0.3 FP, beta=0.7 FN)\ntversky = SparseTverskyLoss(alpha=0.3, beta=0.7, ignore_class_ids=2)\n\n# CenterlineDice β topology preservation via soft-skeletonisation\ncldice = SparseCenterlineDiceLoss(target_class_ids=1, ignore_class_ids=2, iters=20)\n\n# Combined\nloss = tversky + 0.2 * cldice\nRationale: Tversky with Ξ²=0.7 penalizes missed surfaces (FN) more than false positives [3], aligning with the competition metric's sensitivity to gaps. CenterlineDice [4] preserves topological continuity via soft-skeletonisation on predictions β directly targeting the TopoScore sub-metric. The ignore_class_ids=2 automatically masks out unlabeled regions.
Data is stored as TFRecords (780 samples). Training augmentation includes:
\nmin_valid_ratio=0.3 β rejects patches >70% ignore-class)nonzero=True) β adapts to each volume's actual intensity range rather than fixed [0,255]β[0,1] scalingEarlier experiments used CosineDecayRestarts(m_mul=0.9), but the learning rate collapsed too early (~epoch 300). Switching to a single-cycle CosineDecay over the full 800-epoch budget produced much more stable convergence.
SlidingWindowInference(\n roi_size=(160, 160, 160), # must match ViT positional embeddings (5Β³=125 patches)\n overlap=0.5,\n blend_mode=\"gaussian\", # smooth boundary artifacts\n)\nTTA_AXES = [(), (0,), (1,), (2,), (0,1), (1,2)]\n# Original + 3 single-axis flips + 2 pair flips β averaged\npred = 0.80 * pred_best_val_dice + 0.20 * pred_best_train_loss\nTwo checkpoints from the same training run β best_val_dice generalizes better, best_train_loss is more expressive. Blending consistently improved over either single checkpoint.
ink_prob -= 0.15 Γ ignore_prob β leverages class-2 predictions to reduce false positives in uncertain regionsThis competition was my first serious experience training on TPUs. It required several specific codebase adjustments to run successfully:
\n1. JAX backend over TensorFlow:\nJAX is significantly more stable and faster on TPUs than TF eager mode.
\nimport os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n2. Disabling Preallocation (Critical for Inference):\nBy default, JAX reserves ~75% of VRAM upfront. For a 160Β³ 3D volume on inference, this guarantees an OOM error before the model even spins up.
\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\nos.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n3. Distributed Training Strategy:\nScaling across 32 cores requires the DataParallel distribution. Tuning the batch size is critical; too large wastes HBM, too small underutilizes the v4-32 hardware.
import jax\nimport keras\n\njax.distributed.initialize()\ndevices = keras.distribution.list_devices()\ndata_parallel = keras.distribution.DataParallel(devices=devices)\nkeras.distribution.set_distribution(data_parallel)\ndata_parallel.auto_shard_dataset = False # Essential for custom tf.data pipelines\n4. Flash Attention Incompatibility:\nDistributed training over multiple TPU workers currently requires disabling Flash Attention to prevent XlaRuntimeError.
keras.config.disable_flash_attention()\n5. Custom HBM Management:\n160Β³ volumes plus TransUNet push even the v4-32's generous memory boundaries. I wrote a custom Keras callback using jax.local_devices()[i].memory_stats() to monitor actual HBM usage dynamically during training.
6. Preemption Resilience:\nKaggle TPU sessions get preempted frequently. An auto-resume script that saves checkpoint weights and CSV logs per-epoch is mandatory, not optional.
\n---### Competition Strategy
\n| Category | \nTool | \n
|---|---|
| Framework | \nKeras 3 + JAX | \n
| Hardware | \nKaggle TPU v4-32 | \n
| Model library | \nmedicai (TransUNet, losses, metrics, transforms) | \n
| Data pipeline | \nTFRecords + tf.data | \n
| 3D transforms | \nmedicai.transforms (MONAI-inspired) | \n
| Inference | \nSlidingWindowInference + TTA | \n
| Post-processing | \nscipy.ndimage | \n
I will definitely focus more on reviewing open-source projects first, and I will continue to give my best, even in difficult situations.\nLastly, I would really appreciate it if you could leave your advice in the comments.
", "rawMarkdown": "# 123rd place β not where I aimed, but here's what I learned\n\n- **Competition:** [Vesuvius Challenge Surface Detection](https://www.kaggle.com/competitions/vesuvius-challenge-surface-detection)\n- **Result:** 123rd / 1427 Β· Score: 0.579 Β· π₯ Bronze Medal\n- **Training code:** [kaggle.com/code/osman0/vesuvius-surface-train-tpu](https://www.kaggle.com/code/osman0/vesuvius-surface-train-tpu)\n\n---\n\n## Personal Reflection\n\nI'll be honest. I didn't do as well as I had hoped. I was aiming much higher like top 10 But the result didn't reflect that effort, and I think I know whys.\n\nI'll be real about where I was coming from: I'm relatively new to competitive machine learning, TPUs, and JAX all at once. There was a steep learning curve just to get anything running, let alone running well.\n\nI won't pretend the motivation was purely academic either. Top 10 meant prize money, and I genuinely needed it. That drove me to put in a lot of hours β probably more than I should have at certain points. I tried hard. I just didn't have the experience yet to convert that effort into a top result.\n\nAnd that's okay. I learned an enormous amount. I went from knowing nothing about distributed TPU training to having a working end-to-end 3D segmentation pipeline. That's real progress, even if the leaderboard didn't fully show it. Next time, I'll come in with that foundation already in place.\n\nA big thank you to [Innat](https://www.kaggle.com/ipythonx) β his [notebook](https://www.kaggle.com/code/ipythonx/train-vesuvius-surface-3d-detection-on-tpu) was what got me started with the `medicai` framework and TPU-compatible 3D training.\n\n---\n\n## Acknowledgements\n\nThis training was done as part of the **TRC (TPU Research Cloud) program**. I'm genuinely grateful to Google for providing access to TPU v4-32 hardware through TRC training 160Β³ 3D models at this scale would have been completely infeasible otherwise. Thank you.\n\n---\n\n## TL;DR\n\nI built a 3D segmentation pipeline to detect ancient papyrus surfaces from CT scans of the Herculaneum scrolls. My final solution uses **TransUNet** [[1]](#references) with a **SEResNeXt** [[2]](#references) encoder, trained on **160Β³ voxel** patches with a **Tversky** [[3]](#references) **+ CenterlineDice** [[4]](#references) loss on **TPU v4-32**. Inference combines **sliding window**, **6-fold TTA**, **dual checkpoint ensemble**, and **morphological post-processing**. My biggest personal takeaway was learning to train 3D medical segmentation models end-to-end on TPUs with JAX/Keras.\n\n---\n\n## Solution Architecture\n\n\n\n*End-to-end pipeline: 3D data prep β TransUNet training β dual checkpoint blending β post-processing β binary ink mask.*\n\n### Model: TransUNet + SEResNeXt\n\n```\nInput (160Β³Γ1) β SEResNeXt CNN Encoder β Transformer Bottleneck β UNet Decoder β Softmax (3 classes)\n```\n\n**Why this architecture?**\n- **160Β³ input** β 4.6Γ larger receptive field than 96Β³, critical for thin papyrus surfaces spanning hundreds of voxels\n- **SEResNeXt encoder** [[2]](#references) β squeeze-excitation blocks provide channel-wise attention, proven on 3D CT data\n- **Hybrid CNN + Transformer** [[1]](#references) β CNN encoder captures local features, transformer bottleneck captures global context\n- **HBM-efficient** β CNN encoder is lighter on TPU HBM than Swin attention maps at 160Β³ resolution\n\nI experimented with two encoder sizes:\n\n| Variant | Training | Inference |\n|---------|----------|-----------|\n| `seresnext50` | β (train-tpu.py) | β |\n| `seresnext101` | β (final model) | β (inference-kaggle6.py) |\n\n---\n\n## Training Details\n\n| Parameter | Value |\n|-----------|-------|\n| **Hardware** | Kaggle TPU v4-32 (32 cores) |\n| **Framework** | Keras 3 + JAX backend |\n| **Batch size** | `total_devices // 2` (16 on TPU v4-32) |\n| **Epochs** | 800 (CosineDecay + EarlyStopping patience=60) |\n| **Base LR** | 2e-4 with sqrt batch scaling |\n| **LR Schedule** | Single-cycle CosineDecay, 5-epoch warmup, alpha=0.005 |\n| **Optimizer** | AdamW (weight_decay=1e-5, clipnorm=1.0) |\n| **Loss** | SparseTverskyLoss + 0.2 Γ SparseCenterlineDiceLoss |\n| **Val split** | 5 TFRecord shards (~30 samples), seed=42 |\n\n### Loss Function\n\n```python\n# Tversky loss β recall-focused (alpha=0.3 FP, beta=0.7 FN)\ntversky = SparseTverskyLoss(alpha=0.3, beta=0.7, ignore_class_ids=2)\n\n# CenterlineDice β topology preservation via soft-skeletonisation\ncldice = SparseCenterlineDiceLoss(target_class_ids=1, ignore_class_ids=2, iters=20)\n\n# Combined\nloss = tversky + 0.2 * cldice\n```\n\n**Rationale:** Tversky with Ξ²=0.7 penalizes missed surfaces (FN) more than false positives [[3]](#references), aligning with the competition metric's sensitivity to gaps. CenterlineDice [[4]](#references) preserves topological continuity via soft-skeletonisation on predictions β directly targeting the TopoScore sub-metric. The `ignore_class_ids=2` automatically masks out unlabeled regions.\n\n### Data Pipeline & Augmentation\n\nData is stored as TFRecords (780 samples). Training augmentation includes:\n\n- **RandSpatialCrop** (160Β³, `min_valid_ratio=0.3` β rejects patches >70% ignore-class)\n- **RandFlip** (all 3 axes, prob=0.5)\n- **RandRotate90** (prob=0.4) + **RandRotate** (factor=0.2, prob=0.25)\n- **NormalizeIntensity** (`nonzero=True`) β adapts to each volume's actual intensity range rather than fixed [0,255]β[0,1] scaling\n- **RandShiftIntensity** (offsets=0.10, prob=0.5)\n- **RandCutOut** (volume mode, prob=0.2) β occlusion regularisation\n\n### LR Schedule Evolution\n\nEarlier experiments used `CosineDecayRestarts(m_mul=0.9)`, but the learning rate collapsed too early (~epoch 300). Switching to a **single-cycle CosineDecay** over the full 800-epoch budget produced much more stable convergence.\n\n### TPU-Specific Implementation\n\n- **DataParallel** distribution across 32 TPU v4 cores with auto-sharding\n- **Flash attention disabled** β required for distributed training compatibility\n- **Custom HBM Memory Monitor callback** β logs TPU HBM usage and host RAM every N epochs\n- **Auto-resume from preemption** β checkpoint + CSV log continuation for Kaggle TPU interrupts\n\n---\n\n## Inference Pipeline\n\n### Sliding Window Inference\n```python\nSlidingWindowInference(\n roi_size=(160, 160, 160), # must match ViT positional embeddings (5Β³=125 patches)\n overlap=0.5,\n blend_mode=\"gaussian\", # smooth boundary artifacts\n)\n```\n\n### Test-Time Augmentation (6-fold)\n```python\nTTA_AXES = [(), (0,), (1,), (2,), (0,1), (1,2)]\n# Original + 3 single-axis flips + 2 pair flips β averaged\n```\n\n### Dual Checkpoint Blend\n```python\npred = 0.80 * pred_best_val_dice + 0.20 * pred_best_train_loss\n```\n\nTwo checkpoints from the same training run β `best_val_dice` generalizes better, `best_train_loss` is more expressive. Blending consistently improved over either single checkpoint.\n\n### Post-Processing\n1. **Ignore-class suppression**: `ink_prob -= 0.15 Γ ignore_prob` β leverages class-2 predictions to reduce false positives in uncertain regions\n2. **Thresholding**: prob β₯ 0.5 β binary mask\n3. **Morphological closing** (radius=1) β fills small gaps\n4. **Connected component filtering** (min_size=50) β removes noise\n\n---\n## Learnings\n\n### TPU for 3D Medical Segmentation\n\nThis competition was my first serious experience training on TPUs. It required several specific codebase adjustments to run successfully:\n\n**1. JAX backend over TensorFlow:**\nJAX is significantly more stable and faster on TPUs than TF eager mode.\n```python\nimport os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n```\n\n**2. Disabling Preallocation (Critical for Inference):**\nBy default, JAX reserves ~75% of VRAM upfront. For a 160Β³ 3D volume on inference, this guarantees an OOM error before the model even spins up.\n```python\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\nos.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n```\n\n**3. Distributed Training Strategy:**\nScaling across 32 cores requires the `DataParallel` distribution. Tuning the batch size is critical; too large wastes HBM, too small underutilizes the v4-32 hardware.\n```python\nimport jax\nimport keras\n\njax.distributed.initialize()\ndevices = keras.distribution.list_devices()\ndata_parallel = keras.distribution.DataParallel(devices=devices)\nkeras.distribution.set_distribution(data_parallel)\ndata_parallel.auto_shard_dataset = False # Essential for custom tf.data pipelines\n```\n\n**4. Flash Attention Incompatibility:**\nDistributed training over multiple TPU workers currently requires disabling Flash Attention to prevent `XlaRuntimeError`.\n```python\nkeras.config.disable_flash_attention()\n```\n\n**5. Custom HBM Management:**\n160Β³ volumes plus TransUNet push even the v4-32's generous memory boundaries. I wrote a custom Keras callback using `jax.local_devices()[i].memory_stats()` to monitor actual HBM usage dynamically during training.\n\n**6. Preemption Resilience:**\nKaggle TPU sessions get preempted frequently. An auto-resume script that saves checkpoint weights and CSV logs per-epoch is mandatory, not optional.\n\n---### Competition Strategy\n- **~60 submissions** β aggressive but deliberate experimentation\n- **Fast iteration** β TPU speed enabled multiple experiments per day\n- **Post-processing matters** β the inference pipeline is as important as the model itself\n- **Loss-metric alignment** β designing loss functions that mirror competition metrics (Surface Dice + TopoScore β Tversky + clDice) is crucial\n\n### General Insights\n- In 3D segmentation, **input resolution** often matters more than model complexity\n- **Ensemble techniques** have low cost but high return\n- **Weight loading issues** across library versions are common β always save version metadata alongside checkpoints\n\n---\n## Tech Stack\n\n| Category | Tool |\n|----------|------|\n| Framework | Keras 3 + JAX |\n| Hardware | Kaggle TPU v4-32 |\n| Model library | medicai (TransUNet, losses, metrics, transforms) |\n| Data pipeline | TFRecords + tf.data |\n| 3D transforms | medicai.transforms (MONAI-inspired) |\n| Inference | SlidingWindowInference + TTA |\n| Post-processing | scipy.ndimage |\n\n---\n## Conclusion\n\nI will definitely focus more on reviewing open-source projects first, and I will continue to give my best, even in difficult situations.\nLastly, I would really appreciate it if you could leave your advice in the comments.", "votes": null }, { "id": "3416232", "postDate": "03/02/2026 11:03:18", "content": "@osman0 Great work.
\nThroughout this competition, we identified several missing components that could strengthen medicai, and we managed to incorporate some of them within the limited timeframe.
We also gained a deeper appreciation for the generalization strength of the nnUNet framework. Whether we like it or not, nnUNet clearly outperformed others on the leaderboard. Iβm planning to integrate this framework into medicai through a mid- to high-level API to make it easier to use (ticket). However, this wonβt be a straightforward task. We need to ensure a proper and validated translation of nnUNetβs design into medicai, while maintaining overall API consistency.
Now, you've understanding about segmentation, I would suggest to explote 5th place solution, quite interesting and solid kaggle thing.
\nAnother valuable lesson from this competition was learning how to effectively interact with LLMs. Many participants found them incredibly helpful, and understanding how to communicate with them properly made a significant difference; discussion.
", "rawMarkdown": "osman0 Great work. \n\nThroughout this competition, we identified several missing components that could strengthen `medicai`, and we managed to incorporate some of them within the limited timeframe.\n\nWe also gained a deeper appreciation for the generalization strength of the `nnUNet` framework. Whether we like it or not, `nnUNet` clearly outperformed others on the leaderboard. Iβm planning to integrate this framework into `medicai` through a mid- to high-level API to make it easier to use ([ticket](https://github.com/innat/medic-ai/issues/30)). However, this wonβt be a straightforward task. We need to ensure a proper and validated translation of `nnUNet`βs design into `medicai`, while maintaining overall API consistency.\n\nNow, you've understanding about segmentation, I would suggest to explote [5th place](https://www.kaggle.com/competitions/vesuvius-challenge-surface-detection/writeups/5th-place-solution) solution, quite interesting and solid kaggle thing. \n\nAnother valuable lesson from this competition was learning how to effectively interact with LLMs. Many participants found them incredibly helpful, and understanding how to communicate with them properly made a significant difference; [discussion](https://www.kaggle.com/competitions/vesuvius-challenge-surface-detection/writeups/bronze-medal-chatgpt-vibe-coding).", "votes": null }, { "id": "3416238", "postDate": "03/02/2026 11:34:17", "content": "Thank you so much for your kind words and for reaching out!
\nI especially want to thank you because you and your work were incredibly helpful in my journey to learn about TPUs. I also had the opportunity to try out some of the models from your medic-ai repository, and it was a truly amazing experience for me.
\nI really appreciate your suggestion about the 5th place solutionβI will definitely take a close look at it. and olso def I look for your nnUNet framework.
\nThanks again.
", "rawMarkdown": "Thank you so much for your kind words and for reaching out!\n\nI especially want to thank you because you and your work were incredibly helpful in my journey to learn about TPUs. I also had the opportunity to try out some of the models from your medic-ai repository, and it was a truly amazing experience for me.\n\nI really appreciate your suggestion about the 5th place solutionβI will definitely take a close look at it. and olso def I look for your nnUNet framework. \n\nThanks again.", "votes": null } ], "comments": [ { "id": 3416232, "author_name": "ipythonx", "author_url": "", "post_date": "03/02/2026 11:03:18", "content": "@osman0 Great work.
\nThroughout this competition, we identified several missing components that could strengthen medicai, and we managed to incorporate some of them within the limited timeframe.
We also gained a deeper appreciation for the generalization strength of the nnUNet framework. Whether we like it or not, nnUNet clearly outperformed others on the leaderboard. Iβm planning to integrate this framework into medicai through a mid- to high-level API to make it easier to use (ticket). However, this wonβt be a straightforward task. We need to ensure a proper and validated translation of nnUNetβs design into medicai, while maintaining overall API consistency.
Now, you've understanding about segmentation, I would suggest to explote 5th place solution, quite interesting and solid kaggle thing.
\nAnother valuable lesson from this competition was learning how to effectively interact with LLMs. Many participants found them incredibly helpful, and understanding how to communicate with them properly made a significant difference; discussion.
", "votes": null, "replies": [ { "id": 3416238, "author_name": "osman0", "author_url": "", "post_date": "03/02/2026 11:34:17", "content": "Thank you so much for your kind words and for reaching out!
\nI especially want to thank you because you and your work were incredibly helpful in my journey to learn about TPUs. I also had the opportunity to try out some of the models from your medic-ai repository, and it was a truly amazing experience for me.
\nI really appreciate your suggestion about the 5th place solutionβI will definitely take a close look at it. and olso def I look for your nnUNet framework.
\nThanks again.
", "votes": null, "replies": [] } ] } ], "raw_markdown_by_id": { "3416203": "# 123rd place β not where I aimed, but here's what I learned\n\n- **Competition:** [Vesuvius Challenge Surface Detection](https://www.kaggle.com/competitions/vesuvius-challenge-surface-detection)\n- **Result:** 123rd / 1427 Β· Score: 0.579 Β· π₯ Bronze Medal\n- **Training code:** [kaggle.com/code/osman0/vesuvius-surface-train-tpu](https://www.kaggle.com/code/osman0/vesuvius-surface-train-tpu)\n\n---\n\n## Personal Reflection\n\nI'll be honest. I didn't do as well as I had hoped. I was aiming much higher like top 10 But the result didn't reflect that effort, and I think I know whys.\n\nI'll be real about where I was coming from: I'm relatively new to competitive machine learning, TPUs, and JAX all at once. There was a steep learning curve just to get anything running, let alone running well.\n\nI won't pretend the motivation was purely academic either. Top 10 meant prize money, and I genuinely needed it. That drove me to put in a lot of hours β probably more than I should have at certain points. I tried hard. I just didn't have the experience yet to convert that effort into a top result.\n\nAnd that's okay. I learned an enormous amount. I went from knowing nothing about distributed TPU training to having a working end-to-end 3D segmentation pipeline. That's real progress, even if the leaderboard didn't fully show it. Next time, I'll come in with that foundation already in place.\n\nA big thank you to [Innat](https://www.kaggle.com/ipythonx) β his [notebook](https://www.kaggle.com/code/ipythonx/train-vesuvius-surface-3d-detection-on-tpu) was what got me started with the `medicai` framework and TPU-compatible 3D training.\n\n---\n\n## Acknowledgements\n\nThis training was done as part of the **TRC (TPU Research Cloud) program**. I'm genuinely grateful to Google for providing access to TPU v4-32 hardware through TRC training 160Β³ 3D models at this scale would have been completely infeasible otherwise. Thank you.\n\n---\n\n## TL;DR\n\nI built a 3D segmentation pipeline to detect ancient papyrus surfaces from CT scans of the Herculaneum scrolls. My final solution uses **TransUNet** [[1]](#references) with a **SEResNeXt** [[2]](#references) encoder, trained on **160Β³ voxel** patches with a **Tversky** [[3]](#references) **+ CenterlineDice** [[4]](#references) loss on **TPU v4-32**. Inference combines **sliding window**, **6-fold TTA**, **dual checkpoint ensemble**, and **morphological post-processing**. My biggest personal takeaway was learning to train 3D medical segmentation models end-to-end on TPUs with JAX/Keras.\n\n---\n\n## Solution Architecture\n\n\n\n*End-to-end pipeline: 3D data prep β TransUNet training β dual checkpoint blending β post-processing β binary ink mask.*\n\n### Model: TransUNet + SEResNeXt\n\n```\nInput (160Β³Γ1) β SEResNeXt CNN Encoder β Transformer Bottleneck β UNet Decoder β Softmax (3 classes)\n```\n\n**Why this architecture?**\n- **160Β³ input** β 4.6Γ larger receptive field than 96Β³, critical for thin papyrus surfaces spanning hundreds of voxels\n- **SEResNeXt encoder** [[2]](#references) β squeeze-excitation blocks provide channel-wise attention, proven on 3D CT data\n- **Hybrid CNN + Transformer** [[1]](#references) β CNN encoder captures local features, transformer bottleneck captures global context\n- **HBM-efficient** β CNN encoder is lighter on TPU HBM than Swin attention maps at 160Β³ resolution\n\nI experimented with two encoder sizes:\n\n| Variant | Training | Inference |\n|---------|----------|-----------|\n| `seresnext50` | β (train-tpu.py) | β |\n| `seresnext101` | β (final model) | β (inference-kaggle6.py) |\n\n---\n\n## Training Details\n\n| Parameter | Value |\n|-----------|-------|\n| **Hardware** | Kaggle TPU v4-32 (32 cores) |\n| **Framework** | Keras 3 + JAX backend |\n| **Batch size** | `total_devices // 2` (16 on TPU v4-32) |\n| **Epochs** | 800 (CosineDecay + EarlyStopping patience=60) |\n| **Base LR** | 2e-4 with sqrt batch scaling |\n| **LR Schedule** | Single-cycle CosineDecay, 5-epoch warmup, alpha=0.005 |\n| **Optimizer** | AdamW (weight_decay=1e-5, clipnorm=1.0) |\n| **Loss** | SparseTverskyLoss + 0.2 Γ SparseCenterlineDiceLoss |\n| **Val split** | 5 TFRecord shards (~30 samples), seed=42 |\n\n### Loss Function\n\n```python\n# Tversky loss β recall-focused (alpha=0.3 FP, beta=0.7 FN)\ntversky = SparseTverskyLoss(alpha=0.3, beta=0.7, ignore_class_ids=2)\n\n# CenterlineDice β topology preservation via soft-skeletonisation\ncldice = SparseCenterlineDiceLoss(target_class_ids=1, ignore_class_ids=2, iters=20)\n\n# Combined\nloss = tversky + 0.2 * cldice\n```\n\n**Rationale:** Tversky with Ξ²=0.7 penalizes missed surfaces (FN) more than false positives [[3]](#references), aligning with the competition metric's sensitivity to gaps. CenterlineDice [[4]](#references) preserves topological continuity via soft-skeletonisation on predictions β directly targeting the TopoScore sub-metric. The `ignore_class_ids=2` automatically masks out unlabeled regions.\n\n### Data Pipeline & Augmentation\n\nData is stored as TFRecords (780 samples). Training augmentation includes:\n\n- **RandSpatialCrop** (160Β³, `min_valid_ratio=0.3` β rejects patches >70% ignore-class)\n- **RandFlip** (all 3 axes, prob=0.5)\n- **RandRotate90** (prob=0.4) + **RandRotate** (factor=0.2, prob=0.25)\n- **NormalizeIntensity** (`nonzero=True`) β adapts to each volume's actual intensity range rather than fixed [0,255]β[0,1] scaling\n- **RandShiftIntensity** (offsets=0.10, prob=0.5)\n- **RandCutOut** (volume mode, prob=0.2) β occlusion regularisation\n\n### LR Schedule Evolution\n\nEarlier experiments used `CosineDecayRestarts(m_mul=0.9)`, but the learning rate collapsed too early (~epoch 300). Switching to a **single-cycle CosineDecay** over the full 800-epoch budget produced much more stable convergence.\n\n### TPU-Specific Implementation\n\n- **DataParallel** distribution across 32 TPU v4 cores with auto-sharding\n- **Flash attention disabled** β required for distributed training compatibility\n- **Custom HBM Memory Monitor callback** β logs TPU HBM usage and host RAM every N epochs\n- **Auto-resume from preemption** β checkpoint + CSV log continuation for Kaggle TPU interrupts\n\n---\n\n## Inference Pipeline\n\n### Sliding Window Inference\n```python\nSlidingWindowInference(\n roi_size=(160, 160, 160), # must match ViT positional embeddings (5Β³=125 patches)\n overlap=0.5,\n blend_mode=\"gaussian\", # smooth boundary artifacts\n)\n```\n\n### Test-Time Augmentation (6-fold)\n```python\nTTA_AXES = [(), (0,), (1,), (2,), (0,1), (1,2)]\n# Original + 3 single-axis flips + 2 pair flips β averaged\n```\n\n### Dual Checkpoint Blend\n```python\npred = 0.80 * pred_best_val_dice + 0.20 * pred_best_train_loss\n```\n\nTwo checkpoints from the same training run β `best_val_dice` generalizes better, `best_train_loss` is more expressive. Blending consistently improved over either single checkpoint.\n\n### Post-Processing\n1. **Ignore-class suppression**: `ink_prob -= 0.15 Γ ignore_prob` β leverages class-2 predictions to reduce false positives in uncertain regions\n2. **Thresholding**: prob β₯ 0.5 β binary mask\n3. **Morphological closing** (radius=1) β fills small gaps\n4. **Connected component filtering** (min_size=50) β removes noise\n\n---\n## Learnings\n\n### TPU for 3D Medical Segmentation\n\nThis competition was my first serious experience training on TPUs. It required several specific codebase adjustments to run successfully:\n\n**1. JAX backend over TensorFlow:**\nJAX is significantly more stable and faster on TPUs than TF eager mode.\n```python\nimport os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n```\n\n**2. Disabling Preallocation (Critical for Inference):**\nBy default, JAX reserves ~75% of VRAM upfront. For a 160Β³ 3D volume on inference, this guarantees an OOM error before the model even spins up.\n```python\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\nos.environ[\"XLA_PYTHON_CLIENT_ALLOCATOR\"] = \"platform\"\n```\n\n**3. Distributed Training Strategy:**\nScaling across 32 cores requires the `DataParallel` distribution. Tuning the batch size is critical; too large wastes HBM, too small underutilizes the v4-32 hardware.\n```python\nimport jax\nimport keras\n\njax.distributed.initialize()\ndevices = keras.distribution.list_devices()\ndata_parallel = keras.distribution.DataParallel(devices=devices)\nkeras.distribution.set_distribution(data_parallel)\ndata_parallel.auto_shard_dataset = False # Essential for custom tf.data pipelines\n```\n\n**4. Flash Attention Incompatibility:**\nDistributed training over multiple TPU workers currently requires disabling Flash Attention to prevent `XlaRuntimeError`.\n```python\nkeras.config.disable_flash_attention()\n```\n\n**5. Custom HBM Management:**\n160Β³ volumes plus TransUNet push even the v4-32's generous memory boundaries. I wrote a custom Keras callback using `jax.local_devices()[i].memory_stats()` to monitor actual HBM usage dynamically during training.\n\n**6. Preemption Resilience:**\nKaggle TPU sessions get preempted frequently. An auto-resume script that saves checkpoint weights and CSV logs per-epoch is mandatory, not optional.\n\n---### Competition Strategy\n- **~60 submissions** β aggressive but deliberate experimentation\n- **Fast iteration** β TPU speed enabled multiple experiments per day\n- **Post-processing matters** β the inference pipeline is as important as the model itself\n- **Loss-metric alignment** β designing loss functions that mirror competition metrics (Surface Dice + TopoScore β Tversky + clDice) is crucial\n\n### General Insights\n- In 3D segmentation, **input resolution** often matters more than model complexity\n- **Ensemble techniques** have low cost but high return\n- **Weight loading issues** across library versions are common β always save version metadata alongside checkpoints\n\n---\n## Tech Stack\n\n| Category | Tool |\n|----------|------|\n| Framework | Keras 3 + JAX |\n| Hardware | Kaggle TPU v4-32 |\n| Model library | medicai (TransUNet, losses, metrics, transforms) |\n| Data pipeline | TFRecords + tf.data |\n| 3D transforms | medicai.transforms (MONAI-inspired) |\n| Inference | SlidingWindowInference + TTA |\n| Post-processing | scipy.ndimage |\n\n---\n## Conclusion\n\nI will definitely focus more on reviewing open-source projects first, and I will continue to give my best, even in difficult situations.\nLastly, I would really appreciate it if you could leave your advice in the comments.", "3416232": "osman0 Great work. \n\nThroughout this competition, we identified several missing components that could strengthen `medicai`, and we managed to incorporate some of them within the limited timeframe.\n\nWe also gained a deeper appreciation for the generalization strength of the `nnUNet` framework. Whether we like it or not, `nnUNet` clearly outperformed others on the leaderboard. Iβm planning to integrate this framework into `medicai` through a mid- to high-level API to make it easier to use ([ticket](https://github.com/innat/medic-ai/issues/30)). However, this wonβt be a straightforward task. We need to ensure a proper and validated translation of `nnUNet`βs design into `medicai`, while maintaining overall API consistency.\n\nNow, you've understanding about segmentation, I would suggest to explote [5th place](https://www.kaggle.com/competitions/vesuvius-challenge-surface-detection/writeups/5th-place-solution) solution, quite interesting and solid kaggle thing. \n\nAnother valuable lesson from this competition was learning how to effectively interact with LLMs. Many participants found them incredibly helpful, and understanding how to communicate with them properly made a significant difference; [discussion](https://www.kaggle.com/competitions/vesuvius-challenge-surface-detection/writeups/bronze-medal-chatgpt-vibe-coding).", "3416238": "Thank you so much for your kind words and for reaching out!\n\nI especially want to thank you because you and your work were incredibly helpful in my journey to learn about TPUs. I also had the opportunity to try out some of the models from your medic-ai repository, and it was a truly amazing experience for me.\n\nI really appreciate your suggestion about the 5th place solutionβI will definitely take a close look at it. and olso def I look for your nnUNet framework. \n\nThanks again." }, "source": "meta" }