r""" Swin MLP Args: img_size (int | tuple(int)): Input image size. Default 224 patch_size (int | tuple(int)): Patch size. Default: 4 in_chans (int): Number of input image channels. Default: 3 num_classes (int): Number of classes for classification head. Default:
| 346 | |
| 347 | |
| 348 | class SwinMLP(nn.Module): |
| 349 | r""" Swin MLP |
| 350 | |
| 351 | Args: |
| 352 | img_size (int | tuple(int)): Input image size. Default 224 |
| 353 | patch_size (int | tuple(int)): Patch size. Default: 4 |
| 354 | in_chans (int): Number of input image channels. Default: 3 |
| 355 | num_classes (int): Number of classes for classification head. Default: 1000 |
| 356 | embed_dim (int): Patch embedding dimension. Default: 96 |
| 357 | depths (tuple(int)): Depth of each Swin MLP layer. |
| 358 | num_heads (tuple(int)): Number of attention heads in different layers. |
| 359 | window_size (int): Window size. Default: 7 |
| 360 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 |
| 361 | drop_rate (float): Dropout rate. Default: 0 |
| 362 | drop_path_rate (float): Stochastic depth rate. Default: 0.1 |
| 363 | norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. |
| 364 | ape (bool): If True, add absolute position embedding to the patch embedding. Default: False |
| 365 | patch_norm (bool): If True, add normalization after patch embedding. Default: True |
| 366 | use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False |
| 367 | """ |
| 368 | |
| 369 | def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, |
| 370 | embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], |
| 371 | window_size=7, mlp_ratio=4., drop_rate=0., drop_path_rate=0.1, |
| 372 | norm_layer=nn.LayerNorm, ape=False, patch_norm=True, |
| 373 | use_checkpoint=False, **kwargs): |
| 374 | super().__init__() |
| 375 | |
| 376 | self.num_classes = num_classes |
| 377 | self.num_layers = len(depths) |
| 378 | self.embed_dim = embed_dim |
| 379 | self.ape = ape |
| 380 | self.patch_norm = patch_norm |
| 381 | self.num_features = int(embed_dim * 2 ** (self.num_layers - 1)) |
| 382 | self.mlp_ratio = mlp_ratio |
| 383 | |
| 384 | # split image into non-overlapping patches |
| 385 | self.patch_embed = PatchEmbed( |
| 386 | img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, |
| 387 | norm_layer=norm_layer if self.patch_norm else None) |
| 388 | num_patches = self.patch_embed.num_patches |
| 389 | patches_resolution = self.patch_embed.patches_resolution |
| 390 | self.patches_resolution = patches_resolution |
| 391 | |
| 392 | # absolute position embedding |
| 393 | if self.ape: |
| 394 | self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) |
| 395 | trunc_normal_(self.absolute_pos_embed, std=.02) |
| 396 | |
| 397 | self.pos_drop = nn.Dropout(p=drop_rate) |
| 398 | |
| 399 | # stochastic depth |
| 400 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule |
| 401 | |
| 402 | # build layers |
| 403 | self.layers = nn.ModuleList() |
| 404 | for i_layer in range(self.num_layers): |
| 405 | layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), |