(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.1, drop_path=0.0)
| 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): |
| 102 | super().__init__() |
| 103 | self.norm1 = nn.LayerNorm(embed_dim) |
| 104 | self.attn = MultiHeadAttention(embed_dim, num_heads, dropout) |
| 105 | self.norm2 = nn.LayerNorm(embed_dim) |
| 106 | self.mlp = FeedForward(embed_dim, int(embed_dim * mlp_ratio), dropout) |
| 107 | |
| 108 | from .modules import DropPath |
| 109 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 110 | |
| 111 | def forward(self, x, mask=None): |
| 112 | x = x + self.drop_path(self.attn(self.norm1(x), mask)) |
nothing calls this directly
no test coverage detected