論文：MobileOne: 改良された 1 ミリ秒モバイル バックボーン
正式実装：https://github.com/apple/ml-mobileone
サードパーティ実装：メインの mmpretrain/mobileone.py · open-mmlab/mmpretrain · GitHub
序文
モバイル デバイス向けの効率的な深層学習アーキテクチャの設計と展開は大きく進歩しており、多くの軽量モデルは浮動小数点演算 (FLOPS) とパラメータ数を削減しながら精度を向上させ続けています。ただし、レイテンシの観点からは、これらの指標はモデルの効率とあまり相関しておらず、FLOP などの指標では、メモリ アクセス コストやモデルの並列度が考慮されておらず、推論中に失われる可能性があります。遅延に大きな影響を与えます。パラメーターの量もレイテンシーとあまり相関しません。たとえば、パラメーターを共有するとモデルのサイズは小さくなりますが、FLOP が増加する可能性があります。さらに、スキップ接続や分岐などの小さなパラメータを使用した操作では、大量のメモリ アクセス コストが発生します。
この記事の目的は、レイテンシに影響を与える主要なボトルネックを見つけて最適化し、モデルの精度を向上させながらレイテンシ コストを削減することです。特に小規模なネットワークをトレーニングする場合、最適化がもう 1 つのボトルネックになります。これは、トレーニングと推論のためのネットワーク構造を分離する、つまり構造再パラメーター化手法によって軽減できます。さらに、最適化のボトルネックは、トレーニング中に正則化を動的に遅くすることでさらに軽減され、すでに小さいモデルが過度に正則化されるのを防ぎます。
この論文では、発見された構造と最適化のボトルネックに基づいて、新しいネットワーク構造 MobileOne を設計します。そのバリアントは、iPhone12 で 1ms 未満の遅延で SOTA 精度を達成します。MobileOne と以前の構造再パラメータ化モデルの主な違いは、過剰パラメータ化ブランチとモデル スケーリング戦略の導入です。
この記事への寄稿
- この論文では、モバイル デバイス上で 1 ミリ秒未満の遅延内で効率的なモデル アーキテクチャで最先端の画像分類精度を達成する新しいアーキテクチャである MobileOne を設計します。このモデルのパフォーマンスはデスクトップ CPU にも一般化されます。
- このペーパーでは、モバイル デバイスでの高い遅延コストにつながる、アクティベーションと分岐におけるパフォーマンスのボトルネックを分析します。
- このペーパーでは、構造の再パラメーター化と動的緩和の正則化の影響を分析します。これらを組み合わせることで、小規模モデルをトレーニングするときに発生する最適化のボトルネックを軽減できます。
- この論文で設計された MobileOne は、ターゲット検出やセマンティック セグメンテーションなどの他のダウンストリーム タスクにも十分に拡張でき、そのパフォーマンスは以前の SOTA モデルよりも優れています。
メソッドの紹介
メトリックの相関関係
著者はまず、一般的に使用される 2 つの指標 FLOP およびパラメータ量とモバイルデバイスの遅延の相関関係を分析し、いくつかの一般的な軽量モデルの FLOP およびパラメータ量と iPhone12 の遅延の関係を以下の図に示します。 FLOP が少ないことやパラメータ数が少ないことは、実際のレイテンシが低いことを意味するわけではありません。
主要なボトルネック
著者は遅延のボトルネックをネットワーク構造から分析しており、その第一は活性化関数である。遅延に対するアクティベーション関数の影響を分析するために、著者は 30 層のニューラル ネットワークを構築し、さまざまなアクティベーション関数を使用して iPhone 12 でベンチマーク テストを実行しました。結果を表 3 に示します。 ReLUの遅延が最も少ないとのことで、筆者はMobileOneではReLUアクティベーション機能のみを使用しています。
著者は、ブロック構造が遅延に与える影響も分析しました。レイテンシーに影響を与える 2 つの重要な要素は、メモリ アクセス コストとモデルの並列処理です。マルチブランチ構造では、グラフ内の次のテンソルを計算するために各ブランチのアクティブ化を保存する必要があるため、メモリ アクセス コストが大幅に増加します。ネットワークの分岐数が少なければ、このメモリのボトルネックを回避できます。さらに、SE ブロック内のグローバル平均プーリング操作などの同期操作を実行する必要があるブロック構造の場合、同期コストも遅延に影響します。そこで著者は推論中に分岐を持たない構造を採用し、メモリアクセスのコストを削減しました。さらに、精度を向上させるために、SE ブロックは MobileOne の最大のバリアントでのみ使用されます。
MobileOne アーキテクチャ
モバイルワンブロック
MobileOne ブロックの構造は深さ方向の層と点方向の層に分解され、オーバーパラメータ化分岐も導入されます。基本ブロックの構造は、MobileNet-V1 の設計に基づいており、3x3 の深さ方向の畳み込みに続いて 1x1 の点方向の畳み込みが続き、その後、バッチノルム ブランチと複数回複製できるブランチを含む再パラメータ化可能なスキップ接続ブランチが追加されました。図 3 に示すように
このうち、過剰パラメータ化係数 \(k\) はハイパーパラメータであり、推論段階では構造の再パラメータ化によって複数の分岐がマージされ、単一分岐の構造のみが保持されます。
モデルのスケーリング
最近の作品の中には、スケール モデルの幅、深さ、解像度を拡大縮小することでパフォーマンスを向上させるものもあります。MobileOne の深度スケーリング戦略は MobileNet-V2 の戦略と似ています。つまり、浅い初期段階では使用するレイヤーの数が少なくなります。これは、浅いレイヤーの特徴マップ解像度が大きく、解像度が小さい深層レイヤーよりもはるかに遅いためです。幅については、表 2 に示すように、著者は 5 つの異なるスケールを採用しています。
解像度については、モバイル デバイスでの実行パフォーマンスに役立たないため、作成者はスケールしませんでした。
トレーニング
小規模なモデルでは、過学習に対処するために必要な正則化が大規模なモデルよりも少なくなります。学習率については、著者はコサイン スケジュール戦略を採用し、重み減衰係数にも同じ戦略を使用します。さらに、著者はプログレッシブ学習カリキュラムのアイデアも採用しており、表 5 はさまざまなトレーニング設定による精度の向上を示しています。
実験結果
筆者は、ImageNet-1K 上で MobileOne の分類効果を評価しています｡モデルの具体的な内容は以下の通りです: ラベル平滑正則化を採用し、平滑化係数を 0.1 に設定し、初期学習率を 0.1 に設定し、コサインアニーリングを行いますスケジュールは学習率を調整するために使用されます。重み減衰係数は \(10^{-4}\) に設定され、コサイン スケジュールも \(10^{-5}\) に削減されます。MobileOne のより大きなバリアント、つまり S2、S3 をトレーニングする場合のみです。 、S4 AutoAugment データ強化が採用され、自動拡張の強度と入力解像度は EfficientNetv2 の方法が採用され、トレーニング プロセス中に徐々に増加します。より小さいバリアント、つまり S0 と S1 の場合は、基本的なデータ強化方法であるランダム スケーリングとトリミングと水平反転が採用され、すべてのバリアントで EMA (指数移動平均) が使用され、減衰定数は 0.9995 に設定されます。実験結果は以下の通りです
結果は遅延ごとにグループ化されており、MobileOne のさまざまなバージョンが最小の遅延で最高の精度を達成していることがわかります。
また、著者は MobileOne をバックボーンとして使用して、ターゲット検出タスクとセマンティック セグメンテーション タスクのパフォーマンスを比較しました。結果は次のとおりです。
バックボーンとしての MobileOne は、優れた汎化パフォーマンスに加えて、ターゲット検出タスクとセマンティック セグメンテーション タスクの両方で他の軽量モデルよりも優れていることがわかります。
コード分析
ここでは mmdetection での実装を例として、S0 ～ S4 の仕様を表 2 と併せて示します。このうち、num_blocks は各ステージのブロック数であり、表 2 の stage1、stage7、8 は含まれず、stage4、5 はマージされるため、リストの要素は 4 つだけになります。width_factor は幅、つまりチャネルのスケーリング係数を表します。これは、表 2 の \(\alpha\) です。num_conv_branches は、ブロック内の過剰パラメータ化分岐の繰り返し数、つまり表 2 の \(k\) を示します。num_se_blocksは、各ステージのブロックで使用されるSEレイヤの数を示す。これら 4 つのリストの要素の数は 4 で、それぞれ表 2 の stage2、stage3、stage45、および stage6 に対応します。
arch_zoo = {
's0':
dict(
num_blocks=[2, 8, 10, 1],
width_factor=[0.75, 1.0, 1.0, 2.0],
num_conv_branches=[4, 4, 4, 4],
num_se_blocks=[0, 0, 0, 0]),
's1':
dict(
num_blocks=[2, 8, 10, 1],
width_factor=[1.5, 1.5, 2.0, 2.5],
num_conv_branches=[1, 1, 1, 1],
num_se_blocks=[0, 0, 0, 0]),
's2':
dict(
num_blocks=[2, 8, 10, 1],
width_factor=[1.5, 2.0, 2.5, 4.0],
num_conv_branches=[1, 1, 1, 1],
num_se_blocks=[0, 0, 0, 0]),
's3':
dict(
num_blocks=[2, 8, 10, 1],
width_factor=[2.0, 2.5, 3.0, 4.0],
num_conv_branches=[1, 1, 1, 1],
num_se_blocks=[0, 0, 0, 0]),
's4':
dict(
num_blocks=[2, 8, 10, 1],
width_factor=[3.0, 3.5, 3.5, 4.0],
num_conv_branches=[1, 1, 1, 1],
num_se_blocks=[0, 0, 5, 1])
}MobileOneBlockのコードは次のとおりです。図 3 に示すように、完全なブロックには 3x3 の深さ方向のブロックと 1x1 の点方向のブロックが含まれており、以下の実装は単なる 1 つのブロックであるため、2 回呼び出す必要があることに注意してください。
forward 関数を見てください。self.branch_normは BN ブランチであり、stride=1 で入力チャネルと出力チャネルの数が同じ場合にのみ存在します。self.branch_scaleは 3x3 深さ方向ブロック内の 1x1 分岐ですが、1x1 点方向ブロックにはそのような分岐はありません。self.branch_conv_listは \(k\) 回繰り返すオーバーパラメータ化ブランチ、つまり 3x3 深さ方向の変換または 1x1 の点方向の変換です。
class MobileOneBlock(BaseModule):
"""MobileOne block for MobileOne backbone.
Args:
in_channels (int): The input channels of the block.
out_channels (int): The output channels of the block.
kernel_size (int): The kernel size of the convs in the block. If the
kernel size is large than 1, there will be a ``branch_scale`` in
the block.
num_convs (int): Number of the convolution branches in the block.
stride (int): Stride of convolution layers. Defaults to 1.
padding (int): Padding of the convolution layers. Defaults to 1.
dilation (int): Dilation of the convolution layers. Defaults to 1.
groups (int): Groups of the convolution layers. Defaults to 1.
se_cfg (None or dict): The configuration of the se module.
Defaults to None.
norm_cfg (dict): Configuration to construct and config norm layer.
Defaults to ``dict(type='BN')``.
act_cfg (dict): Config dict for activation layer.
Defaults to ``dict(type='ReLU')``.
deploy (bool): Whether the model structure is in the deployment mode.
Defaults to False.
init_cfg (dict or list[dict], optional): Initialization config dict.
Defaults to None.
"""
def __init__(self,
in_channels: int,
out_channels: int,
kernel_size: int,
num_convs: int,
stride: int = 1,
padding: int = 1,
dilation: int = 1,
groups: int = 1,
se_cfg: Optional[dict] = None,
conv_cfg: Optional[dict] = None,
norm_cfg: Optional[dict] = dict(type='BN'),
act_cfg: Optional[dict] = dict(type='ReLU'),
deploy: bool = False,
init_cfg: Optional[dict] = None):
super(MobileOneBlock, self).__init__(init_cfg)
assert se_cfg is None or isinstance(se_cfg, dict)
if se_cfg is not None:
self.se = SELayer(channels=out_channels, **se_cfg)
else:
self.se = nn.Identity()
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.num_conv_branches = num_convs
self.stride = stride
self.padding = padding
self.se_cfg = se_cfg
self.conv_cfg = conv_cfg
self.norm_cfg = norm_cfg
self.act_cfg = act_cfg
self.deploy = deploy
self.groups = groups
self.dilation = dilation
if deploy:
self.branch_reparam = build_conv_layer(
conv_cfg,
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
groups=self.groups,
stride=stride,
padding=padding,
dilation=dilation,
bias=True)
else:
# judge if input shape and output shape are the same.
# If true, add a normalized identity shortcut.
if out_channels == in_channels and stride == 1:
self.branch_norm = build_norm_layer(norm_cfg, in_channels)[1]
else:
self.branch_norm = None
self.branch_scale = None
if kernel_size > 1:
self.branch_scale = self.create_conv_bn(kernel_size=1)
self.branch_conv_list = ModuleList()
for _ in range(num_convs):
self.branch_conv_list.append(
self.create_conv_bn(
kernel_size=kernel_size,
padding=padding,
dilation=dilation))
self.act = build_activation_layer(act_cfg)
def create_conv_bn(self, kernel_size, dilation=1, padding=0):
"""cearte a (conv + bn) Sequential layer."""
conv_bn = Sequential()
conv_bn.add_module(
'conv',
build_conv_layer(
self.conv_cfg,
in_channels=self.in_channels,
out_channels=self.out_channels,
kernel_size=kernel_size,
groups=self.groups,
stride=self.stride,
dilation=dilation,
padding=padding,
bias=False))
conv_bn.add_module(
'norm',
build_norm_layer(self.norm_cfg, num_features=self.out_channels)[1])
return conv_bn
def forward(self, x):
def _inner_forward(inputs):
if self.deploy:
return self.branch_reparam(inputs)
inner_out = 0
if self.branch_norm is not None:
inner_out = self.branch_norm(inputs)
if self.branch_scale is not None:
inner_out += self.branch_scale(inputs)
for branch_conv in self.branch_conv_list:
inner_out += branch_conv(inputs)
return inner_out
return self.act(self.se(_inner_forward(x)))
def switch_to_deploy(self):
"""Switch the model structure from training mode to deployment mode."""
if self.deploy:
return
assert self.norm_cfg['type'] == 'BN', \
"Switch is not allowed when norm_cfg['type'] != 'BN'."
reparam_weight, reparam_bias = self.reparameterize()
self.branch_reparam = build_conv_layer(
self.conv_cfg,
self.in_channels,
self.out_channels,
kernel_size=self.kernel_size,
stride=self.stride,
padding=self.padding,
dilation=self.dilation,
groups=self.groups,
bias=True)
self.branch_reparam.weight.data = reparam_weight
self.branch_reparam.bias.data = reparam_bias
for param in self.parameters():
param.detach_()
delattr(self, 'branch_conv_list')
if hasattr(self, 'branch_scale'):
delattr(self, 'branch_scale')
delattr(self, 'branch_norm')
self.deploy = True
def reparameterize(self):
"""Fuse all the parameters of all branches.
Returns:
tuple[torch.Tensor, torch.Tensor]: Parameters after fusion of all
branches. the first element is the weights and the second is
the bias.
"""
weight_conv, bias_conv = 0, 0
for branch_conv in self.branch_conv_list:
weight, bias = self._fuse_conv_bn(branch_conv)
weight_conv += weight
bias_conv += bias
weight_scale, bias_scale = 0, 0
if self.branch_scale is not None:
weight_scale, bias_scale = self._fuse_conv_bn(self.branch_scale)
# Pad scale branch kernel to match conv branch kernel size.
pad = self.kernel_size // 2
weight_scale = F.pad(weight_scale, [pad, pad, pad, pad])
weight_norm, bias_norm = 0, 0
if self.branch_norm:
tmp_conv_bn = self._norm_to_conv(self.branch_norm)
weight_norm, bias_norm = self._fuse_conv_bn(tmp_conv_bn)
return (weight_conv + weight_scale + weight_norm,
bias_conv + bias_scale + bias_norm)
def _fuse_conv_bn(self, branch):
"""Fuse the parameters in a branch with a conv and bn.
Args:
branch (mmcv.runner.Sequential): A branch with conv and bn.
Returns:
tuple[torch.Tensor, torch.Tensor]: The parameters obtained after
fusing the parameters of conv and bn in one branch.
The first element is the weight and the second is the bias.
"""
if branch is None:
return 0, 0
kernel = branch.conv.weight
running_mean = branch.norm.running_mean
running_var = branch.norm.running_var
gamma = branch.norm.weight
beta = branch.norm.bias
eps = branch.norm.eps
std = (running_var + eps).sqrt()
fused_weight = (gamma / std).reshape(-1, 1, 1, 1) * kernel
fused_bias = beta - running_mean * gamma / std
return fused_weight, fused_bias
def _norm_to_conv(self, branch_nrom):
"""Convert a norm layer to a conv-bn sequence towards
``self.kernel_size``.
Args:
branch (nn.BatchNorm2d): A branch only with bn in the block.
Returns:
(mmcv.runner.Sequential): a sequential with conv and bn.
"""
input_dim = self.in_channels // self.groups
conv_weight = torch.zeros(
(self.in_channels, input_dim, self.kernel_size, self.kernel_size),
dtype=branch_nrom.weight.dtype)
for i in range(self.in_channels):
conv_weight[i, i % input_dim, self.kernel_size // 2,
self.kernel_size // 2] = 1
conv_weight = conv_weight.to(branch_nrom.weight.device)
tmp_conv = self.create_conv_bn(kernel_size=self.kernel_size)
tmp_conv.conv.weight.data = conv_weight
tmp_conv.norm = branch_nrom
return tmp_conv
表 2 の stage1 のコードは次のとおりです. stage1 にはブロックが 1 つだけあり、深さ方向のブロックが 3x3 のみで、次の 1x1 点方向のブロックがないことに注意してください。
self.stage0 = MobileOneBlock(
self.in_channels,
channels,
stride=2,
kernel_size=3,
num_convs=1,
conv_cfg=conv_cfg,
norm_cfg=norm_cfg,
act_cfg=act_cfg,
deploy=deploy)次のステップでは、Arch_zooで指定されたバージョン パラメータに従って stage2 ～ stage6 を移動します。
self.stages = []
for i, num_blocks in enumerate(self.arch['num_blocks']):
planes = int(base_channels[i] * self.arch['width_factor'][i])
stage = self._make_stage(planes, num_blocks,
arch['num_se_blocks'][i],
arch['num_conv_branches'][i])
stage_name = f'stage{i + 1}'
self.add_module(stage_name, stage)
self.stages.append(stage_name)このうちself._make_stageは各ステージをビルドするもので、コードは以下の通りです。stage1 とは異なり、ここでは各ブロックに深さ方向のブロックと点方向のブロックが含まれていることがわかります。stride=2の場合、各ステージの最初のブロックの深さ方向のブロックでダウンサンプリングが行われます。
def _make_stage(self, planes, num_blocks, num_se, num_conv_branches):
strides = [2] + [1] * (num_blocks - 1)
if num_se > num_blocks:
raise ValueError('Number of SE blocks cannot '
'exceed number of layers.')
blocks = []
for i in range(num_blocks):
use_se = False
if i >= (num_blocks - num_se):
use_se = True
blocks.append(
# Depthwise conv
MobileOneBlock(
in_channels=self.in_planes,
out_channels=self.in_planes,
kernel_size=3,
num_convs=num_conv_branches,
stride=strides[i],
padding=1,
groups=self.in_planes,
se_cfg=self.se_cfg if use_se else None,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg,
deploy=self.deploy))
blocks.append(
# Pointwise conv
MobileOneBlock(
in_channels=self.in_planes,
out_channels=planes,
kernel_size=1,
num_convs=num_conv_branches,
stride=1,
padding=0,
se_cfg=self.se_cfg if use_se else None,
conv_cfg=self.conv_cfg,
norm_cfg=self.norm_cfg,
act_cfg=self.act_cfg,
deploy=self.deploy))
self.in_planes = planes
return Sequential(*blocks)