collie.utils.dict_as_params

dict_as_params(input_keys, output_keys)

使用字典作为参数输入的辅助函数

从输入的字典中顺次取出 input_keys 作为模型的输入，并且将模型的输出以 output_keys 为 key 放入字典中作为输出。在这一过程中多余的 key 并不会被丢 弃。

可以用于 nn.LayerNorm 这些在流水线并行中一般不需要改变 forward 过程但需要 改变输入输出结构的模型，使用该函数可以避免频繁地重写这些模型，并且可以适应流水线 LayerSpec 的初始化。

dict_as_params(input_keys="input_ids", output_keys="hidden_states")(nn.Embedding, vocab_size, hidden_size)

LayerSpec(
    dict_as_params(input_keys="input_ids", output_keys="hidden_states"),
    nn.Embbedding, vocab_size, hidden_size
)

参数:

• input_keys (Union[str, Sequence[str]]) – 该模型输入需要的 key。dict_as_params 会从输入的字典中 依次取出 input_keys 的内容传入模型。

• output_keys (Union[str, Sequence[str]]) – 该模型输出对应的 key。dict_as_params 会依次将模型的 输出和 output_keys 进行对应，并放入字典中作为最终的输出。

备注

在使用该函数时，请您注意输入输出顺序和 input_keys output_keys 顺序 的对应关系，避免将错误的 key 赋给了对应的张量。