TransformerCrossAttentionLayer

class torchtune.modules.TransformerCrossAttentionLayer(attn: MultiHeadAttention, mlp: Module, *, ca_norm: Optional[Module] = None, mlp_norm: Optional[Module] = None, ca_scale: Optional[Module] = None, mlp_scale: Optional[Module] = None)

遵循与 TransformerSelfAttentionLayer 相同约定的交叉注意力 Transformer 层。归一化应用于注意力层和 FF 层之前。

参数:

• attn (MultiHeadAttention) – 注意力模块。

• mlp (nn.Module) – 前馈模块。

• ca_norm (Optional[nn.Module]) – 在交叉注意力之前应用的归一化。

• mlp_norm (Optional[nn.Module]) – 在前馈层之前应用的归一化。

• ca_scale (Optional[nn.Module]) – 用于缩放交叉注意力输出的模块。

• mlp_scale (Optional[nn.Module]) – 用于缩放前馈输出的模块。

引发:

AssertionError – 如果设置了 attn.pos_embeddings。

caches_are_enabled() → bool

检查 `self.attn` 上的键值缓存是否已启用。参见 :func:~torchtune.modules.TransformerDecoder.caches_are_enabled`。

caches_are_setup() → bool

检查 `self.attn` 上的键值缓存是否已设置。参见 :func:~torchtune.modules.TransformerDecoder.caches_are_setup`。

forward(x: Tensor, *, encoder_input: Optional[Tensor] = None, encoder_mask: Optional[Tensor] = None, **kwargs: Dict) → Tensor

参数:

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

• encoder_input (Optional[torch.Tensor]) – 可选的来自编码器的输入嵌入。形状为 [batch_size x token_sequence x embed_dim]

• encoder_mask (Optional[torch.Tensor]) – 布尔张量，定义了 token 和编码器嵌入之间的关系矩阵。位置 i,j 为 True 表示 token i 可以关注解码器中的嵌入 j。掩码形状为 [batch_size x token_sequence x embed_sequence]。默认为 None。

• **kwargs (Dict) – 与自注意力无关的 Transformer 层输入。

返回:

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

[batch_size x seq_length x embed_dim]

返回类型:

torch.Tensor

reset_cache()

重置键值缓存。

setup_caches(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) – 此参数在此层中被忽略。