| 130 | |
| 131 | |
| 132 | class FFNLayer(nn.Module): |
| 133 | |
| 134 | def __init__(self, d_model, dim_feedforward=2048, dropout=0.0, |
| 135 | activation="relu", normalize_before=False): |
| 136 | super().__init__() |
| 137 | # Implementation of Feedforward model |
| 138 | self.linear1 = nn.Linear(d_model, dim_feedforward) |
| 139 | self.dropout = nn.Dropout(dropout) |
| 140 | self.linear2 = nn.Linear(dim_feedforward, d_model) |
| 141 | |
| 142 | self.norm = nn.LayerNorm(d_model) |
| 143 | |
| 144 | self.activation = _get_activation_fn(activation) |
| 145 | self.normalize_before = normalize_before |
| 146 | |
| 147 | self._reset_parameters() |
| 148 | |
| 149 | def _reset_parameters(self): |
| 150 | for p in self.parameters(): |
| 151 | if p.dim() > 1: |
| 152 | nn.init.xavier_uniform_(p) |
| 153 | |
| 154 | def with_pos_embed(self, tensor, pos: Optional[Tensor]): |
| 155 | return tensor if pos is None else tensor + pos |
| 156 | |
| 157 | def forward_post(self, tgt): |
| 158 | tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) |
| 159 | tgt = tgt + self.dropout(tgt2) |
| 160 | tgt = self.norm(tgt) |
| 161 | return tgt |
| 162 | |
| 163 | def forward_pre(self, tgt): |
| 164 | tgt2 = self.norm(tgt) |
| 165 | tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) |
| 166 | tgt = tgt + self.dropout(tgt2) |
| 167 | return tgt |
| 168 | |
| 169 | def forward(self, tgt): |
| 170 | if self.normalize_before: |
| 171 | return self.forward_pre(tgt) |
| 172 | return self.forward_post(tgt) |
| 173 | |
| 174 | |
| 175 | def _get_activation_fn(activation): |
nothing calls this directly
no outgoing calls
no test coverage detected