r""" Focal Modulation Networks (FocalNets) 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 c
| 374 | |
| 375 | |
| 376 | class FocalNet(nn.Module): |
| 377 | r""" Focal Modulation Networks (FocalNets) |
| 378 | |
| 379 | Args: |
| 380 | img_size (int | tuple(int)): Input image size. Default 224 |
| 381 | patch_size (int | tuple(int)): Patch size. Default: 4 |
| 382 | in_chans (int): Number of input image channels. Default: 3 |
| 383 | num_classes (int): Number of classes for classification head. Default: 1000 |
| 384 | embed_dim (int): Patch embedding dimension. Default: 96 |
| 385 | depths (tuple(int)): Depth of each Focal Transformer layer. |
| 386 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 |
| 387 | drop_rate (float): Dropout rate. Default: 0 |
| 388 | drop_path_rate (float): Stochastic depth rate. Default: 0.1 |
| 389 | norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. |
| 390 | patch_norm (bool): If True, add normalization after patch embedding. Default: True |
| 391 | use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False |
| 392 | focal_levels (list): How many focal levels at all stages. Note that this excludes the finest-grain level. Default: [1, 1, 1, 1] |
| 393 | focal_windows (list): The focal window size at all stages. Default: [7, 5, 3, 1] |
| 394 | use_conv_embed (bool): Whether use convolutional embedding. We noted that using convolutional embedding usually improve the performance, but we do not use it by default. Default: False |
| 395 | use_layerscale (bool): Whether use layerscale proposed in CaiT. Default: False |
| 396 | layerscale_value (float): Value for layer scale. Default: 1e-4 |
| 397 | use_postln (bool): Whether use layernorm after modulation (it helps stablize training of large models) |
| 398 | """ |
| 399 | |
| 400 | def __init__(self, |
| 401 | img_size=32, |
| 402 | patch_size=2, |
| 403 | in_chans=3, |
| 404 | num_classes=10, |
| 405 | embed_dim=96, |
| 406 | depths=[2, 2, 6, 2], |
| 407 | mlp_ratio=4., |
| 408 | drop_rate=0., |
| 409 | drop_path_rate=0.1, |
| 410 | norm_layer=nn.LayerNorm, |
| 411 | patch_norm=True, |
| 412 | use_checkpoint=False, |
| 413 | focal_levels=[2, 2, 2, 2], |
| 414 | focal_windows=[3, 3, 3, 3], |
| 415 | use_conv_embed=False, |
| 416 | use_layerscale=False, |
| 417 | layerscale_value=1e-4, |
| 418 | use_postln=False, |
| 419 | use_postln_in_modulation=False, |
| 420 | normalize_modulator=False, |
| 421 | **kwargs): |
| 422 | super().__init__() |
| 423 | |
| 424 | self.num_layers = len(depths) |
| 425 | embed_dim = [embed_dim * (2 ** i) for i in range(self.num_layers)] |
| 426 | |
| 427 | self.num_classes = num_classes |
| 428 | self.embed_dim = embed_dim |
| 429 | self.patch_norm = patch_norm |
| 430 | self.num_features = embed_dim[-1] |
| 431 | self.mlp_ratio = mlp_ratio |
| 432 | |
| 433 | # split image into patches using either non-overlapped embedding or overlapped embedding |