pdebench-swe-rdb_model-M.yaml

model_type: pdeformer
model:
  graphormer:
    num_node_type: 128
    num_in_degree: 32
    num_out_degree: 32
    num_spatial: 16
    num_encoder_layers: 12
    embed_dim: 768
    ffn_embed_dim: 1536
    num_heads: 32
    pre_layernorm: True
  scalar_encoder:
    dim_hidden: 256
    num_layers: 3
  function_encoder:
    type: cnn2dv3
    num_branches: 4
    resolution: 128
    conv2d_input_txyz: False
    cnn_keep_nchw: True
  multi_inr:
    enable: False
  inr:
    type: poly_inr
    num_layers: 12
    dim_hidden: 256
    poly_inr:
      enable_affine: False
      enable_shift: True
      enable_scale: True
      modify_he_init: False
      affine_act_fn: identity  # {identity, lrelu, sin}
      activation_fn: sin  # {lrelu, sin}
  hypernet:
    dim_hidden: 512
    num_layers: 2
    shared: False  # whether the parameters of all INR layers are generated by the same hypernet
  load_ckpt: path/to/your/downloaded/model-M.ckpt
  # You can download from https://ai.gitee.com/functoreality/PDEformer2-M/blob/master/model-M.ckpt
data:
  path: ../data_download  # or any path/to/your/data_download
  type: single_pde
  num_workers: 8
  num_samples_per_file:
    train: 1
    test: 100
  pde_dag:
    max_n_scalar_nodes: 80
    max_n_function_nodes: 4
    disconn_attn_bias: -inf
  single_pde:
    param_name: rdb
    regularize_ratio: 0.
    train: [1]
train:
  total_batch_size: 1
  num_txyz_samp_pts: 8192
  lr_init: 5.e-6
  epochs: 600
  loss:
    type: RMSE
    normalize: True
    normalize_eps: 0.05
  optimizer: Adam  # {Adam, AdamW}
  weight_decay: 0.0
  lr_scheduler:
    type: mstep
    milestones: [1]
    decay: 0.9
    enable_warmup: False
  grad_clip_value: 1  # -1 means no gradient clipping
eval:
  total_batch_size: 4
  interval: 60
  plot_num_per_type: 1
  dataset_per_type: 1
record_dir: "exp/benchmark_data/pdebench-rdb/train-1/model-M"