| 460 | return x #.reshape(B, C, H, W) |
| 461 | |
| 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) |
| 493 | # x = x + self.pos_embed(x) |
| 494 | x = x + self.drop_path(self.token_mixer(self.norm1(x), relative_pos_enc)) |
| 495 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 496 | return x |
| 497 | |
| 498 | def forward(self, x, relative_pos_enc=None): |
| 499 | if self.grad_checkpoint and x.requires_grad: |
| 500 | x = checkpoint.checkpoint(self._forward_impl, x, relative_pos_enc) |
| 501 | else: |
| 502 | x = self._forward_impl(x, relative_pos_enc) |
| 503 | return x |
| 504 | |
| 505 | class RetBlock(nn.Module): |
| 506 | |