Build model depending on its size Args: setting (str): model size (small/base/large) image_size (int, int): ihabe height and width patch_size (int): patch size for the vit Returns: model (BinModel): the built model to be trained
(setting , image_size, patch_size)
| 94 | return rec_images, patches |
| 95 | |
| 96 | def build_model(setting , image_size, patch_size): |
| 97 | """ |
| 98 | Build model depending on its size |
| 99 | |
| 100 | Args: |
| 101 | setting (str): model size (small/base/large) |
| 102 | image_size (int, int): ihabe height and width |
| 103 | patch_size (int): patch size for the vit |
| 104 | Returns: |
| 105 | model (BinModel): the built model to be trained |
| 106 | """ |
| 107 | # define hyperparameters for the models depending on size |
| 108 | hyper_params = {"base": [6, 8, 768], |
| 109 | "small": [3, 4, 512], |
| 110 | "large": [12, 16, 1024]} |
| 111 | |
| 112 | encoder_layers = hyper_params[setting][0] |
| 113 | encoder_heads = hyper_params[setting][1] |
| 114 | encoder_dim = hyper_params[setting][2] |
| 115 | |
| 116 | # define encoder |
| 117 | v = ViT( |
| 118 | image_size = image_size, |
| 119 | patch_size = patch_size, |
| 120 | num_classes = 1000, |
| 121 | dim = encoder_dim, |
| 122 | depth = encoder_layers, |
| 123 | heads = encoder_heads, |
| 124 | mlp_dim = 2048 |
| 125 | ) |
| 126 | |
| 127 | # define full model |
| 128 | model = BinModel( |
| 129 | encoder = v, |
| 130 | decoder_dim = encoder_dim, |
| 131 | decoder_depth = encoder_layers, |
| 132 | decoder_heads = encoder_heads |
| 133 | ) |
| 134 | return model |
| 135 | |
| 136 | class multitask_ViT_model(BaseModel): |
| 137 | def __init__(self, opt): |