FeedForward module Args: dim (int): Input dimension. hidden_dim (int): Hidden dimension of the feedforward layer. multiple_of (int): Value to ensure hidden dimension is a multiple of this value. ffn_dim_multiplier (Optional[float]): Custom multiplier for hid
| 229 | |
| 230 | |
| 231 | class FeedForward(nn.Module): |
| 232 | """ |
| 233 | FeedForward module |
| 234 | |
| 235 | Args: |
| 236 | dim (int): Input dimension. |
| 237 | hidden_dim (int): Hidden dimension of the feedforward layer. |
| 238 | multiple_of (int): Value to ensure hidden dimension is a multiple of this value. |
| 239 | ffn_dim_multiplier (Optional[float]): Custom multiplier for hidden dimension. Defaults to None. |
| 240 | |
| 241 | Attributes: |
| 242 | w1 (Linear): Linear transformation for the first layer. |
| 243 | w2 (Linear): Linear transformation for the second layer. |
| 244 | w3 (Linear): Linear transformation for the third layer. |
| 245 | |
| 246 | """ |
| 247 | |
| 248 | def __init__( |
| 249 | self, |
| 250 | dim: int, |
| 251 | hidden_dim: int, |
| 252 | multiple_of: int, |
| 253 | ffn_dim_multiplier: Optional[float], |
| 254 | ): |
| 255 | super().__init__() |
| 256 | hidden_dim = int(2 * hidden_dim / 3) |
| 257 | # custom dim factor multiplier |
| 258 | if ffn_dim_multiplier is not None: |
| 259 | hidden_dim = int(ffn_dim_multiplier * hidden_dim) |
| 260 | hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) |
| 261 | |
| 262 | self.w1 = nn.Linear(dim, hidden_dim, bias=False) |
| 263 | self.w2 = nn.Linear(hidden_dim, dim, bias=False) |
| 264 | self.w3 = nn.Linear(dim, hidden_dim, bias=False) |
| 265 | |
| 266 | def forward(self, x): |
| 267 | return self.w2(F.silu(self.w1(x)) * self.w3(x)) |
| 268 | |
| 269 | def init_weights(self, init_std: float): |
| 270 | nn.init.trunc_normal_(self.w1.weight, mean=0.0, std=0.02) |
| 271 | for linear in (self.w2, self.w3): |
| 272 | nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std) |
| 273 | |
| 274 | |
| 275 | class TransformerBlock(nn.Module): |