1212 lines
39 KiB
Python
1212 lines
39 KiB
Python
# copyright (c) 2023 PaddlePaddle Authors. All Rights Reserve.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# Code was based on https://github.com/micronDLA/MobileViTv3/blob/main/MobileViTv3-v1/cvnets/models/classification/mobilevit.py
|
|
# reference: https://arxiv.org/abs/2209.15159
|
|
|
|
import math
|
|
from functools import partial
|
|
from typing import Dict, Optional, Tuple, Union
|
|
|
|
import paddle
|
|
import paddle.nn as nn
|
|
import paddle.nn.functional as F
|
|
|
|
from ....utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
|
|
|
|
MODEL_URLS = {
|
|
"MobileViTv3_XXS": "",
|
|
"MobileViTv3_XS": "",
|
|
"MobileViTv3_S": "",
|
|
"MobileViTv3_x0_5": "",
|
|
"MobileViTv3_x0_75": "",
|
|
"MobileViTv3_x1_0": "",
|
|
}
|
|
|
|
layer_norm_2d = partial(nn.GroupNorm, num_groups=1)
|
|
|
|
|
|
def make_divisible(v, divisor=8, min_value=None):
|
|
if min_value is None:
|
|
min_value = divisor
|
|
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
|
|
if new_v < 0.9 * v:
|
|
new_v += divisor
|
|
return new_v
|
|
|
|
|
|
class InvertedResidual(nn.Layer):
|
|
"""
|
|
Inverted residual block (MobileNetv2): https://arxiv.org/abs/1801.04381
|
|
"""
|
|
|
|
def __init__(self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
stride: int,
|
|
expand_ratio: Union[int, float],
|
|
dilation: int=1) -> None:
|
|
assert stride in [1, 2]
|
|
super(InvertedResidual, self).__init__()
|
|
self.stride = stride
|
|
|
|
hidden_dim = make_divisible(int(round(in_channels * expand_ratio)), 8)
|
|
self.use_res_connect = self.stride == 1 and in_channels == out_channels
|
|
|
|
block = nn.Sequential()
|
|
if expand_ratio != 1:
|
|
block.add_sublayer(
|
|
name="exp_1x1",
|
|
sublayer=nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
in_channels, hidden_dim, 1, bias_attr=False)),
|
|
('norm', nn.BatchNorm2D(hidden_dim)), ('act', nn.Silu())))
|
|
|
|
block.add_sublayer(
|
|
name="conv_3x3",
|
|
sublayer=nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
hidden_dim,
|
|
hidden_dim,
|
|
3,
|
|
bias_attr=False,
|
|
stride=stride,
|
|
padding=dilation,
|
|
dilation=dilation,
|
|
groups=hidden_dim)), ('norm', nn.BatchNorm2D(hidden_dim)),
|
|
('act', nn.Silu())))
|
|
|
|
block.add_sublayer(
|
|
name="red_1x1",
|
|
sublayer=nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
hidden_dim, out_channels, 1, bias_attr=False)),
|
|
('norm', nn.BatchNorm2D(out_channels))))
|
|
|
|
self.block = block
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.exp = expand_ratio
|
|
self.dilation = dilation
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
if self.use_res_connect:
|
|
return x + self.block(x)
|
|
else:
|
|
return self.block(x)
|
|
|
|
|
|
class MultiHeadAttention(nn.Layer):
|
|
def __init__(self,
|
|
dim,
|
|
num_heads=8,
|
|
qkv_bias=False,
|
|
qk_scale=None,
|
|
attn_drop=0.):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
head_dim = dim // num_heads
|
|
self.scale = qk_scale or head_dim**-0.5
|
|
|
|
self.qkv_proj = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
|
|
self.attn_drop = nn.Dropout(attn_drop)
|
|
self.out_proj = nn.Linear(dim, dim, bias_attr=qkv_bias)
|
|
|
|
def forward(self, x):
|
|
# B = paddle.shape(x)[0]
|
|
N, C = x.shape[1:]
|
|
qkv = self.qkv_proj(x).reshape((-1, N, 3, self.num_heads,
|
|
C // self.num_heads)).transpose(
|
|
(2, 0, 3, 1, 4))
|
|
q, k, v = qkv[0], qkv[1], qkv[2]
|
|
|
|
attn = (q.matmul(k.transpose((0, 1, 3, 2)))) * self.scale
|
|
attn = nn.functional.softmax(attn, axis=-1)
|
|
attn = self.attn_drop(attn)
|
|
|
|
x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((-1, N, C))
|
|
x = self.out_proj(x)
|
|
return x
|
|
|
|
|
|
class TransformerEncoder(nn.Layer):
|
|
"""
|
|
This class defines the Transformer encoder (pre-norm) as described in "Attention is all you need" paper
|
|
https://arxiv.org/abs/1706.03762
|
|
"""
|
|
|
|
def __init__(self,
|
|
embed_dim: int,
|
|
ffn_latent_dim: int,
|
|
num_heads: Optional[int]=8,
|
|
attn_dropout: Optional[float]=0.0,
|
|
dropout: Optional[float]=0.1,
|
|
ffn_dropout: Optional[float]=0.0,
|
|
transformer_norm_layer: nn.Layer=nn.LayerNorm):
|
|
super(TransformerEncoder, self).__init__()
|
|
|
|
self.pre_norm_mha = nn.Sequential(
|
|
transformer_norm_layer(embed_dim),
|
|
MultiHeadAttention(
|
|
embed_dim, num_heads, attn_drop=attn_dropout, qkv_bias=True),
|
|
nn.Dropout(p=dropout))
|
|
|
|
self.pre_norm_ffn = nn.Sequential(
|
|
transformer_norm_layer(embed_dim),
|
|
nn.Linear(embed_dim, ffn_latent_dim),
|
|
nn.Silu(),
|
|
nn.Dropout(p=ffn_dropout),
|
|
nn.Linear(ffn_latent_dim, embed_dim),
|
|
nn.Dropout(p=dropout))
|
|
self.embed_dim = embed_dim
|
|
self.ffn_dim = ffn_latent_dim
|
|
self.ffn_dropout = ffn_dropout
|
|
|
|
def forward(self, x):
|
|
# Multi-head attention
|
|
x = x + self.pre_norm_mha(x)
|
|
|
|
# Feed forward network
|
|
x = x + self.pre_norm_ffn(x)
|
|
return x
|
|
|
|
|
|
class MobileViTv3Block(nn.Layer):
|
|
"""
|
|
MobileViTv3 block
|
|
"""
|
|
|
|
def __init__(self,
|
|
in_channels: int,
|
|
transformer_dim: int,
|
|
ffn_dim: int,
|
|
n_transformer_blocks: Optional[int]=2,
|
|
head_dim: Optional[int]=32,
|
|
attn_dropout: Optional[float]=0.1,
|
|
dropout: Optional[int]=0.1,
|
|
ffn_dropout: Optional[int]=0.1,
|
|
patch_h: Optional[int]=8,
|
|
patch_w: Optional[int]=8,
|
|
transformer_norm_layer: nn.Layer=nn.LayerNorm,
|
|
conv_ksize: Optional[int]=3,
|
|
dilation: Optional[int]=1,
|
|
var_ffn: Optional[bool]=False,
|
|
no_fusion: Optional[bool]=False):
|
|
|
|
# For MobileViTv3: Normal 3x3 convolution --> Depthwise 3x3 convolution
|
|
padding = (conv_ksize - 1) // 2 * dilation
|
|
conv_3x3_in = nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
in_channels,
|
|
in_channels,
|
|
conv_ksize,
|
|
bias_attr=False,
|
|
padding=padding,
|
|
dilation=dilation,
|
|
groups=in_channels)), ('norm', nn.BatchNorm2D(in_channels)),
|
|
('act', nn.Silu()))
|
|
conv_1x1_in = nn.Sequential(('conv', nn.Conv2D(
|
|
in_channels, transformer_dim, 1, bias_attr=False)))
|
|
|
|
conv_1x1_out = nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
transformer_dim, in_channels, 1, bias_attr=False)),
|
|
('norm', nn.BatchNorm2D(in_channels)), ('act', nn.Silu()))
|
|
conv_3x3_out = None
|
|
|
|
# For MobileViTv3: input+global --> local+global
|
|
if not no_fusion:
|
|
#input_ch = tr_dim + in_ch
|
|
conv_3x3_out = nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
transformer_dim + in_channels,
|
|
in_channels,
|
|
1,
|
|
bias_attr=False)), ('norm', nn.BatchNorm2D(in_channels)),
|
|
('act', nn.Silu()))
|
|
|
|
super().__init__()
|
|
self.local_rep = nn.Sequential()
|
|
self.local_rep.add_sublayer(name="conv_3x3", sublayer=conv_3x3_in)
|
|
self.local_rep.add_sublayer(name="conv_1x1", sublayer=conv_1x1_in)
|
|
|
|
assert transformer_dim % head_dim == 0
|
|
num_heads = transformer_dim // head_dim
|
|
|
|
ffn_dims = [ffn_dim] * n_transformer_blocks
|
|
|
|
global_rep = [
|
|
TransformerEncoder(
|
|
embed_dim=transformer_dim,
|
|
ffn_latent_dim=ffn_dims[block_idx],
|
|
num_heads=num_heads,
|
|
attn_dropout=attn_dropout,
|
|
dropout=dropout,
|
|
ffn_dropout=ffn_dropout,
|
|
transformer_norm_layer=transformer_norm_layer)
|
|
for block_idx in range(n_transformer_blocks)
|
|
]
|
|
global_rep.append(transformer_norm_layer(transformer_dim))
|
|
self.global_rep = nn.Sequential(*global_rep)
|
|
|
|
self.conv_proj = conv_1x1_out
|
|
|
|
self.fusion = conv_3x3_out
|
|
|
|
self.patch_h = patch_h
|
|
self.patch_w = patch_w
|
|
self.patch_area = self.patch_w * self.patch_h
|
|
|
|
self.cnn_in_dim = in_channels
|
|
self.cnn_out_dim = transformer_dim
|
|
self.n_heads = num_heads
|
|
self.ffn_dim = ffn_dim
|
|
self.dropout = dropout
|
|
self.attn_dropout = attn_dropout
|
|
self.ffn_dropout = ffn_dropout
|
|
self.dilation = dilation
|
|
self.ffn_max_dim = ffn_dims[0]
|
|
self.ffn_min_dim = ffn_dims[-1]
|
|
self.var_ffn = var_ffn
|
|
self.n_blocks = n_transformer_blocks
|
|
self.conv_ksize = conv_ksize
|
|
|
|
def unfolding(self, feature_map):
|
|
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)
|
|
|
|
interpolate = False
|
|
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([0, 2, 1, 3])
|
|
# [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([0, 3, 2, 1])
|
|
# [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
|
|
|
|
def folding(self, patches, info_dict):
|
|
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.reshape([
|
|
info_dict["batch_size"], self.patch_area,
|
|
info_dict["total_patches"], -1
|
|
])
|
|
|
|
batch_size, pixels, num_patches, channels = patches.shape
|
|
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([0, 3, 2, 1])
|
|
|
|
# [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([0, 2, 1, 3])
|
|
# [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
|
|
|
|
def forward(self, x):
|
|
res = x
|
|
|
|
# For MobileViTv3: Normal 3x3 convolution --> Depthwise 3x3 convolution
|
|
fm_conv = self.local_rep(x)
|
|
|
|
# convert feature map to patches
|
|
patches, info_dict = self.unfolding(fm_conv)
|
|
|
|
# learn global representations
|
|
patches = self.global_rep(patches)
|
|
|
|
# [B x Patch x Patches x C] --> [B x C x Patches x Patch]
|
|
fm = self.folding(patches=patches, info_dict=info_dict)
|
|
|
|
fm = self.conv_proj(fm)
|
|
|
|
if self.fusion is not None:
|
|
# For MobileViTv3: input+global --> local+global
|
|
fm = self.fusion(paddle.concat((fm_conv, fm), axis=1))
|
|
|
|
# For MobileViTv3: Skip connection
|
|
fm = fm + res
|
|
|
|
return fm
|
|
|
|
|
|
class LinearSelfAttention(nn.Layer):
|
|
def __init__(self, embed_dim, attn_dropout=0.0, bias=True):
|
|
super().__init__()
|
|
self.qkv_proj = nn.Conv2D(
|
|
embed_dim, 1 + (2 * embed_dim), 1, bias_attr=bias)
|
|
self.attn_dropout = nn.Dropout(p=attn_dropout)
|
|
self.out_proj = nn.Conv2D(embed_dim, embed_dim, 1, bias_attr=bias)
|
|
|
|
def forward(self, x):
|
|
# [B, C, P, N] --> [B, h + 2d, P, N]
|
|
qkv = self.qkv_proj(x)
|
|
|
|
# Project x into query, key and value
|
|
# Query --> [B, 1, P, N]
|
|
# value, key --> [B, d, P, N]
|
|
query, key, value = paddle.split(
|
|
qkv, [1, self.embed_dim, self.embed_dim], axis=1)
|
|
|
|
# apply softmax along N dimension
|
|
context_scores = F.softmax(query, axis=-1)
|
|
# Uncomment below line to visualize context scores
|
|
# self.visualize_context_scores(context_scores=context_scores)
|
|
context_scores = self.attn_dropout(context_scores)
|
|
|
|
# Compute context vector
|
|
# [B, d, P, N] x [B, 1, P, N] -> [B, d, P, N]
|
|
context_vector = key * context_scores
|
|
# [B, d, P, N] --> [B, d, P, 1]
|
|
context_vector = paddle.sum(context_vector, axis=-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
|
|
|
|
|
|
class LinearAttnFFN(nn.Layer):
|
|
def __init__(self,
|
|
embed_dim: int,
|
|
ffn_latent_dim: int,
|
|
attn_dropout: Optional[float]=0.0,
|
|
dropout: Optional[float]=0.1,
|
|
ffn_dropout: Optional[float]=0.0,
|
|
norm_layer: Optional[str]=layer_norm_2d) -> None:
|
|
super().__init__()
|
|
attn_unit = LinearSelfAttention(
|
|
embed_dim=embed_dim, attn_dropout=attn_dropout, bias=True)
|
|
|
|
self.pre_norm_attn = nn.Sequential(
|
|
norm_layer(num_channels=embed_dim),
|
|
attn_unit,
|
|
nn.Dropout(p=dropout))
|
|
|
|
self.pre_norm_ffn = nn.Sequential(
|
|
norm_layer(num_channels=embed_dim),
|
|
nn.Conv2D(embed_dim, ffn_latent_dim, 1),
|
|
nn.Silu(),
|
|
nn.Dropout(p=ffn_dropout),
|
|
nn.Conv2D(ffn_latent_dim, embed_dim, 1),
|
|
nn.Dropout(p=dropout))
|
|
|
|
def forward(self, x):
|
|
# self-attention
|
|
x = x + self.pre_norm_attn(x)
|
|
# Feed forward network
|
|
x = x + self.pre_norm_ffn(x)
|
|
return x
|
|
|
|
|
|
class MobileViTv3Block_v2(nn.Layer):
|
|
"""
|
|
This class defines the `MobileViTv3 block`
|
|
"""
|
|
|
|
def __init__(self,
|
|
in_channels: int,
|
|
attn_unit_dim: int,
|
|
ffn_multiplier: float=2.0,
|
|
n_attn_blocks: Optional[int]=2,
|
|
attn_dropout: Optional[float]=0.0,
|
|
dropout: Optional[float]=0.0,
|
|
ffn_dropout: Optional[float]=0.0,
|
|
patch_h: Optional[int]=8,
|
|
patch_w: Optional[int]=8,
|
|
conv_ksize: Optional[int]=3,
|
|
dilation: Optional[int]=1,
|
|
attn_norm_layer: Optional[str]=layer_norm_2d):
|
|
cnn_out_dim = attn_unit_dim
|
|
|
|
padding = (conv_ksize - 1) // 2 * dilation
|
|
conv_3x3_in = nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
in_channels,
|
|
in_channels,
|
|
conv_ksize,
|
|
bias_attr=False,
|
|
padding=padding,
|
|
dilation=dilation,
|
|
groups=in_channels)), ('norm', nn.BatchNorm2D(in_channels)),
|
|
('act', nn.Silu()))
|
|
conv_1x1_in = nn.Sequential(('conv', nn.Conv2D(
|
|
in_channels, cnn_out_dim, 1, bias_attr=False)))
|
|
|
|
super().__init__()
|
|
self.local_rep = nn.Sequential(conv_3x3_in, conv_1x1_in)
|
|
|
|
self.global_rep, attn_unit_dim = self._build_attn_layer(
|
|
d_model=attn_unit_dim,
|
|
ffn_mult=ffn_multiplier,
|
|
n_layers=n_attn_blocks,
|
|
attn_dropout=attn_dropout,
|
|
dropout=dropout,
|
|
ffn_dropout=ffn_dropout,
|
|
attn_norm_layer=attn_norm_layer)
|
|
|
|
# MobileViTv3: input changed from just global to local+global
|
|
self.conv_proj = nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
2 * cnn_out_dim, in_channels, 1, bias_attr=False)),
|
|
('norm', nn.BatchNorm2D(in_channels)))
|
|
|
|
self.patch_h = patch_h
|
|
self.patch_w = patch_w
|
|
|
|
def _build_attn_layer(self,
|
|
d_model: int,
|
|
ffn_mult: float,
|
|
n_layers: int,
|
|
attn_dropout: float,
|
|
dropout: float,
|
|
ffn_dropout: float,
|
|
attn_norm_layer: nn.Layer):
|
|
|
|
# ensure that dims are multiple of 16
|
|
ffn_dims = [ffn_mult * d_model // 16 * 16] * n_layers
|
|
|
|
global_rep = [
|
|
LinearAttnFFN(
|
|
embed_dim=d_model,
|
|
ffn_latent_dim=ffn_dims[block_idx],
|
|
attn_dropout=attn_dropout,
|
|
dropout=dropout,
|
|
ffn_dropout=ffn_dropout,
|
|
norm_layer=attn_norm_layer) for block_idx in range(n_layers)
|
|
]
|
|
global_rep.append(attn_norm_layer(num_channels=d_model))
|
|
|
|
return nn.Sequential(*global_rep), d_model
|
|
|
|
def unfolding(self, feature_map):
|
|
batch_size, in_channels, img_h, img_w = feature_map.shape
|
|
|
|
# [B, C, H, W] --> [B, C, P, N]
|
|
patches = F.unfold(
|
|
feature_map,
|
|
kernel_sizes=(self.patch_h, self.patch_w),
|
|
stride=(self.patch_h, self.patch_w))
|
|
patches = patches.reshape(
|
|
[batch_size, in_channels, self.patch_h * self.patch_w, -1])
|
|
|
|
return patches, (img_h, img_w)
|
|
|
|
def folding(self, patches, output_size: Tuple[int, int]):
|
|
batch_size, in_dim, patch_size, n_patches = patches.shape
|
|
|
|
# [B, C, P, N]
|
|
patches = patches.reshape([batch_size, in_dim * patch_size, n_patches])
|
|
|
|
feature_map = F.fold(
|
|
patches,
|
|
output_sizes=output_size,
|
|
kernel_sizes=(self.patch_h, self.patch_w),
|
|
stride=(self.patch_h, self.patch_w))
|
|
|
|
return feature_map
|
|
|
|
def forward(self, x):
|
|
fm_conv = self.local_rep(x)
|
|
|
|
# convert feature map to patches
|
|
patches, output_size = self.unfolding(fm_conv)
|
|
|
|
# learn global representations on all patches
|
|
patches = self.global_rep(patches)
|
|
|
|
# [B x Patch x Patches x C] --> [B x C x Patches x Patch]
|
|
fm = self.folding(patches=patches, output_size=output_size)
|
|
|
|
# MobileViTv3: local+global instead of only global
|
|
fm = self.conv_proj(paddle.concat((fm, fm_conv), axis=1))
|
|
|
|
# MobileViTv3: skip connection
|
|
fm = fm + x
|
|
|
|
return fm
|
|
|
|
|
|
class MobileViTv3(nn.Layer):
|
|
"""
|
|
MobileViTv3:
|
|
"""
|
|
|
|
def __init__(self,
|
|
mobilevit_config: Dict,
|
|
dropout=0.1,
|
|
class_num=1000,
|
|
classifier_dropout=0.1,
|
|
output_stride=None,
|
|
mobilevit_v2_based=False):
|
|
super().__init__()
|
|
self.round_nearest = 8
|
|
self.dilation = 1
|
|
self.dropout = dropout
|
|
self.mobilevit_v2_based = mobilevit_v2_based
|
|
|
|
dilate_l4 = dilate_l5 = False
|
|
if output_stride == 8:
|
|
dilate_l4 = True
|
|
dilate_l5 = True
|
|
elif output_stride == 16:
|
|
dilate_l5 = True
|
|
|
|
# store model configuration in a dictionary
|
|
in_channels = mobilevit_config["layer0"]["img_channels"]
|
|
out_channels = mobilevit_config["layer0"]["out_channels"]
|
|
self.conv_1 = nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
in_channels,
|
|
out_channels,
|
|
3,
|
|
bias_attr=False,
|
|
stride=2,
|
|
padding=1)), ('norm', nn.BatchNorm2D(out_channels)),
|
|
('act', nn.Silu()))
|
|
|
|
in_channels = out_channels
|
|
self.layer_1, out_channels = self._make_layer(
|
|
input_channel=in_channels, cfg=mobilevit_config["layer1"])
|
|
|
|
in_channels = out_channels
|
|
self.layer_2, out_channels = self._make_layer(
|
|
input_channel=in_channels, cfg=mobilevit_config["layer2"])
|
|
|
|
in_channels = out_channels
|
|
self.layer_3, out_channels = self._make_layer(
|
|
input_channel=in_channels, cfg=mobilevit_config["layer3"])
|
|
|
|
in_channels = out_channels
|
|
self.layer_4, out_channels = self._make_layer(
|
|
input_channel=in_channels,
|
|
cfg=mobilevit_config["layer4"],
|
|
dilate=dilate_l4)
|
|
|
|
in_channels = out_channels
|
|
self.layer_5, out_channels = self._make_layer(
|
|
input_channel=in_channels,
|
|
cfg=mobilevit_config["layer5"],
|
|
dilate=dilate_l5)
|
|
|
|
in_channels = out_channels
|
|
exp_channels = min(mobilevit_config["last_layer_exp_factor"] *
|
|
in_channels, 960)
|
|
if self.mobilevit_v2_based:
|
|
self.conv_1x1_exp = nn.Identity()
|
|
else:
|
|
self.conv_1x1_exp = nn.Sequential(
|
|
('conv', nn.Conv2D(
|
|
in_channels, exp_channels, 1, bias_attr=False)),
|
|
('norm', nn.BatchNorm2D(exp_channels)), ('act', nn.Silu()))
|
|
|
|
self.classifier = nn.Sequential()
|
|
self.classifier.add_sublayer(
|
|
name="global_pool",
|
|
sublayer=nn.Sequential(nn.AdaptiveAvgPool2D(1), nn.Flatten()))
|
|
if 0.0 < classifier_dropout < 1.0:
|
|
self.classifier.add_sublayer(
|
|
name="dropout", sublayer=nn.Dropout(p=classifier_dropout))
|
|
self.classifier.add_sublayer(
|
|
name="fc", sublayer=nn.Linear(exp_channels, class_num))
|
|
|
|
# weight initialization
|
|
self.apply(self._init_weights)
|
|
|
|
def _init_weights(self, m):
|
|
if isinstance(m, nn.Conv2D):
|
|
fan_out = m.weight.shape[0] * m.weight.shape[2] * m.weight.shape[3]
|
|
nn.initializer.KaimingNormal(fan_in=fan_out)(m.weight)
|
|
if m.bias is not None:
|
|
nn.initializer.Constant(0)(m.bias)
|
|
elif isinstance(m, nn.BatchNorm2D):
|
|
nn.initializer.Constant(1)(m.weight)
|
|
nn.initializer.Constant(0)(m.bias)
|
|
elif isinstance(m, nn.Linear):
|
|
nn.initializer.TruncatedNormal(std=.02)(m.weight)
|
|
if m.bias is not None:
|
|
nn.initializer.Constant(0)(m.bias)
|
|
|
|
def _make_layer(self, input_channel, cfg, dilate=False):
|
|
block_type = cfg.get("block_type", "mobilevit")
|
|
if block_type.lower() == "mobilevit":
|
|
return self._make_mit_layer(
|
|
input_channel=input_channel, cfg=cfg, dilate=dilate)
|
|
else:
|
|
return self._make_mobilenet_layer(
|
|
input_channel=input_channel, cfg=cfg)
|
|
|
|
def _make_mit_layer(self, input_channel, cfg, dilate=False):
|
|
prev_dilation = self.dilation
|
|
block = []
|
|
stride = cfg.get("stride", 1)
|
|
|
|
if stride == 2:
|
|
if dilate:
|
|
self.dilation *= 2
|
|
stride = 1
|
|
|
|
layer = InvertedResidual(
|
|
in_channels=input_channel,
|
|
out_channels=cfg.get("out_channels"),
|
|
stride=stride,
|
|
expand_ratio=cfg.get("mv_expand_ratio", 4),
|
|
dilation=prev_dilation)
|
|
|
|
block.append(layer)
|
|
input_channel = cfg.get("out_channels")
|
|
|
|
if self.mobilevit_v2_based:
|
|
block.append(
|
|
MobileViTv3Block_v2(
|
|
in_channels=input_channel,
|
|
attn_unit_dim=cfg["attn_unit_dim"],
|
|
ffn_multiplier=cfg.get("ffn_multiplier"),
|
|
n_attn_blocks=cfg.get("attn_blocks", 1),
|
|
ffn_dropout=0.,
|
|
attn_dropout=0.,
|
|
dilation=self.dilation,
|
|
patch_h=cfg.get("patch_h", 2),
|
|
patch_w=cfg.get("patch_w", 2)))
|
|
else:
|
|
head_dim = cfg.get("head_dim", 32)
|
|
transformer_dim = cfg["transformer_channels"]
|
|
ffn_dim = cfg.get("ffn_dim")
|
|
if head_dim is None:
|
|
num_heads = cfg.get("num_heads", 4)
|
|
if num_heads is None:
|
|
num_heads = 4
|
|
head_dim = transformer_dim // num_heads
|
|
|
|
assert transformer_dim % head_dim == 0, (
|
|
"Transformer input dimension should be divisible by head dimension. "
|
|
"Got {} and {}.".format(transformer_dim, head_dim))
|
|
|
|
block.append(
|
|
MobileViTv3Block(
|
|
in_channels=input_channel,
|
|
transformer_dim=transformer_dim,
|
|
ffn_dim=ffn_dim,
|
|
n_transformer_blocks=cfg.get("transformer_blocks", 1),
|
|
patch_h=cfg.get("patch_h", 2),
|
|
patch_w=cfg.get("patch_w", 2),
|
|
dropout=self.dropout,
|
|
ffn_dropout=0.,
|
|
attn_dropout=0.,
|
|
head_dim=head_dim))
|
|
|
|
return nn.Sequential(*block), input_channel
|
|
|
|
def _make_mobilenet_layer(self, input_channel, cfg):
|
|
output_channels = cfg.get("out_channels")
|
|
num_blocks = cfg.get("num_blocks", 2)
|
|
expand_ratio = cfg.get("expand_ratio", 4)
|
|
block = []
|
|
|
|
for i in range(num_blocks):
|
|
stride = cfg.get("stride", 1) if i == 0 else 1
|
|
|
|
layer = InvertedResidual(
|
|
in_channels=input_channel,
|
|
out_channels=output_channels,
|
|
stride=stride,
|
|
expand_ratio=expand_ratio)
|
|
block.append(layer)
|
|
input_channel = output_channels
|
|
return nn.Sequential(*block), input_channel
|
|
|
|
def extract_features(self, x):
|
|
x = self.conv_1(x)
|
|
x = self.layer_1(x)
|
|
x = self.layer_2(x)
|
|
x = self.layer_3(x)
|
|
|
|
x = self.layer_4(x)
|
|
x = self.layer_5(x)
|
|
x = self.conv_1x1_exp(x)
|
|
return x
|
|
|
|
def forward(self, x):
|
|
x = self.extract_features(x)
|
|
x = self.classifier(x)
|
|
return x
|
|
|
|
|
|
def _load_pretrained(pretrained, model, model_url, use_ssld=False):
|
|
if pretrained is False:
|
|
pass
|
|
elif pretrained is True:
|
|
load_dygraph_pretrain_from_url(model, model_url, use_ssld=use_ssld)
|
|
elif isinstance(pretrained, str):
|
|
load_dygraph_pretrain(model, pretrained)
|
|
else:
|
|
raise RuntimeError(
|
|
"pretrained type is not available. Please use `string` or `boolean` type."
|
|
)
|
|
|
|
|
|
def MobileViTv3_S(pretrained=False, use_ssld=False, **kwargs):
|
|
mv2_exp_mult = 4
|
|
mobilevit_config = {
|
|
"layer0": {
|
|
"img_channels": 3,
|
|
"out_channels": 16,
|
|
},
|
|
"layer1": {
|
|
"out_channels": 32,
|
|
"expand_ratio": mv2_exp_mult,
|
|
"num_blocks": 1,
|
|
"stride": 1,
|
|
"block_type": "mv2"
|
|
},
|
|
"layer2": {
|
|
"out_channels": 64,
|
|
"expand_ratio": mv2_exp_mult,
|
|
"num_blocks": 3,
|
|
"stride": 2,
|
|
"block_type": "mv2"
|
|
},
|
|
"layer3": { # 28x28
|
|
"out_channels": 128,
|
|
"transformer_channels": 144,
|
|
"ffn_dim": 288,
|
|
"transformer_blocks": 2,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"layer4": { # 14x14
|
|
"out_channels": 256,
|
|
"transformer_channels": 192,
|
|
"ffn_dim": 384,
|
|
"transformer_blocks": 4,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"layer5": { # 7x7
|
|
"out_channels": 320,
|
|
"transformer_channels": 240,
|
|
"ffn_dim": 480,
|
|
"transformer_blocks": 3,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"last_layer_exp_factor": 4
|
|
}
|
|
|
|
model = MobileViTv3(mobilevit_config, **kwargs)
|
|
|
|
_load_pretrained(
|
|
pretrained, model, MODEL_URLS["MobileViTv3_S"], use_ssld=use_ssld)
|
|
return model
|
|
|
|
|
|
def MobileViTv3_XS(pretrained=False, use_ssld=False, **kwargs):
|
|
mv2_exp_mult = 4
|
|
mobilevit_config = {
|
|
"layer0": {
|
|
"img_channels": 3,
|
|
"out_channels": 16,
|
|
},
|
|
"layer1": {
|
|
"out_channels": 32,
|
|
"expand_ratio": mv2_exp_mult,
|
|
"num_blocks": 1,
|
|
"stride": 1,
|
|
"block_type": "mv2"
|
|
},
|
|
"layer2": {
|
|
"out_channels": 48,
|
|
"expand_ratio": mv2_exp_mult,
|
|
"num_blocks": 3,
|
|
"stride": 2,
|
|
"block_type": "mv2"
|
|
},
|
|
"layer3": { # 28x28
|
|
"out_channels": 96,
|
|
"transformer_channels": 96,
|
|
"ffn_dim": 192,
|
|
"transformer_blocks": 2,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"layer4": { # 14x14
|
|
"out_channels": 160,
|
|
"transformer_channels": 120,
|
|
"ffn_dim": 240,
|
|
"transformer_blocks": 4,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"layer5": { # 7x7
|
|
"out_channels": 160,
|
|
"transformer_channels": 144,
|
|
"ffn_dim": 288,
|
|
"transformer_blocks": 3,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"last_layer_exp_factor": 4
|
|
}
|
|
|
|
model = MobileViTv3(mobilevit_config, **kwargs)
|
|
|
|
_load_pretrained(
|
|
pretrained, model, MODEL_URLS["MobileViTv3_XS"], use_ssld=use_ssld)
|
|
return model
|
|
|
|
|
|
def MobileViTv3_XXS(pretrained=False, use_ssld=False, **kwargs):
|
|
mv2_exp_mult = 2
|
|
mobilevit_config = {
|
|
"layer0": {
|
|
"img_channels": 3,
|
|
"out_channels": 16,
|
|
},
|
|
"layer1": {
|
|
"out_channels": 16,
|
|
"expand_ratio": mv2_exp_mult,
|
|
"num_blocks": 1,
|
|
"stride": 1,
|
|
"block_type": "mv2"
|
|
},
|
|
"layer2": {
|
|
"out_channels": 24,
|
|
"expand_ratio": mv2_exp_mult,
|
|
"num_blocks": 3,
|
|
"stride": 2,
|
|
"block_type": "mv2"
|
|
},
|
|
"layer3": { # 28x28
|
|
"out_channels": 64,
|
|
"transformer_channels": 64,
|
|
"ffn_dim": 128,
|
|
"transformer_blocks": 2,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"layer4": { # 14x14
|
|
"out_channels": 80,
|
|
"transformer_channels": 80,
|
|
"ffn_dim": 160,
|
|
"transformer_blocks": 4,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"layer5": { # 7x7
|
|
"out_channels": 128,
|
|
"transformer_channels": 96,
|
|
"ffn_dim": 192,
|
|
"transformer_blocks": 3,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": mv2_exp_mult,
|
|
"head_dim": None,
|
|
"num_heads": 4,
|
|
"block_type": "mobilevit"
|
|
},
|
|
"last_layer_exp_factor": 4
|
|
}
|
|
|
|
model = MobileViTv3(mobilevit_config, **kwargs)
|
|
|
|
_load_pretrained(
|
|
pretrained, model, MODEL_URLS["MobileViTv3_XXS"], use_ssld=use_ssld)
|
|
return model
|
|
|
|
|
|
def MobileViTv3_x1_0(pretrained=False, use_ssld=False, **kwargs):
|
|
mobilevit_config = {
|
|
"layer0": {
|
|
"img_channels": 3,
|
|
"out_channels": 32,
|
|
},
|
|
"layer1": {
|
|
"out_channels": 64,
|
|
"expand_ratio": 2,
|
|
"num_blocks": 1,
|
|
"stride": 1,
|
|
"block_type": "mv2",
|
|
},
|
|
"layer2": {
|
|
"out_channels": 128,
|
|
"expand_ratio": 2,
|
|
"num_blocks": 2,
|
|
"stride": 2,
|
|
"block_type": "mv2",
|
|
},
|
|
"layer3": { # 28x28
|
|
"out_channels": 256,
|
|
"attn_unit_dim": 128,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 2,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"layer4": { # 14x14
|
|
"out_channels": 384,
|
|
"attn_unit_dim": 192,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 4,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"layer5": { # 7x7
|
|
"out_channels": 512,
|
|
"attn_unit_dim": 256,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 3,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"last_layer_exp_factor": 4,
|
|
}
|
|
|
|
model = MobileViTv3(mobilevit_config, mobilevit_v2_based=True, **kwargs)
|
|
|
|
_load_pretrained(
|
|
pretrained, model, MODEL_URLS["MobileViTv3_x1_0"], use_ssld=use_ssld)
|
|
return model
|
|
|
|
|
|
def MobileViTv3_x0_75(pretrained=False, use_ssld=False, **kwargs):
|
|
mobilevit_config = {
|
|
"layer0": {
|
|
"img_channels": 3,
|
|
"out_channels": 24,
|
|
},
|
|
"layer1": {
|
|
"out_channels": 48,
|
|
"expand_ratio": 2,
|
|
"num_blocks": 1,
|
|
"stride": 1,
|
|
"block_type": "mv2",
|
|
},
|
|
"layer2": {
|
|
"out_channels": 96,
|
|
"expand_ratio": 2,
|
|
"num_blocks": 2,
|
|
"stride": 2,
|
|
"block_type": "mv2",
|
|
},
|
|
"layer3": { # 28x28
|
|
"out_channels": 192,
|
|
"attn_unit_dim": 96,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 2,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"layer4": { # 14x14
|
|
"out_channels": 288,
|
|
"attn_unit_dim": 144,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 4,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"layer5": { # 7x7
|
|
"out_channels": 384,
|
|
"attn_unit_dim": 192,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 3,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"last_layer_exp_factor": 4,
|
|
}
|
|
|
|
model = MobileViTv3(mobilevit_config, mobilevit_v2_based=True, **kwargs)
|
|
|
|
_load_pretrained(
|
|
pretrained, model, MODEL_URLS["MobileViTv3_x0_75"], use_ssld=use_ssld)
|
|
return model
|
|
|
|
|
|
def MobileViTv3_x0_5(pretrained=False, use_ssld=False, **kwargs):
|
|
mobilevit_config = {
|
|
"layer0": {
|
|
"img_channels": 3,
|
|
"out_channels": 16,
|
|
},
|
|
"layer1": {
|
|
"out_channels": 32,
|
|
"expand_ratio": 2,
|
|
"num_blocks": 1,
|
|
"stride": 1,
|
|
"block_type": "mv2",
|
|
},
|
|
"layer2": {
|
|
"out_channels": 64,
|
|
"expand_ratio": 2,
|
|
"num_blocks": 2,
|
|
"stride": 2,
|
|
"block_type": "mv2",
|
|
},
|
|
"layer3": { # 28x28
|
|
"out_channels": 128,
|
|
"attn_unit_dim": 64,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 2,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"layer4": { # 14x14
|
|
"out_channels": 192,
|
|
"attn_unit_dim": 96,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 4,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"layer5": { # 7x7
|
|
"out_channels": 256,
|
|
"attn_unit_dim": 128,
|
|
"ffn_multiplier": 2,
|
|
"attn_blocks": 3,
|
|
"patch_h": 2,
|
|
"patch_w": 2,
|
|
"stride": 2,
|
|
"mv_expand_ratio": 2,
|
|
"block_type": "mobilevit",
|
|
},
|
|
"last_layer_exp_factor": 4,
|
|
}
|
|
|
|
model = MobileViTv3(mobilevit_config, mobilevit_v2_based=True, **kwargs)
|
|
|
|
_load_pretrained(
|
|
pretrained, model, MODEL_URLS["MobileViTv3_x0_5"], use_ssld=use_ssld)
|
|
return model
|