| 168 | |
| 169 | |
| 170 | class Block(nn.Module): |
| 171 | |
| 172 | def __init__( |
| 173 | self, |
| 174 | dim, |
| 175 | num_heads, |
| 176 | mlp_ratio=4.0, |
| 177 | qkv_bias=False, |
| 178 | qk_scale=None, |
| 179 | drop=0.0, |
| 180 | attn_drop=0.0, |
| 181 | drop_path=0.0, |
| 182 | act_layer=nn.GELU, |
| 183 | norm_layer=nn.LayerNorm, |
| 184 | ): |
| 185 | super().__init__() |
| 186 | self.norm1 = norm_layer(dim) |
| 187 | self.attn = Attention( |
| 188 | dim, |
| 189 | num_heads=num_heads, |
| 190 | qkv_bias=qkv_bias, |
| 191 | qk_scale=qk_scale, |
| 192 | attn_drop=attn_drop, |
| 193 | proj_drop=drop, |
| 194 | ) |
| 195 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
| 196 | self.drop_path = DropPath( |
| 197 | drop_path) if drop_path > 0.0 else nn.Identity() |
| 198 | self.norm2 = norm_layer(dim) |
| 199 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 200 | self.mlp = Mlp(in_features=dim, |
| 201 | hidden_features=mlp_hidden_dim, |
| 202 | act_layer=act_layer, |
| 203 | drop=drop) |
| 204 | |
| 205 | def forward(self, x): |
| 206 | x = x + self.drop_path(self.attn(self.norm1(x))) |
| 207 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 208 | return x |
| 209 | |
| 210 | |
| 211 | class PatchEmbed(nn.Module): |