【MobileViT】

小橘AI

已于 2023-06-26 10:25:39 修改

阅读量4.2k

点赞数 6

分类专栏：论文阅读文章标签：机器学习深度学习人工智能

于 2022-07-11 11:29:06 首次发布

本文链接：https://blog.csdn.net/szhizhi/article/details/125717935

版权

论文阅读专栏收录该内容

8 篇文章 0 订阅

订阅专栏

MobileViT v1

轻量级的卷积神经网络在空间上局部建模，如果想要学习全局表征，可以采用基于自注意的视觉Transformer（ViT），但ViTs的参数量比较大，因此作者提出了Mobile V i T。

知识背景

自注意

加权融合，QKV都是输入x乘以对应的W权值矩阵得到，W权值会更新学习；
在这里插入图片描述

dk代表K的维度，同样的有dq、dv，这样做对权值矩阵进行一次缩放再送入softmax，因为原输入乘以权值矩阵后，得到的输出矩阵中元素方差很大，会使得softmax的分布变得陡峭影响梯度稳定计算，进行一次缩放后softmax分布能变得平缓，进而在之后的训练过程中保持梯度稳定。

ViT

将输入reshape为一系列patch，然后将其投影到固定的维度空间中得到Xp，然后使用一组Transformer块学习patch间的表示。

在这里插入图片描述

输入x∈R(H* W* C），reshape为一系列patches，Xf∈R（N* PC），P=w* h, 是patch中的像素数，N为patch数，通过Linear缩放维度为Xp∈R(N*d),

缺点：

1、忽略了CNN固有的空间归纳偏置，因此需要更多的参数来学习视觉特征；(这个地方我理解就是把像素点全部混在一起，图像原有的空间位置被忽略了)

2、相比CNN，优化能力更弱，需要大量的数据增强和L2正则化以防止过度拟合；

3、对于密集预测任务，需要昂贵的解码器。

transformer块

InputEmbedding：对输入序列进行语义信息转换，还有位置编码；

vision中transformer：只有编码过程，增加一个可学习类作为最终输入分类的向量，并通过concat方式与原一维图片向量进行拼接；MLP是分类处理部分，利用学习得到的分类向量输入MLP中；（这个学习向量在哪？）

在这里插入图片描述

#将x转化成qkv，
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
qkv = self.to_qkv(x).chunk(3, dim = -1)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)

multi-head attention：h就是多头，允许模型在不同子空间里学习到相关的信息。

在这里插入图片描述

scaled Dot-Product Attention：QK先做矩阵乘法，再进行维度缩放，mask层只在decoder部分使用，经过softmax后与V相乘；QK相乘可以看作图片块之间的关联性，获得注意力权值后再scale到V

在这里插入图片描述

参考【Vision Transformer】超详解+个人心得 - 知乎 (zhihu.com)

基本思想

结合CNN（固有的空间偏置归纳和对数据增强的较低敏感性）和ViT（自适应加权和全局信息处理），有效地将局部和全局信息进行编码，从不同角度学习全局表示。

Mobile ViT块

标准卷积涉及三个操作：展开+局部处理+折叠，利用Transformer将卷积中的局部建模替换为全局建模，这使得MobileViT具有CNN和ViT的性质。

在这里插入图片描述

Mobile ViT块中，n* n卷积后跟一个1* 1卷积，n* n编码局部空间，1* 1卷积体通过学习channel线性组合，将向量投影到高维空间。
长尺度非局部的依赖：dilated convolutions，需要谨慎选择dilation rates；或者权值也应用于填0部分而不仅仅是有效区域。（？）
保留空间固有偏置：将XL∈R(H* W* d)展开成XU∈R(N* P* d)，P=w*h是原来部分数据的长宽，N=HW/P是patch数量,需要注意wh分别需要被WH整除；经过transformer之后得到XG，再被展开成XF(H * W * d).
再使用1*1卷积将其投影到低维空间，并与x concatenation，n * n卷积融合x和经过处理的数据。
XU编码局部信息，XG针对每个局部区域，通过P patches编码全局信息，XG中每个像素可以编码X中所有像素的信息； [外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-FoMaSGgI-1657509646125)(image-20220706161603435.png)]
每个像素能看到其他所有像素，红色像素在transformer中能注意到蓝色像素，同理蓝色也能注意到周边像素，因此红色能注意到所有像素，也就能编码整张图像。（当patch过大时，无法收集到所有的像素信息）

在这里插入图片描述

轻量级在于中间使用transformer学习全局信息而且没有损失空间偏置，ViT需要更多的算力去学习视觉表示；

多尺度采样训练

给定一系列排序的空间分辨率S = {(H1, W1), ··· , (Hn, Wn)}，最大空间分辨率有最小的batch，加快优化更新；在每个GPU第t次迭代中随机抽样一个空间分辨率，然后计算迭代大小；

在这里插入图片描述

相较于以前多尺度采样，这次它不需要自己每隔几个iteration微调得到新的空间分辨率，并且改变batch提高了训练速度；使用多GPU进行训练（我猜不同空间分辨率在不同的GPU上运行）

这个可以提高网络的泛化能力，减少训练和验证之间差距；并且适用于其他网络训练；

实验过程

度量没有使用FLOPs是因为内存访问、并行度、平台特性的问题，导致FLOP无法在移动设备上实现低延迟（仅在Image-1K上对比）。

1、数据集Image-1K，对比ViTs，准确率上升，参数量降低

数据集：Image-1k；300epochs on 8 NVIDIA；batch size=1024

损失函数：cross-entropy loss；学习率：从0.0002到0.002在最初3k iteration，然后用余弦退火到0.0002；

激活函数：Swish；正则化：L2；

在这里插入图片描述

2、主干网络能力对比

在物体检测和语义分割任务上，用MobileViT做骨干网络，和MobileNet对比；

物体检测：用SSD做后续特征处理，并用深度分离卷积替代了SSD中的标准卷积；在MS-COCO上做的训练；

语义分割：用DeepLabv3做后续特征处理，在PASCAL VOC 2012上做的训练；参数量降低，准确率也有部分降低；

在这里插入图片描述

3、在移动设备上对比

运行速度比MobileNetv2要慢

理由：GPU上有专用的CUDA核做transformer，但这些在ViTs中被用来out-of-the-box来提高其在GPU上的可伸缩性核效率；CNN受益于几个设备级别的优化，包括卷积层的批量归一化融合；

在这里插入图片描述

MobileViT v2

主要思想

降低多头自注意时间复杂度有两个方向（tokens就是patches）：

1、在自注意层引入sparsity，在输入序列中每个token引入tokens一个子集；使用预定义模式限制token输入（不接受所有的tokens而是接受子集，缺点训练样本少性能下降很快）或者使用局部敏感的hash分组tokens（大型序列上才能看到提升）；

2、通过低秩矩阵估计得到近似自注意矩阵，由线性连接将自注意操作分解成多个更小的自注意操作（Linformer使用batch-wise矩阵乘法）；

主要是为了解决v1版本的高延迟问题，本论文用分离自注意代替多头自注意提高效率，使用element-wise操作替代batch-wise矩阵乘法；

在这里插入图片描述

MHA

dh=d/h,最后输出k个d维tokens，这个输出会在做一次矩阵乘法变成k*d维向量，作为最后的输出；

在这里插入图片描述

Separable self-attention

计算latent token L的上下文得分，这些分数重新加权输入token并生成上下文向量，这个向量编码了全局信息。

在这里插入图片描述

分支L用矩阵(b)L将x中每个d维向量映射到标量，计算(b)L与x的距离得到一个k维向量，这个k维向量softmax后就是上下文得分cs；

分支K直接矩阵相乘得到输出Xk,与cs相乘并相加k层，得到cv，cv类似于MHA的a矩阵，也编码了所有x的输入；

在这里插入图片描述

分支V线性映射并由ReLU激活得到Xv，然后与cv element-wise相乘，最后通过线性层得到最后的输出。

在这里插入图片描述

实验过程

时间复杂度不能解释所有操作的成本，因此使用了不同k在CPU上运行；只针对Transformer这个块，分离MHA确实降低了延迟，准确率略有下降。

在ImageNet-21k-P上预训练，在ImageNet-1K上微调，预训练初始化使用ImageNet-1k的权重提高收敛速度；

1、在手机上，MobileViT比MobileFormer速度要快，但是在GPU上两者一样。

2、ConvNexT速度比MobileViTv2快，因为手机对CNN模型有优化；

3、分辨率上升时，ConvNexT和MobileViTv2之间的差距减小了，表明ViT模型有更好的缩放性能；

MobileViT v2代码讲解

网络架构

在这里插入图片描述

mobilevit block

在这里插入图片描述

将X输入3*3卷积：

 conv_3x3_in = ConvLayer(
            opts=opts,
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=conv_ksize,
            stride=1,
            use_norm=True,
            use_act=True,
            dilation=dilation, #长尺度非局部的依赖
        )

再用1*1卷积做高维映射：

 conv_1x1_in = ConvLayer(
            opts=opts,
            in_channels=in_channels,
            out_channels=transformer_dim, #得到[B,C,H,W]形状向量
            kernel_size=1,
            stride=1,
            use_norm=False,
            use_act=False,
        )

Unfold过程：

 def unfolding(self, feature_map: Tensor) -> Tuple[Tensor, Dict]:
        patch_w, patch_h = self.patch_w, self.patch_h
        patch_area = int(patch_w * patch_h)
        batch_size, in_channels, orig_h, orig_w = feature_map.shape

        new_h = int(math.ceil(orig_h / self.patch_h) * self.patch_h)
        new_w = int(math.ceil(orig_w / self.patch_w) * self.patch_w) 
        #math.ceil返回ori/patch最小整数值

        interpolate = False #如果不能整除H或W，transformer中需要有特殊处理
        if new_w != orig_w or new_h != orig_h:
            # Note: Padding can be done, but then it needs to be handled in attention function.
            feature_map = F.interpolate(
                feature_map, size=(new_h, new_w), mode="bilinear", align_corners=False
            )
            interpolate = True

        # number of patches along width and height
        num_patch_w = new_w // patch_w  # n_w
        num_patch_h = new_h // patch_h  # n_h
        num_patches = num_patch_h * num_patch_w  # N

        # [B, C, H, W] --> [B * C * n_h, p_h, n_w, p_w] 先分成魔方那个形状
        reshaped_fm = feature_map.reshape(
            batch_size * in_channels * num_patch_h, patch_h, num_patch_w, patch_w
        )
        # [B * C * n_h, p_h, n_w, p_w] --> [B * C * n_h, n_w, p_h, p_w]
        #改变顺序是便于后续压扁
        transposed_fm = reshaped_fm.transpose(1, 2)
        # [B * C * n_h, n_w, p_h, p_w] --> [B, C, N, P] where P = p_h * p_w and N = n_h * n_w
        reshaped_fm = transposed_fm.reshape(
            batch_size, in_channels, num_patches, patch_area
        )
        # [B, C, N, P] --> [B, P, N, C]
        transposed_fm = reshaped_fm.transpose(1, 3)
        # [B, P, N, C] --> [BP, N, C]
        patches = transposed_fm.reshape(batch_size * patch_area, num_patches, -1)

        info_dict = {
            "orig_size": (orig_h, orig_w),
            "batch_size": batch_size,
            "interpolate": interpolate,
            "total_patches": num_patches,
            "num_patches_w": num_patch_w,
            "num_patches_h": num_patch_h,
        }

        return patches, info_dict

transformer过后，Flod过程（就是unflod逆过程）：

    def folding(self, patches: Tensor, info_dict: Dict) -> Tensor:
        n_dim = patches.dim()
        assert n_dim == 3, "Tensor should be of shape BPxNxC. Got: {}".format(
            patches.shape
        )
        # [BP, N, C] --> [B, P, N, C]
        patches = patches.contiguous().view(
            info_dict["batch_size"], self.patch_area, info_dict["total_patches"], -1
        )

        batch_size, pixels, num_patches, channels = patches.size()
        num_patch_h = info_dict["num_patches_h"]
        num_patch_w = info_dict["num_patches_w"]

        # [B, P, N, C] --> [B, C, N, P]
        patches = patches.transpose(1, 3)

        # [B, C, N, P] --> [B*C*n_h, n_w, p_h, p_w]
        feature_map = patches.reshape(
            batch_size * channels * num_patch_h, num_patch_w, self.patch_h, self.patch_w
        )
        # [B*C*n_h, n_w, p_h, p_w] --> [B*C*n_h, p_h, n_w, p_w]
        feature_map = feature_map.transpose(1, 2)
        # [B*C*n_h, p_h, n_w, p_w] --> [B, C, H, W]
        feature_map = feature_map.reshape(
            batch_size, channels, num_patch_h * self.patch_h, num_patch_w * self.patch_w
        )
        if info_dict["interpolate"]:
            feature_map = F.interpolate(
                feature_map,
                size=info_dict["orig_size"],
                mode="bilinear",
                align_corners=False,
            )
        return feature_map

transformer块

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-MVWXyEYh-1657509646133)(image-20220708161840912.png)]

        qkv = self.qkv_proj(x) #卷积映射成1+2*（embed_dim) channels

        # Project x into query, key and value
        # Query --> [B, 1, P, N]，x中每个d维向量映射到标量
        # value, key --> [B, d, P, N]
        query, key, value = torch.split(
            qkv, split_size_or_sections=[1, self.embed_dim, self.embed_dim], dim=1
        )

        # apply softmax along N dimension
        context_scores = F.softmax(query, dim=-1)
        # Uncomment below line to visualize context scores
        # self.visualize_context_scores(context_scores=context_scores)
        #tensor -> [tensor, float, float]
        context_scores = self.attn_dropout(context_scores)

        # Compute context vector
        # [B, d, P, N] x [B, 1, P, N] -> [B, d, P, N]
        #其实这里我们可以看到kv两个分支并没有一个线性映射过程
        context_vector = key * context_scores
        # [B, d, P, N] --> [B, d, P, 1]
        # 编码全局信息
        context_vector = torch.sum(context_vector, dim=-1, keepdim=True)

        # combine context vector with values
        # [B, d, P, N] * [B, d, P, 1] --> [B, d, P, N]
        out = F.relu(value) * context_vector.expand_as(value)
        out = self.out_proj(out)
        return out

mobilenetv2块

        block = nn.Sequential()
        if expand_ratio != 1:
            block.add_module(
                name="exp_1x1",
                module=ConvLayer(
                    opts,
                    in_channels=in_channels,
                    out_channels=hidden_dim,
                    kernel_size=1,
                    use_act=False,
                    use_norm=True,
                ),
            )
            block.add_module(name="act_fn_1", module=act_fn)

        block.add_module(
            name="conv_3x3",
            module=ConvLayer(
                opts,
                in_channels=hidden_dim,
                out_channels=hidden_dim,
                stride=stride,
                kernel_size=3,
                groups=hidden_dim,
                use_act=False,
                use_norm=True,
                dilation=dilation,
            ),
        )
        block.add_module(name="act_fn_2", module=act_fn)

        if use_se:
            se = SqueezeExcitation(
                opts=opts,
                in_channels=hidden_dim,
                squeeze_factor=4,
                scale_fn_name="hard_sigmoid",
            )
            block.add_module(name="se", module=se)

        block.add_module(
            name="red_1x1",
            module=ConvLayer(
                opts,
                in_channels=hidden_dim,
                out_channels=out_channels,
                kernel_size=1,
                use_act=False,
                use_norm=True,
            ),
        )

多尺度采样

    #设置最小和最大size    
    if is_training:
            self.img_batch_tuples = _image_batch_pairs(
                crop_size_h=self.crop_size_h,
                crop_size_w=self.crop_size_w,
                batch_size_gpu0=self.batch_size_gpu0,
                n_gpus=self.num_replicas,
                max_scales=self.max_img_scales,
                check_scale_div_factor=self.check_scale_div_factor,
                min_crop_size_w=self.min_crop_size_w,
                max_crop_size_w=self.max_crop_size_w,
                min_crop_size_h=self.min_crop_size_h,
                max_crop_size_h=self.max_crop_size_h,
            )
            self.img_batch_tuples = [
                (h, w, self.batch_size_gpu0) for h, w, b in self.img_batch_tuples
            ]

    def __iter__(self):
        if self.shuffle:
            random.seed(self.epoch)
            indices_rank_i = self.img_indices[
                self.rank : len(self.img_indices) : self.num_replicas
            ]
            random.shuffle(indices_rank_i)
        else:
            indices_rank_i = self.img_indices[
                self.rank : len(self.img_indices) : self.num_replicas
            ]

        start_index = 0
        n_samples_rank_i = len(indices_rank_i)
        while start_index < n_samples_rank_i:
            crop_h, crop_w, batch_size = random.choice(self.img_batch_tuples) #根据tuples设置随机大小

            end_index = min(start_index + batch_size, n_samples_rank_i)
            batch_ids = indices_rank_i[start_index:end_index]
            n_batch_samples = len(batch_ids)
            if n_batch_samples != batch_size:
                batch_ids += indices_rank_i[: (batch_size - n_batch_samples)]
            start_index += batch_size

            if len(batch_ids) > 0: #设置batch大小
                batch = [(crop_h, crop_w, b_id) for b_id in batch_ids]
                yield batch

实验结果

原代码所需参数特别多，利用作者提供的yaml文件进行训练
我训练用的是imagenetMINI（因为没有服务器上没有更大的数据集）

可能出现的问题：
1、tensor

2022-07-15 15:19:19 - ERROR   - Unable to stack the tensors. Error: expected Tensor as element 0 in argument 0, but got numpy.int64
2022-07-15 15:19:19 - ERROR   - Exiting!!!
2022-07-15 15:19:19 - LOGS    - Exception occurred that interrupted the training. DataLoader worker (pid(s) 22338) exited unexpectedly
DataLoader worker (pid(s) 22338) exited unexpectedly

#需要将yaml文件中的imagenet_opencv
name: "imagenet_opencv"
#改成
name: "imagenet"

2、resize为空，在yaml文件中添加resize

  resize:
    enable: true
    size: 288 # shorter size is 288
    interpolation: "bicubic"
  center_crop:
    enable: true
    size: 256

3、load ckpt出现问题，好像是因为优化函数 Adam，pytorch升级之后出现的问题，至少要降回1.11.0之前版本，这个问题据说会在1.12.1版本解决；

#强制将优化参数改为ture，这个会导致训练速度下降
#这个我试了没用，会报错
optim.param_groups[0]['capturable'] = True

#下一个问题采用这个没用，新建ckpt之后还需要删除这个，否则报错
#我选择直接将之前的ckpt删除，从头开始训练，先看看acc再决定要不要优化

还有解决方法是load ckpt时不加载优化器（我还没有试

checkpoint = torch.load('/path/to/last.ckpt')
lightning_module.load_state_dict(checkpoint['state_dict'])```

###实验结果
MINI，效果不好，怀疑是因为数据集太小了而且很耗时；

2022-07-20 03:42:59 - LOGS    - Epoch: 299 [   40665/10000000], loss: 2.5460, LR: [0.0002, 0.0002], Avg. batch load time: 16.448, Elapsed time: 17.50
2022-07-20 03:45:23 - LOGS    - *** Training summary for epoch 299
         loss=1.8571
2022-07-20 03:45:41 - LOGS    - Epoch: 299 [     100/    1962], loss: 3.7218, top1: 45.5000, top5: 64.0000, LR: [0.0002, 0.0002], Avg. batch load time: 0.000, Elapsed time: 15.85
2022-07-20 03:45:45 - LOGS    - *** Validation summary for epoch 299
         loss=4.2378 || top1=35.4250 || top5=56.7000
2022-07-20 03:46:03 - LOGS    - Epoch: 299 [     100/    1962], loss: 3.6208, top1: 46.0000, top5: 68.5000, LR: [0.0002, 0.0002], Avg. batch load time: 0.000, Elapsed time: 16.23
2022-07-20 03:46:07 - LOGS    - *** Validation (Ema) summary for epoch 299
         loss=4.1501 || top1=36.9500 || top5=58.5000
2022-07-20 03:46:08 - INFO    - Checkpoints saved at: classification_mbvit2/run_1
================================================================================================================================
2022-07-20 03:46:08 - LOGS    - Training took 17:28:53.65

换了train训练集，ImageNet，top1已经有68%了，但我认为这个top1不太有说服力，因为val跑的是MINI里面的val（论文里面有78.4%）

不知道ImageNet本身有没有val数据集，现在MINI的top1有80.4%；

2022-08-01 08:32:09 - LOGS    - *** Training summary for epoch 299
         loss=2.2574
2022-08-01 08:32:42 - LOGS    - Epoch: 299 [     100/     981], loss: 1.6015, top1: 82.0000, top5: 96.2500, LR: [0.0002, 0.0002], Avg. batch load time: 0.000, Elapsed time: 30.86
2022-08-01 08:32:53 - LOGS    - *** Validation summary for epoch 299
         loss=1.8304 || top1=77.6750 || top5=93.7250
2022-08-01 08:33:13 - LOGS    - Epoch: 299 [     100/     981], loss: 1.5313, top1: 84.5000, top5: 96.7500, LR: [0.0002, 0.0002], Avg. batch load time: 0.000, Elapsed time: 18.14
2022-08-01 08:33:15 - LOGS    - *** Validation (Ema) summary for epoch 299
         loss=1.7275 || top1=80.4000 || top5=94.5500
2022-08-01 08:33:17 - LOGS    - Best EMA checkpoint with score 80.40 saved at classification_mbvit2_imagenet/run_1/checkpoint_ema_best.pt