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)

交叉注意力 Transformer 层，遵循与 TransformerSelfAttentionLayer 相同的约定。归一化应用于注意力和 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]) – 用于缩放前馈输出的模块。

Raises:

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 层输入。

Returns:

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

[batch_size x seq_length x embed_dim]

Return type:

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