JanardhanM/no-reference-iqa

1#
2# For licensing see accompanying LICENSE file.
3# Copyright (C) 2023 Apple Inc. All Rights Reserved.
4#
5
6import argparse
7from typing import Dict, Optional, Tuple
8
9from torch import nn
10
11from cvnets.layers import ConvLayer2d, Dropout, GlobalPool, LinearLayer
12from cvnets.models import MODEL_REGISTRY
13from cvnets.models.classification.base_image_encoder import BaseImageEncoder
14from cvnets.models.classification.config.mobilevit import get_configuration
15from cvnets.modules import InvertedResidual, MobileViTBlock
16from utils import logger
17
18
19@MODEL_REGISTRY.register(name="mobilevit", type="classification")
20class MobileViT(BaseImageEncoder):
21    """
22    This class implements the `MobileViT architecture <https://arxiv.org/abs/2110.02178?context=cs.LG>`_
23    """
24
25    def __init__(self, opts, *args, **kwargs) -> None:
26        num_classes = getattr(opts, "model.classification.n_classes", 1000)
27        classifier_dropout = getattr(
28            opts, "model.classification.classifier_dropout", 0.0
29        )
30
31        pool_type = getattr(opts, "model.layer.global_pool", "mean")
32        image_channels = 3
33        out_channels = 16
34
35        mobilevit_config = get_configuration(opts=opts)
36
37        super().__init__(opts, *args, **kwargs)
38
39        # store model configuration in a dictionary
40        self.model_conf_dict = dict()
41        self.conv_1 = ConvLayer2d(
42            opts=opts,
43            in_channels=image_channels,
44            out_channels=out_channels,
45            kernel_size=3,
46            stride=2,
47            use_norm=True,
48            use_act=True,
49        )
50
51        self.model_conf_dict["conv1"] = {"in": image_channels, "out": out_channels}
52
53        in_channels = out_channels
54        self.layer_1, out_channels = self._make_layer(
55            opts=opts, input_channel=in_channels, cfg=mobilevit_config["layer1"]
56        )
57        self.model_conf_dict["layer1"] = {"in": in_channels, "out": out_channels}
58
59        in_channels = out_channels
60        self.layer_2, out_channels = self._make_layer(
61            opts=opts, input_channel=in_channels, cfg=mobilevit_config["layer2"]
62        )
63        self.model_conf_dict["layer2"] = {"in": in_channels, "out": out_channels}
64
65        in_channels = out_channels
66        self.layer_3, out_channels = self._make_layer(
67            opts=opts, input_channel=in_channels, cfg=mobilevit_config["layer3"]
68        )
69        self.model_conf_dict["layer3"] = {"in": in_channels, "out": out_channels}
70
71        in_channels = out_channels
72        self.layer_4, out_channels = self._make_layer(
73            opts=opts,
74            input_channel=in_channels,
75            cfg=mobilevit_config["layer4"],
76            dilate=self.dilate_l4,
77        )
78        self.model_conf_dict["layer4"] = {"in": in_channels, "out": out_channels}
79
80        in_channels = out_channels
81        self.layer_5, out_channels = self._make_layer(
82            opts=opts,
83            input_channel=in_channels,
84            cfg=mobilevit_config["layer5"],
85            dilate=self.dilate_l5,
86        )
87        self.model_conf_dict["layer5"] = {"in": in_channels, "out": out_channels}
88
89        in_channels = out_channels
90        exp_channels = min(mobilevit_config["last_layer_exp_factor"] * in_channels, 960)
91        self.conv_1x1_exp = ConvLayer2d(
92            opts=opts,
93            in_channels=in_channels,
94            out_channels=exp_channels,
95            kernel_size=1,
96            stride=1,
97            use_act=True,
98            use_norm=True,
99        )
100
101        self.model_conf_dict["exp_before_cls"] = {
102            "in": in_channels,
103            "out": exp_channels,
104        }
105
106        self.classifier = nn.Sequential()
107        self.classifier.add_module(
108            name="global_pool", module=GlobalPool(pool_type=pool_type, keep_dim=False)
109        )
110        if 0.0 < classifier_dropout < 1.0:
111            self.classifier.add_module(
112                name="dropout", module=Dropout(p=classifier_dropout, inplace=True)
113            )
114        self.classifier.add_module(
115            name="fc",
116            module=LinearLayer(
117                in_features=exp_channels, out_features=num_classes, bias=True
118            ),
119        )
120
121        # check model
122        self.check_model()
123
124        # weight initialization
125        self.reset_parameters(opts=opts)
126
127    @classmethod
128    def add_arguments(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
129        group = parser.add_argument_group(title=cls.__name__)
130        group.add_argument(
131            "--model.classification.mit.mode",
132            type=str,
133            default="small",
134            choices=["xx_small", "x_small", "small"],
135            help="MobileViT mode. Defaults to small",
136        )
137        group.add_argument(
138            "--model.classification.mit.attn-dropout",
139            type=float,
140            default=0.0,
141            help="Dropout in attention layer. Defaults to 0.0",
142        )
143        group.add_argument(
144            "--model.classification.mit.ffn-dropout",
145            type=float,
146            default=0.0,
147            help="Dropout between FFN layers. Defaults to 0.0",
148        )
149        group.add_argument(
150            "--model.classification.mit.dropout",
151            type=float,
152            default=0.0,
153            help="Dropout in Transformer layer. Defaults to 0.0",
154        )
155        group.add_argument(
156            "--model.classification.mit.transformer-norm-layer",
157            type=str,
158            default="layer_norm",
159            help="Normalization layer in transformer. Defaults to LayerNorm",
160        )
161        group.add_argument(
162            "--model.classification.mit.no-fuse-local-global-features",
163            action="store_true",
164            help="Do not combine local and global features in MobileViT block",
165        )
166        group.add_argument(
167            "--model.classification.mit.conv-kernel-size",
168            type=int,
169            default=3,
170            help="Kernel size of Conv layers in MobileViT block",
171        )
172
173        group.add_argument(
174            "--model.classification.mit.head-dim",
175            type=int,
176            default=None,
177            help="Head dimension in transformer",
178        )
179        group.add_argument(
180            "--model.classification.mit.number-heads",
181            type=int,
182            default=None,
183            help="Number of heads in transformer",
184        )
185        return parser
186
187    def _make_layer(
188        self,
189        opts,
190        input_channel,
191        cfg: Dict,
192        dilate: Optional[bool] = False,
193        *args,
194        **kwargs
195    ) -> Tuple[nn.Sequential, int]:
196        block_type = cfg.get("block_type", "mobilevit")
197        if block_type.lower() == "mobilevit":
198            return self._make_mit_layer(
199                opts=opts, input_channel=input_channel, cfg=cfg, dilate=dilate
200            )
201        else:
202            return self._make_mobilenet_layer(
203                opts=opts, input_channel=input_channel, cfg=cfg
204            )
205
206    @staticmethod
207    def _make_mobilenet_layer(
208        opts, input_channel: int, cfg: Dict, *args, **kwargs
209    ) -> Tuple[nn.Sequential, int]:
210        output_channels = cfg.get("out_channels")
211        num_blocks = cfg.get("num_blocks", 2)
212        expand_ratio = cfg.get("expand_ratio", 4)
213        block = []
214
215        for i in range(num_blocks):
216            stride = cfg.get("stride", 1) if i == 0 else 1
217
218            layer = InvertedResidual(
219                opts=opts,
220                in_channels=input_channel,
221                out_channels=output_channels,
222                stride=stride,
223                expand_ratio=expand_ratio,
224            )
225            block.append(layer)
226            input_channel = output_channels
227        return nn.Sequential(*block), input_channel
228
229    def _make_mit_layer(
230        self,
231        opts,
232        input_channel,
233        cfg: Dict,
234        dilate: Optional[bool] = False,
235        *args,
236        **kwargs
237    ) -> Tuple[nn.Sequential, int]:
238        prev_dilation = self.dilation
239        block = []
240        stride = cfg.get("stride", 1)
241
242        if stride == 2:
243            if dilate:
244                self.dilation *= 2
245                stride = 1
246
247            layer = InvertedResidual(
248                opts=opts,
249                in_channels=input_channel,
250                out_channels=cfg.get("out_channels"),
251                stride=stride,
252                expand_ratio=cfg.get("mv_expand_ratio", 4),
253                dilation=prev_dilation,
254            )
255
256            block.append(layer)
257            input_channel = cfg.get("out_channels")
258
259        head_dim = cfg.get("head_dim", 32)
260        transformer_dim = cfg["transformer_channels"]
261        ffn_dim = cfg.get("ffn_dim")
262        if head_dim is None:
263            num_heads = cfg.get("num_heads", 4)
264            if num_heads is None:
265                num_heads = 4
266            head_dim = transformer_dim // num_heads
267
268        if transformer_dim % head_dim != 0:
269            logger.error(
270                "Transformer input dimension should be divisible by head dimension. "
271                "Got {} and {}.".format(transformer_dim, head_dim)
272            )
273
274        block.append(
275            MobileViTBlock(
276                opts=opts,
277                in_channels=input_channel,
278                transformer_dim=transformer_dim,
279                ffn_dim=ffn_dim,
280                n_transformer_blocks=cfg.get("transformer_blocks", 1),
281                patch_h=cfg.get("patch_h", 2),
282                patch_w=cfg.get("patch_w", 2),
283                dropout=getattr(opts, "model.classification.mit.dropout", 0.1),
284                ffn_dropout=getattr(opts, "model.classification.mit.ffn_dropout", 0.0),
285                attn_dropout=getattr(
286                    opts, "model.classification.mit.attn_dropout", 0.1
287                ),
288                head_dim=head_dim,
289                no_fusion=getattr(
290                    opts,
291                    "model.classification.mit.no_fuse_local_global_features",
292                    False,
293                ),
294                conv_ksize=getattr(
295                    opts, "model.classification.mit.conv_kernel_size", 3
296                ),
297            )
298        )
299
300        return nn.Sequential(*block), input_channel