FusionLayer

class torchtune.modules.model_fusion.FusionLayer(layer: Module, fusion_layer: Module, fusion_first: bool = True)

如 Flamingo：用于少样本学习的视觉语言模型 中所述的融合层。

深度融合模型架构通过将编码器输出注入到 LLM 的中间层来组合预训练的编码器模型和预训练的语言模型。这使得语言模型能够将编码器输出解释为文本并“理解”您可以训练编码器的任何模态。为了使语言模型能够适应编码器输出，FusionLayer 将一个新的可学习层融合到现有的解码器（语言模型）层中。此附加层可以获取编码器嵌入并学习将其与解码器中的标记嵌入相结合。该模块支持在原始层之前或之后融合新层，在 Flamingo 中，新层融合在原始层之前。

原始层被包装在 FusionLayer 中，以便它保留其原始的 state_dict 键，并且预训练的检查点不会被破坏。新的层参数可以通过 fusion_params 获得，以便单独控制它们是否可训练。

示例

>>> # Original decoder style transformer
>>> layer = nn.TransformerSelfAttentionLayer(...)
>>> model = TransformerDecoder(layers=layer, num_layers=32, ...)
>>>
>>> # Fuse a cross attention layer to each self attention layer to adapt for the encoder
>>> fusion_layer = nn.TransformerCrossAttentionLayer(...)
>>> fused_layer = FusionLayer(layer, fusion_layer)
>>> model = TransformerDecoder(layers=fused_layer, num_layers=32, ...)
>>>
>>> # Original decoder state_dict still works
>>> model.load_state_dict(..., strict=False)

参数：

• layer (nn.Module) – 原始解码器层

• fusion_layer (nn.Module) – 新的融合层

• fusion_first (bool) – 布尔值，用于在解码器层之前或之后插入融合层。

property cache_enabled: bool

检查键值缓存是否已设置。

forward(x: Tensor, **kwargs: Dict) → Tensor

参数：

• x (torch.Tensor) – 输入张量，形状为 [batch_size x seq_length x embed_dim]

• **kwargs (Dict) – 所有其他层参数

返回值：

与输入形状相同的输出张量

[batch_size x seq_length x embed_dim]`

返回类型：

张量

fusion_params() → List[str]

返回融合层的参数。

reset_cache()

重置两个层的键值缓存。

setup_cache(batch_size: int, dtype: dtype, *, encoder_max_seq_len: int, decoder_max_seq_len: int) → None

为两个层设置键值缓存。

参数：

• batch_size (int) – 缓存的批次大小。

• dtype (torch.dpython:type) – 缓存的数据类型。

• encoder_max_seq_len (int) – 交叉注意力层的最大缓存序列长度。

• decoder_max_seq_len (int) – 自注意力层的最大缓存序列长度。