CNN-ViT/config.py at main · kangchengX/CNN-ViT · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
from typing import Literal


def get_mobilevit_config(
    config_arch: Literal['mobilevit_xxs', 'mobilevit_xs', 'mobilevit_s'],
    num_classes: int ,
    image_size: int | None = 224,
    dropout: float | None = 0.5
):
    """
    Get config for MobileVit.

    Args:
        config_arch: architecture of the model. Must be `'mobilevit_xxs'`, `'mobilevit_xs'`or `'mobilevit_s'`.
        num_classes (int): number of the classes.
        image_size (int): image_size = image_height = image_width. Default to `224`.
        dropout (float): dropout percentage. Default to `0.5`.
    """
    if config_arch == 'mobilevit_xxs':
        config = {
            'channels' : [16, 16, 24, 24, 24, 48, 64, 80],
            'dims' : [64, 80, 96],
            'mlp_dims' : [128, 160, 192],
            'num_classes' : num_classes,
            'image_size' : image_size,
            'expansion_factor' : 2,
            'patch_size' : 2,
            'dropout' : dropout,
            'last_conv_expansion_factor' : 4
        }

    elif config_arch == 'mobilevit_xs':
        config = {
            'channels' : [16, 32, 48, 48, 48, 64, 80, 96],
            'dims' : [96, 120, 144],
            'mlp_dims' : [192, 240, 288],
            'num_classes' : num_classes,
            'image_size' : image_size,
            'expansion_factor' : 4,
            'patch_size' : 2 ,
            'dropout' : dropout,
            'last_conv_expansion_factor' : 4
        }

    elif config_arch == 'mobilevit_s':
        config = {
            'channels' : [16, 32, 64, 64, 64, 96, 128, 160],
            'dims' : [144, 192, 240],
            'mlp_dims' : [288, 384, 480],
            'num_classes' : num_classes,
            'image_size' : image_size,
            'expansion_factor' : 4,
            'patch_size' : 2,
            'dropout' : dropout,
            'last_conv_expansion_factor' : 4
        }

    return config


def get_resnet_convblock_config(config_arch: Literal['resnet18', 'resnet34']):
    if config_arch == 'resnet18':
        config = [(64, 2), (128, 2), (256, 2), (512, 2)]
    elif config_arch == 'resnet34':
        config = [(64, 3), (128, 4), (256, 6), (512, 3)]
    else:
        raise ValueError(f"Unknown config arch {config_arch}")

    return config