FeedForward

class torchtune.modules.FeedForward(*, gate_proj: Module, down_proj: Module, up_proj: Optional[Module] = None, activation: Module = SiLU())

此类实现了从 Llama2 派生的前馈网络。

参数:

• gate_proj (nn.Module) – 从输入维度到隐藏维度的投影，通过激活函数并与 up_proj 相乘。

• down_proj (nn.Module) – 到输出维度的最终投影。

• up_proj (Optional[nn.Module]) – 从输入维度到隐藏维度的投影，与 activation(gate_proj) 相乘。

• activation (nn.Module) – 要使用的激活函数。默认为 nn.SiLU()。

forward(x: Tensor) → Tensor

参数:

x (torch.Tensor) – 输入张量，形状为 (..., in_dim)，其中 in_dim 是 gate_proj 和 up_proj 的输入维度。

返回:

输出张量，形状为 (..., out_dim)，其中 out_dim 是 down_proj 的输出维度。

返回类型:

torch.Tensor