Inspired by torch.nn.TransformerEncoderLayer and rwightman's timm package.
| 118 | |
| 119 | |
| 120 | class TransformerEncoderLayer(nn.Module): |
| 121 | """ |
| 122 | Inspired by torch.nn.TransformerEncoderLayer and |
| 123 | rwightman's timm package. |
| 124 | """ |
| 125 | |
| 126 | def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, |
| 127 | attention_dropout=0.1, drop_path_rate=0.1): |
| 128 | super().__init__() |
| 129 | |
| 130 | self.pre_norm = nn.LayerNorm(d_model) |
| 131 | self.self_attn = Attention(dim=d_model, num_heads=nhead, |
| 132 | attention_dropout=attention_dropout, projection_dropout=dropout) |
| 133 | |
| 134 | self.linear1 = nn.Linear(d_model, dim_feedforward) |
| 135 | self.dropout1 = nn.Dropout(dropout) |
| 136 | self.norm1 = nn.LayerNorm(d_model) |
| 137 | self.linear2 = nn.Linear(dim_feedforward, d_model) |
| 138 | self.dropout2 = nn.Dropout(dropout) |
| 139 | |
| 140 | self.drop_path = DropPath(drop_path_rate) |
| 141 | |
| 142 | self.activation = F.gelu |
| 143 | |
| 144 | def forward(self, src, *args, **kwargs): |
| 145 | src = src + self.drop_path(self.self_attn(self.pre_norm(src))) |
| 146 | src = self.norm1(src) |
| 147 | src2 = self.linear2(self.dropout1(self.activation(self.linear1(src)))) |
| 148 | src = src + self.drop_path(self.dropout2(src2)) |
| 149 | return src |
| 150 | |
| 151 | |
| 152 | class DropPath(nn.Module): |