(self, dim, num_heads=8, attention_dropout=0.1, projection_dropout=0.1)
| 89 | |
| 90 | class Attention(nn.Module): |
| 91 | def __init__(self, dim, num_heads=8, attention_dropout=0.1, projection_dropout=0.1): |
| 92 | super().__init__() |
| 93 | self.heads = num_heads |
| 94 | head_dim = dim // self.heads |
| 95 | self.scale = head_dim ** -0.5 |
| 96 | |
| 97 | self.qkv = nn.Linear(dim, dim * 3, bias=False) |
| 98 | self.attn_drop = nn.Dropout(attention_dropout) |
| 99 | self.proj = nn.Linear(dim, dim) |
| 100 | self.proj_drop = nn.Dropout(projection_dropout) |
| 101 | |
| 102 | def forward(self, x): |
| 103 | B, N, C = x.shape |