前馈网络

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