GeoBrain/first-break-picking

First-Break Picking HF Release

This directory is a compact release package for first-break-picking model weights. It keeps only the best checkpoint and the exact training config for each run.

Contents

• 60 best.pt checkpoint files.
• 60 config.yaml files.
• No intermediate epoch_*.pt checkpoints.
• Current package size: about 4.1 GB.

Two experiment groups are included:

Directory	Meaning	Runs
first_break_picking/	Models trained on all available first-break SEG-Y pairs together	12
first_break_picking_single_dataset_grouped/by_dataset/	Models trained separately on each individual SEG-Y pair	48

Directory Layout

first_break_picking_hf_release/
  README.md
  first_break_picking/
    first_break_pick_<model>_geomseg_seed<seed>/
      config.yaml
      checkpoints/
        best.pt
  first_break_picking_single_dataset_grouped/
    by_dataset/
      <dataset>/
        <model>/
          seed<seed>/
            config.yaml
            checkpoints/
              best.pt

The multi-dataset runs use:

models: unet, res_unet, atten_unet, dncnn_seg
seeds: 42, 43, 44

The single-dataset runs use:

datasets: brunswick_valid, dongbei, halfmile_valid, lalor_valid
models: unet, res_unet, atten_unet, dncnn_seg
seeds: 42, 43, 44

Task Definition

The benchmark treats first-break picking as binary step-mask segmentation.

• Input: single-channel SEG-Y amplitude patches.
• Label: binary step mask, with 0 before the first break and 1 from the first-break sample onward.
• Prediction: single-channel logits with the same spatial shape as the mask.
• Pick extraction: the first time index where sigmoid(logit) >= 0.5.

Dataset Configuration

All configs point to:

data.root: /home/dataset-local/dataset/first_break_picking/segy_with_masks
data.data_dir: data
data.label_dir: label

The available SEG-Y input files are:

Brunswick_valid.sgy
Dongbei.segy
Halfmile_valid.sgy
Lalor_valid.sgy

For multi-dataset runs, data.files: null, so all SEG-Y pairs under data/ are used. For single-dataset runs, data.files contains exactly one input SEG-Y filename.

Common data parameters:

Parameter	Value
label_threshold	0.5
prediction_threshold	0.5
validate_labels	true
label_check_traces	2048
max_patches_per_split	null
split.train	0.8
split.val	0.1
split.test	0.1
split.shuffle_ffids	true

Gather segmentation parameters:

Parameter	Value
gather_segment.enabled	true
gather_segment.line_id_header	INLINE_3D
gather_segment.infer_line_from_geometry	true
gather_segment.distance_floor	1000.0
gather_segment.median_multiplier	5.0

Patch and loader parameters:

Parameter	Multi-dataset	Single-dataset
patch.trace	128	128
patch.time	512	512
patch.trace_stride	64	64
patch.time_stride	256	256
loader.batch_size	64	64
loader.num_workers	4	1
loader.pin_memory	true	true

Preprocessing

Parameter	Value
normalize_mode	max_abs
normalize_scope	gather
clip_percentile	99.5
normalize_eps	1.0e-6

Only input amplitudes are normalized. Labels remain binary segmentation targets; invalid or padded areas are ignored by the loss and metrics in the training code.

Models

Model	Parameters
unet	in_channels=1, out_channels=1, base_channels=32, depth=4
res_unet	in_channels=1, out_channels=1, base_channels=32, depth=4
atten_unet	in_channels=1, out_channels=1, base_channels=32, depth=4
dncnn_seg	in_channels=1, out_channels=1, depth=17, base_channels=64, kernel_size=3

Training Parameters

Parameter	Value
train.epochs	20
train.grad_clip	1.0
train.log_interval	20
train.eval_interval	1
train.ckpt_interval	1
train.vis_interval	1
train.resume	null
optimizer	adamw
optimizer.lr	1.0e-4
optimizer.weight_decay	1.0e-5
scheduler	cosine
scheduler.min_lr	1.0e-6

Loss:

type: bce_dice
bce_weight: 0.5
dice_weight: 0.5
smooth: 1.0
pos_weight: null

Metrics:

dice
iou
f1
HitRate1px
HitRate3px
HitRate5px
HitRate7px
HitRate9px
MeanAbsoluteError
RootMeanSquaredError
MeanBiasError
GatherCoverage

All metric thresholds are 0.5.

Loading A Checkpoint

Use the matching config.yaml next to each best.pt to reconstruct the model and preprocessing settings.

import torch

checkpoint_path = "first_break_picking/first_break_pick_unet_geomseg_seed42/checkpoints/best.pt"
checkpoint = torch.load(checkpoint_path, map_location="cpu")
print(checkpoint.keys())

The experiment.output_dir values inside config.yaml are the original training output paths. They are preserved for reproducibility and do not need to match this release directory.

Notes For Hugging Face Upload

This folder is intended to be uploaded as a model-weight release package. The raw SEG-Y dataset should be uploaded separately because it is much larger and has a different structure:

segy_with_masks/
  data/
  label/

For this release package, upload the entire first_break_picking_hf_release/ directory so each best.pt remains next to its exact config.yaml.

---

初至拾取 Hugging Face 发布包

该目录是初至拾取模型权重的精简发布包。每一次实验只保留最优模型 checkpoint 和对应的完整训练配置。

内容概览

• 60 个 best.pt checkpoint 文件。
• 60 个 config.yaml 配置文件。
• 不包含中间训练轮次的 epoch_*.pt checkpoint。
• 当前发布包大小约 4.1 GB。

包含两类实验：

目录	含义	运行数量
first_break_picking/	使用全部初至拾取 SEG-Y 数据联合训练的模型	12
first_break_picking_single_dataset_grouped/by_dataset/	在单个 SEG-Y 数据集上分别训练的模型	48

目录结构

first_break_picking_hf_release/
  README.md
  first_break_picking/
    first_break_pick_<model>_geomseg_seed<seed>/
      config.yaml
      checkpoints/
        best.pt
  first_break_picking_single_dataset_grouped/
    by_dataset/
      <dataset>/
        <model>/
          seed<seed>/
            config.yaml
            checkpoints/
              best.pt

联合训练实验包含：

模型: unet, res_unet, atten_unet, dncnn_seg
随机种子: 42, 43, 44

单数据集训练实验包含：

数据集: brunswick_valid, dongbei, halfmile_valid, lalor_valid
模型: unet, res_unet, atten_unet, dncnn_seg
随机种子: 42, 43, 44

任务定义

该 benchmark 将初至拾取建模为二值 step-mask 分割任务。

• 输入：单通道 SEG-Y 振幅 patch。
• 标签：二值 step mask，初至之前为 0，从初至采样点开始为 1。
• 预测：与标签空间尺寸相同的单通道 logits。
• 拾取点提取：取 sigmoid(logit) >= 0.5 的第一个时间采样点作为初至位置。

数据配置

所有配置文件都指向：

data.root: /home/dataset-local/dataset/first_break_picking/segy_with_masks
data.data_dir: data
data.label_dir: label

可用的 SEG-Y 输入文件为：

Brunswick_valid.sgy
Dongbei.segy
Halfmile_valid.sgy
Lalor_valid.sgy

联合训练实验中，data.files: null，表示使用 data/ 下所有 SEG-Y 数据对。单数据集实验中，data.files 只包含一个输入 SEG-Y 文件名。

通用数据参数：

参数	取值
label_threshold	0.5
prediction_threshold	0.5
validate_labels	true
label_check_traces	2048
max_patches_per_split	null
split.train	0.8
split.val	0.1
split.test	0.1
split.shuffle_ffids	true

炮集与接收线切分参数：

参数	取值
gather_segment.enabled	true
gather_segment.line_id_header	INLINE_3D
gather_segment.infer_line_from_geometry	true
gather_segment.distance_floor	1000.0
gather_segment.median_multiplier	5.0

patch 与 DataLoader 参数：

参数	联合训练	单数据集训练
patch.trace	128	128
patch.time	512	512
patch.trace_stride	64	64
patch.time_stride	256	256
loader.batch_size	64	64
loader.num_workers	4	1
loader.pin_memory	true	true

预处理

参数	取值
normalize_mode	max_abs
normalize_scope	gather
clip_percentile	99.5
normalize_eps	1.0e-6

只对输入振幅做归一化。标签保持二值分割目标；无效区域或 padding 区域在训练代码中会被 loss 和 metrics 忽略。

模型

模型	参数
unet	in_channels=1, out_channels=1, base_channels=32, depth=4
res_unet	in_channels=1, out_channels=1, base_channels=32, depth=4
atten_unet	in_channels=1, out_channels=1, base_channels=32, depth=4
dncnn_seg	in_channels=1, out_channels=1, depth=17, base_channels=64, kernel_size=3

训练参数

参数	取值
train.epochs	20
train.grad_clip	1.0
train.log_interval	20
train.eval_interval	1
train.ckpt_interval	1
train.vis_interval	1
train.resume	null
optimizer	adamw
optimizer.lr	1.0e-4
optimizer.weight_decay	1.0e-5
scheduler	cosine
scheduler.min_lr	1.0e-6

损失函数：

type: bce_dice
bce_weight: 0.5
dice_weight: 0.5
smooth: 1.0
pos_weight: null

评价指标：

dice
iou
f1
HitRate1px
HitRate3px
HitRate5px
HitRate7px
HitRate9px
MeanAbsoluteError
RootMeanSquaredError
MeanBiasError
GatherCoverage

所有指标阈值均为 0.5。

加载 Checkpoint

使用每个 best.pt 旁边对应的 config.yaml 来重建模型和预处理设置。

import torch

checkpoint_path = "first_break_picking/first_break_pick_unet_geomseg_seed42/checkpoints/best.pt"
checkpoint = torch.load(checkpoint_path, map_location="cpu")
print(checkpoint.keys())

config.yaml 中的 experiment.output_dir 是原始训练时的输出路径。这里保留这些路径是为了可复现性，不要求它们与当前发布目录一致。

Hugging Face 上传说明

该目录适合作为模型权重发布包上传。原始 SEG-Y 数据集体积更大，结构也不同，建议单独上传：

segy_with_masks/
  data/
  label/

上传当前权重发布包时，应上传整个 first_break_picking_hf_release/ 目录，以保证每个 best.pt 都和对应的 config.yaml 保持在一起。