Efficient Multi-Head Self Attention
| 208 | |
| 209 | |
| 210 | class E_MHSA(nn.Module): |
| 211 | """ |
| 212 | Efficient Multi-Head Self Attention |
| 213 | """ |
| 214 | def __init__(self, dim, out_dim=None, head_dim=32, qkv_bias=True, qk_scale=None, |
| 215 | attn_drop=0, proj_drop=0., sr_ratio=1): |
| 216 | super().__init__() |
| 217 | self.dim = dim |
| 218 | self.out_dim = out_dim if out_dim is not None else dim |
| 219 | self.num_heads = self.dim // head_dim |
| 220 | self.scale = qk_scale or head_dim ** -0.5 |
| 221 | self.q = nn.Linear(dim, self.dim, bias=qkv_bias) |
| 222 | self.k = nn.Linear(dim, self.dim, bias=qkv_bias) |
| 223 | self.v = nn.Linear(dim, self.dim, bias=qkv_bias) |
| 224 | self.proj = nn.Linear(self.dim, self.out_dim) |
| 225 | self.attn_drop = nn.Dropout(attn_drop) |
| 226 | self.proj_drop = nn.Dropout(proj_drop) |
| 227 | |
| 228 | self.sr_ratio = sr_ratio |
| 229 | self.N_ratio = sr_ratio ** 2 |
| 230 | if sr_ratio > 1: |
| 231 | self.sr = nn.AvgPool1d(kernel_size=self.N_ratio, stride=self.N_ratio) |
| 232 | self.norm = nn.BatchNorm1d(dim, eps=NORM_EPS) |
| 233 | self.is_bn_merged = False |
| 234 | |
| 235 | def merge_bn(self, pre_bn): |
| 236 | merge_pre_bn(self.q, pre_bn) |
| 237 | if self.sr_ratio > 1: |
| 238 | merge_pre_bn(self.k, pre_bn, self.norm) |
| 239 | merge_pre_bn(self.v, pre_bn, self.norm) |
| 240 | else: |
| 241 | merge_pre_bn(self.k, pre_bn) |
| 242 | merge_pre_bn(self.v, pre_bn) |
| 243 | self.is_bn_merged = True |
| 244 | |
| 245 | def forward(self, x): |
| 246 | B, N, C = x.shape |
| 247 | q = self.q(x) |
| 248 | q = q.reshape(B, N, self.num_heads, int(C // self.num_heads)).permute(0, 2, 1, 3) |
| 249 | |
| 250 | if self.sr_ratio > 1: |
| 251 | x_ = x.transpose(1, 2) |
| 252 | x_ = self.sr(x_) |
| 253 | if not torch.onnx.is_in_onnx_export() and not self.is_bn_merged: |
| 254 | x_ = self.norm(x_) |
| 255 | x_ = x_.transpose(1, 2) |
| 256 | k = self.k(x_) |
| 257 | k = k.reshape(B, -1, self.num_heads, int(C // self.num_heads)).permute(0, 2, 3, 1) |
| 258 | v = self.v(x_) |
| 259 | v = v.reshape(B, -1, self.num_heads, int(C // self.num_heads)).permute(0, 2, 1, 3) |
| 260 | else: |
| 261 | k = self.k(x) |
| 262 | k = k.reshape(B, -1, self.num_heads, int(C // self.num_heads)).permute(0, 2, 3, 1) |
| 263 | v = self.v(x) |
| 264 | v = v.reshape(B, -1, self.num_heads, int(C // self.num_heads)).permute(0, 2, 1, 3) |
| 265 | attn = (q @ k) * self.scale |
| 266 | |
| 267 | attn = attn.softmax(dim=-1) |