(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.)
| 195 | |
| 196 | class Attention(nn.Module): |
| 197 | def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.): |
| 198 | super().__init__() |
| 199 | self.num_heads = num_heads |
| 200 | head_dim = dim // num_heads |
| 201 | self.scale = qk_scale or head_dim ** -0.5 |
| 202 | |
| 203 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| 204 | self.attn_drop = nn.Dropout(attn_drop) |
| 205 | self.proj = nn.Linear(dim, dim) |
| 206 | self.proj_drop = nn.Dropout(proj_drop) |
| 207 | |
| 208 | def forward(self, x): |
| 209 | B, N, C = x.shape |