Spaces:
Sleeping
Sleeping
File size: 2,319 Bytes
6edd10d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """Position-wise feed-forward sub-layer used in transformer blocks."""
from __future__ import annotations
import torch
import torch.nn as nn
from torch import Tensor
__all__ = ["FeedForwardLayer"]
class FeedForwardLayer(nn.Module):
"""
Position-wise feed-forward layer used in Transformer blocks.
Architecture:
fc1: Linear(d_model -> d_ff)
activation: ReLU
dropout: nn.Dropout(dropout_rate)
fc2: Linear(d_ff -> d_model)
Args:
d_model (int): Dimensionality of model embeddings.
d_ff (int): Hidden dimensionality of feed-forward layer.
dropout_rate (float): Dropout probability between 0 and 1 (exclusive).
Shape:
Input: (B, S, D) where D == d_model
Output: (B, S, D)
"""
def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.1):
super().__init__()
if not isinstance(d_ff, int):
raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
if not isinstance(d_model, int):
raise TypeError(f"d_model must be an int, got {type(d_model)}")
if not isinstance(dropout_rate, float):
raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
if d_ff <= 0:
raise ValueError(f"d_ff must be strictly greater than 0, got {d_ff}")
if d_model <= 0:
raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
if not (0.0 <= dropout_rate < 1.0):
raise ValueError(f"dropout_rate must be in [0,1), got {dropout_rate}")
self.d_model = d_model
self.d_ff = d_ff
self.dropout_rate = dropout_rate
self.fc1 = nn.Linear(d_model, d_ff)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(dropout_rate)
self.fc2 = nn.Linear(d_ff, d_model)
def forward(self, x: Tensor) -> Tensor:
if not isinstance(x, torch.Tensor):
raise TypeError(f"x must be a torch.Tensor, got {type(x)}")
if x.ndim != 3:
raise ValueError(f"x must be 3D of shape (B,S,D); got shape {tuple(x.shape)}")
if x.shape[-1] != self.d_model:
raise ValueError(f"Last dim {x.shape[-1]} must match d_model {self.d_model}")
return self.fc2(self.dropout(self.relu(self.fc1(x))))
|