如何在Hydra中实现神经网络层配置名称的可定制化？

问题背景 #

使用Hydra构建神经网络层配置时，因不同配置仅宽度参数（如C_width、hidden_dims等）存在差异，导致需要创建大量重复配置。希望实现配置名称的参数化，例如将固定名称diffusion128|gcn128|concurrent改为diffusion${width1}|gcn${width2}|concurrent，其中width1和width2可在构建架构时动态定义。

现有配置示例 #

diffusion128|gcn128|concurent:
  name: ProteinEncoder
  instanciate:
    _target_: atomsurf.networks.ProteinEncoderBlock
  kwargs:
    surface_encoder:
      name: DiffusionNetBlockBatch
      instanciate:
        _target_: atomsurf.network_utils.DiffusionNetBlockBatch  # diffusion_net.DiffusionNet
      kwargs:
        C_width: 128
        mlp_hidden_dims: [128, 128]
        dropout: 0.0
        use_bn: true
        init_time: 2.0  # either null (for constant init) or a float
        init_std: 2.0
    graph_encoder:
      name: GCNx2Block
      instanciate:
        _target_: atomsurf.network_utils.GCNx2Block
      kwargs:
        dim_in: 128
        hidden_dims: 128
        dim_out: 128
        dropout: 0.0
        use_bn: true
        use_weighted_edge_distance: false
    communication_block:
      name: ConcurrentCommunication
# [...]

实现方案 #

方法1：Hydra Config Groups + OmegaConf插值 #

这是贴合Hydra原生工作流的方案，通过配置组管理参数，结合OmegaConf字符串插值实现动态命名与参数同步：

1. 创建参数化基础配置（如configs/layers/protein_encoder.yaml）：

defaults:
  - _self_
  - override /params: default_params

# 动态生成配置名称
${params.name_prefix}${params.width1}|gcn${params.width2}|concurrent:
  name: ProteinEncoder
  instanciate:
    _target_: atomsurf.networks.ProteinEncoderBlock
  kwargs:
    surface_encoder:
      name: DiffusionNetBlockBatch
      instanciate:
        _target_: atomsurf.network_utils.DiffusionNetBlockBatch
      kwargs:
        C_width: ${params.width1}
        mlp_hidden_dims: [${params.width1}, ${params.width1}]
        dropout: 0.0
        use_bn: true
        init_time: 2.0
        init_std: 2.0
    graph_encoder:
      name: GCNx2Block
      instanciate:
        _target_: atomsurf.network_utils.GCNx2Block
      kwargs:
        dim_in: ${params.width2}
        hidden_dims: ${params.width2}
        dim_out: ${params.width2}
        dropout: 0.0
        use_bn: true
        use_weighted_edge_distance: false
    communication_block:
      name: ConcurrentCommunication
# [...]

2. 创建参数配置组（configs/params/default_params.yaml）：

width1: 128
width2: 128
name_prefix: diffusion

3. 运行时动态指定参数：
   通过命令行直接覆盖参数，生成对应名称的配置：

python your_script.py params.width1=256 params.width2=256

方法2：OmegaConf自定义插值函数 #

如果需要更灵活的命名逻辑，可自定义OmegaConf插值函数：

1. 注册自定义插值函数：
   在Python代码中注册生成配置名称的函数：

from omegaconf import OmegaConf

def generate_layer_name(width1: int, width2: int) -> str:
    return f"diffusion{width1}|gcn{width2}|concurrent"

OmegaConf.register_new_resolver("layer_name", generate_layer_name)

2. 在配置文件中使用自定义插值：

${layer_name:${width1},${width2}}:
  name: ProteinEncoder
  instanciate:
    _target_: atomsurf.networks.ProteinEncoderBlock
  kwargs:
    surface_encoder:
      kwargs:
        C_width: ${width1}
        mlp_hidden_dims: [${width1}, ${width1}]
    graph_encoder:
      kwargs:
        dim_in: ${width2}
        hidden_dims: ${width2}
        dim_out: ${width2}
# [...]

3. 运行时传入参数：

python your_script.py width1=128 width2=64

方法3：Hydra Composition API动态构建配置 #

直接在Python代码中动态生成配置结构，彻底避免重复YAML文件：

from hydra import compose, initialize_config_dir
from omegaconf import OmegaConf

def build_protein_encoder_config(width1: int, width2: int):
    layer_name = f"diffusion{width1}|gcn{width2}|concurrent"
    return OmegaConf.create({
        layer_name: {
            "name": "ProteinEncoder",
            "instanciate": {
                "_target_": "atomsurf.networks.ProteinEncoderBlock"
            },
            "kwargs": {
                "surface_encoder": {
                    "name": "DiffusionNetBlockBatch",
                    "instanciate": {
                        "_target_": "atomsurf.network_utils.DiffusionNetBlockBatch"
                    },
                    "kwargs": {
                        "C_width": width1,
                        "mlp_hidden_dims": [width1, width1],
                        "dropout": 0.0,
                        "use_bn": True,
                        "init_time": 2.0,
                        "init_std": 2.0
                    }
                },
                "graph_encoder": {
                    "name": "GCNx2Block",
                    "instanciate": {
                        "_target_": "atomsurf.network_utils.GCNx2Block"
                    },
                    "kwargs": {
                        "dim_in": width2,
                        "hidden_dims": width2,
                        "dim_out": width2,
                        "dropout": 0.0,
                        "use_bn": True,
                        "use_weighted_edge_distance": False
                    }
                },
                "communication_block": {
                    "name": "ConcurrentCommunication"
                }
            }
        }
    })

# 使用示例
with initialize_config_dir(config_dir="configs"):
    base_config = compose(config_name="base")
    encoder_config = build_protein_encoder_config(128, 64)
    final_config = OmegaConf.merge(base_config, encoder_config)

内容的提问来源于stack exchange，提问作者schlodinger