(self,
dim=64,
sr_ratio=1,
num_heads=1,
mlp_ratio=4,
norm_cfg=nn.LayerNorm, # dict(type='GN', num_groups=1),
act_cfg=nn.GELU, # dict(type='GELU'),
drop=0,
drop_path=0,
layer_scale_init_value=1e-5,
grad_checkpoint=False)
| 462 | class OSRA_Block(nn.Module): |
| 463 | |
| 464 | def __init__(self, |
| 465 | dim=64, |
| 466 | sr_ratio=1, |
| 467 | num_heads=1, |
| 468 | mlp_ratio=4, |
| 469 | norm_cfg=nn.LayerNorm, # dict(type='GN', num_groups=1), |
| 470 | act_cfg=nn.GELU, # dict(type='GELU'), |
| 471 | drop=0, |
| 472 | drop_path=0, |
| 473 | layer_scale_init_value=1e-5, |
| 474 | grad_checkpoint=False): |
| 475 | |
| 476 | super().__init__() |
| 477 | self.grad_checkpoint = grad_checkpoint |
| 478 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 479 | |
| 480 | # self.pos_embed = DWConv2d(dim, 3, 1, 1) |
| 481 | # self.pos_embed = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim) |
| 482 | self.norm1 = norm_cfg(dim) |
| 483 | self.token_mixer = OSRA_Attention(dim, num_heads=num_heads, |
| 484 | sr_ratio=sr_ratio) |
| 485 | self.norm2 = norm_cfg(dim) |
| 486 | |
| 487 | self.mlp = FeedForward(in_dim=dim, hidden_dim=mlp_hidden_dim, act_layer=act_cfg, dropout=drop) |
| 488 | self.drop_path = DropPath( |
| 489 | drop_path) if drop_path > 0. else nn.Identity() |
| 490 | |
| 491 | def _forward_impl(self, x, relative_pos_enc=None): |
| 492 | # print(x.shape) |
nothing calls this directly
no test coverage detected