| 110 | |
| 111 | |
| 112 | class Block(nn.Module): |
| 113 | |
| 114 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., |
| 115 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): |
| 116 | super().__init__() |
| 117 | self.norm1 = norm_layer(dim) |
| 118 | self.attn = Attention( |
| 119 | dim, |
| 120 | num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, |
| 121 | attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) |
| 122 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
| 123 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 124 | self.norm2 = norm_layer(dim) |
| 125 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 126 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 127 | |
| 128 | self.apply(self._init_weights) |
| 129 | |
| 130 | def _init_weights(self, m): |
| 131 | if isinstance(m, nn.Linear): |
| 132 | trunc_normal_(m.weight, std=.02) |
| 133 | if isinstance(m, nn.Linear) and m.bias is not None: |
| 134 | nn.init.constant_(m.bias, 0) |
| 135 | elif isinstance(m, nn.LayerNorm): |
| 136 | nn.init.constant_(m.bias, 0) |
| 137 | nn.init.constant_(m.weight, 1.0) |
| 138 | elif isinstance(m, nn.Conv2d): |
| 139 | fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels |
| 140 | fan_out //= m.groups |
| 141 | m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) |
| 142 | if m.bias is not None: |
| 143 | m.bias.data.zero_() |
| 144 | |
| 145 | def forward(self, x, H, W): |
| 146 | x = x + self.drop_path(self.attn(self.norm1(x), H, W)) |
| 147 | x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) |
| 148 | |
| 149 | return x |
| 150 | |
| 151 | |
| 152 | class OverlapPatchEmbed(nn.Module): |