| 74 | return x |
| 75 | |
| 76 | class FeedForward(nn.Module): |
| 77 | def __init__(self, embed_dim, hidden_dim, dropout=0.1, activation='gelu'): |
| 78 | super().__init__() |
| 79 | self.fc1 = nn.Linear(embed_dim, hidden_dim) |
| 80 | self.fc2 = nn.Linear(hidden_dim, embed_dim) |
| 81 | self.dropout = nn.Dropout(dropout) |
| 82 | |
| 83 | if activation == 'gelu': |
| 84 | self.activation = nn.GELU() |
| 85 | elif activation == 'relu': |
| 86 | self.activation = nn.ReLU() |
| 87 | elif activation == 'silu': |
| 88 | self.activation = nn.SiLU() |
| 89 | else: |
| 90 | self.activation = nn.GELU() |
| 91 | |
| 92 | def forward(self, x): |
| 93 | x = self.fc1(x) |
| 94 | x = self.activation(x) |
| 95 | x = self.dropout(x) |
| 96 | x = self.fc2(x) |
| 97 | x = self.dropout(x) |
| 98 | return x |
| 99 | |
| 100 | class TransformerBlock(nn.Module): |
| 101 | def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.1, drop_path=0.0): |