(self, embed_dim, window_size=7, num_heads=8)
| 183 | |
| 184 | class LocalAttention(nn.Module): |
| 185 | def __init__(self, embed_dim, window_size=7, num_heads=8): |
| 186 | super().__init__() |
| 187 | self.embed_dim = embed_dim |
| 188 | self.window_size = window_size |
| 189 | self.num_heads = num_heads |
| 190 | self.head_dim = embed_dim // num_heads |
| 191 | self.scale = self.head_dim ** -0.5 |
| 192 | |
| 193 | self.qkv = nn.Linear(embed_dim, embed_dim * 3) |
| 194 | self.proj = nn.Linear(embed_dim, embed_dim) |
| 195 | |
| 196 | def forward(self, x): |
| 197 | B, N, C = x.shape |