(self,
dim,
num_heads,
mlp_ratio=4.,
qkv_bias=False,
qk_scale=None,
drop=0.,
attn_drop=0.,
drop_path_rate=0.,
act_layer=nn.GELU,
norm_layer='nn.LayerNorm',
epsilon=1e-6,
prenorm=False)
| 357 | |
| 358 | class MHSA_Block(nn.Module): |
| 359 | def __init__(self, |
| 360 | dim, |
| 361 | num_heads, |
| 362 | mlp_ratio=4., |
| 363 | qkv_bias=False, |
| 364 | qk_scale=None, |
| 365 | drop=0., |
| 366 | attn_drop=0., |
| 367 | drop_path_rate=0., |
| 368 | act_layer=nn.GELU, |
| 369 | norm_layer='nn.LayerNorm', |
| 370 | epsilon=1e-6, |
| 371 | prenorm=False): |
| 372 | super().__init__() |
| 373 | if isinstance(norm_layer, str): |
| 374 | self.norm1 = eval(norm_layer)(dim, eps=epsilon) |
| 375 | else: |
| 376 | self.norm1 = norm_layer(dim) |
| 377 | |
| 378 | self.mixer = Attention( |
| 379 | dim, |
| 380 | num_heads=num_heads, |
| 381 | qkv_bias=qkv_bias, |
| 382 | qk_scale=qk_scale, |
| 383 | attn_drop=attn_drop, |
| 384 | proj_drop=drop) |
| 385 | |
| 386 | # self.drop_path = DropPath(local_rank,drop_path) if drop_path > 0. else Identity() |
| 387 | self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity() |
| 388 | if isinstance(norm_layer, str): |
| 389 | self.norm2 = eval(norm_layer)(dim, eps=epsilon) |
| 390 | else: |
| 391 | self.norm2 = norm_layer(dim) |
| 392 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 393 | self.mlp_ratio = mlp_ratio |
| 394 | self.mlp = FeedForward(in_dim=dim, hidden_dim=mlp_hidden_dim, act_layer=act_layer, dropout=drop) |
| 395 | self.prenorm = prenorm |
| 396 | |
| 397 | def forward(self, x, size=None): |
| 398 | x = x + self.drop_path(self.mixer(self.norm1(x))) |
nothing calls this directly
no test coverage detected