Pytorch 实现 MobileNet V3 模型,并从 TensorFlow 转化预训练参数

        随着移动终端的普及,以及在其上运行深度学习模型的需求,神经网络小型化越来越得到重视和关注,已经成为研究的热门之一。作为小型化模型的经典代表,MobileNet 系列模型已经先后迭代了 3 代,在保持模型参数量和运算量都极其小的情况下,其性能越来越优异。本文我们将实现最新一代的 MobileNet V3,为了能不花费时间在 ImageNet 数据集上训练而直接使用,我们将从 TensorFlow 官方实现的 MobileNet V3 上转化预训练参数。

        本文将重点关注以下两个方面:

  • 详细解读 MobileNet V3 的网络结构;
  • 详细讲述从 TensorFlow 转化预训练参数的方法;

        本文所有代码见 GitHub: mobilenet_v3

一、MobileNet V3 模型

二、模型实现

三、预训练参数转化

        完全采用手动指定的方式进行,即对于 Pytorch 模型的每一参数,从对应的 TensorFlow 预训练参数里取出,然后赋值给它即可。为了保证转化的准确性,我们的目标是:

  • 原 TensorFlow 预训练模型和转化后的 Pytorch 模型的预测结果要绝对一致

以下,我们详细的来描述怎么从 TensorFlow 转化预训练参数。

1.查看 TensorFlow 预训练模型参数名

        首先到 此页 下载 MobileNet V3 模型的 TensorFlow 预训练模型,下载后请解压。我们以 large dm=1 (float) 预训练模型为例来说明。首先,使用如下代码:

import json
import tensorflow as tf

if __name__ == '__main__':
    checkpoint_path = 'xxx/v3-large_224_1.0_float/ema/model-540000'
    output_path = './mobilenet_v3_large.json'

    reader = tf.train.NewCheckpointReader(checkpoint_path)
    weights = {var: 1 for (var, _) in
               reader.get_variable_to_shape_map().items()}
    
    with open(output_path, 'w') as writer:
        json.dump(weights, writer)

预训练模型中的所有参数名都写到一个 json 文件里,为了不把高维的数据写进去,我们都将确切的值改成了 1。但直接写进去的内容很乱,可以借助 json 串格式化的工具(比如,在线格式化,或者 Google Chrome 浏览器插件 FeHelper)将 mobilenet_v3_large.json 文件里的内容格式化,这样你看到的形式就大概如下了:

格式化之后的 mobilenet_v3_large.json 内容

接着,结合 TensorFlow 官方开源的 MobileNet V3 Large 模型的网络定义

MobileNet V3 large 模型 TensorFlow 网络定义

就基本可以知道整个模型参数命名的具体名字和顺序了:

MobilenetV3/Conv/
MobilenetV3/expanded_conv/
MobilenetV3/expanded_conv_1/
...
MobilenetV3/expanded_conv_14/
MobilenetV3/Conv_1/
MobilenetV3/Conv_2/
MobilenetV3/Logits/Conv2d_1c_1x1

以上是 large 模型的总共 19 个大的命名空间(scope),每个 / 之后会接小的命名空间。对于普通的卷积层,比如 Conv, Conv_1, Conv_2, Logits/Conv2d_1c_1x1 你要关注两个东西:

  • 是否有偏置参数:biases
  • 是否有批标准化:BatchNorm

这既可以帮助你修正你定义的 Pytorch 模型,也可以在转化赋值的时候防止被遗忘。类似的思想可以直接移植到复杂的模块 mbv3_op 对应的命名空间,expanded_conv, expanded_conv_1, ...。举个简单的例子,看 large 模型的第一卷积层:MobilenetV3/Conv/,因为该层使用了批标准化(batch normalization),因此是没有偏置参数的,那么就只有如下的 5 个参数:

MobilenetV3/Conv/weights,
MobilenetV3/BatchNorm/beta,
MobilenetV3/Conv/BatchNorm/gamma
MobilenetV3/Conv/BatchNorm/moving_mean
MobilenetV3/Conv/BatchNorm/moving_variance

其中后 4 个参数对应于批标准化的公式:
\gamma \frac{x - \mu}{\sigma} + \beta. \\
再看 large 模型的最后一个卷积层(分类层):MobilenetV3/Logits/Conv2d_1c_1x1,因为该层没有使用批标准化的正规化函数,因此带有偏置项,就只有两个参数:

MobilenetV3/Logits/Conv2d_1c_1x1/weights
MobilenetV3/Logits/Conv2d_1c_1x1/biases

至于其他复杂模块,分割开单独考虑中间命名空间: project, expand, depthwise, squeeze_excite 之后,其实就是简单的卷积层了,因此也很容易处理。

2.查看 Pytorch 模型结构

        这一步更容易,直接实例化定义的 Pytorch 模型,然后打印出来(这里,模型的所有的层都定义在了属性 _layers 里,见 mobilenet_v3.MobileNet 类):

import mobilenet_v3

large = mobilenet_v3.large()
print(large._layers[:10])
print(large._layers[10:])

因为模型结构很长,所以打印的时候分成了前后两部分。保存在 txt 文件里如下:

MobileNet V3 large 模型网络结构

这一步我们唯一需要关注的就是每一层在网络结构里的下标了,比如 _layers[0] 就是整个网络的第 1 个卷积层模块,而 _layers[0]._layers[0] 是这个模块内的二维卷积层,_layers[0]._layers[1] 是这个模块内的批标准化层。因为 torch.nn.Sequential 的行为和 list 一样,因此它们的顺序是确定不变的,取下标是非常安全的操作。

3.对照参数名逐一赋值

        经过前面两步之后,应该对 TensorFlow 预训练模型 和 Pytorch 定义的模型结构 之间的对应关系应该有所印象了,下面需要将它们严格的对应起来,以便预训练参数转化。

        首先,看第一个卷积模块,它包含一个卷积层、批标准化层和一个激活函数层,其中只有前两者是有训练参数的。而且,根据第一步,我们知道对应的 TensorFlow 模型这一个模块的命名空间是:MobilenetV3/Conv/,因此如果我声明了

import mobilenet_v3

model = mobilenet_v3.large()

large 模型,那么对应的第 1 个卷积模块的二维卷积层是 model._layers[0]._layers[0],批标准化层是 model._layers[0]._layers[1]。它们所含有的参数如下:

model._layers[0]._layers[0].weight
model._layers[0]._layers[1].bias:
model._layers[0]._layers[1].weight
model._layers[0]._layers[1].running_mean
model._layers[0]._layers[1].running_var

即卷积层的权重参数(对于 slim.conv2d(),如果指定了正规化函数,即关键字参数 normalizer_fn 不为 None,那么这个卷积层是没有偏置项的;反之,则有,除非将偏置的初始化函数 biases_initializer 设为 None),和批标准化层的 4 个参数:
\gamma \frac{x - \mu}{\sigma} + \beta. \\
很容易的,你可以从 mobilenet_v3_large.json 里找到对应的 TensorFlow 变量名:

conversion_map_for_root_block = {
    model._layers[0]._layers[0].weight: 
        'MobilenetV3/Conv/weights',
    model._layers[0]._layers[1].bias: 
        'MobilenetV3/Conv/BatchNorm/beta',
    model._layers[0]._layers[1].weight:
        'MobilenetV3/Conv/BatchNorm/gamma',
    model._layers[0]._layers[1].running_mean: 
        'MobilenetV3/Conv/BatchNorm/moving_mean',
    model._layers[0]._layers[1].running_var: 
        'MobilenetV3/Conv/BatchNorm/moving_variance',
}

然后用函数 tf.train.load_variable,按照 TensorFlow 的变量名从预训练模型中取出变量的名字赋值给对应的 Pytorch 变量,比如:

checkpoint_path = 'xxx/v3-large_224_1.0_float/ema/model-540000'

tf_param = tf.train.load_variable(checkpoint_path, 'MobilenetV3/Conv/weights')
tf_param = np.transpose(tf_param, (3, 2, 0, 1))
model._layers[0]._layers[0].weight.data = torch.from_numpy(tf_param)

就将第 1 个卷积层的参数转化好了。这里,唯一需要注意的是,TensorFlow 权重的顺序是 [kernel_size, kernel_size, in_channels, out_channels],而 Pytorch 的顺序是 [out_channels, in_channels, kernel_size, kernel_size],因此要将它们的顺序调整到一致。

        其它参数完全按照一样的方式转化即可。完整的转化代码请见 converter.py

        以上过程结束之后,我们来转化几个模型

1.large 模型

        执行(tf_checkpoint_path 参数指定 TensorFlow 预训练模型参数的保存路径):

python3 tf_weights_to_pth.py --tf_checkpoint_path xxx/v3-large_224_1.0_float/ema/model-540000

将在当前项目路径下生成一个 pretrained_models 文件夹,里面保存了转化后的模型:mobilenet_v3_large.pth,同时将输出测试图片(熊猫图片):

panda.jpg

的分类结果:

large 模型 TensorFlow 原预训练模型和转化的 Pytorch 模型对熊猫图片的识别结果

可以看到两者的结果是一模一样的。类似的,再指定另一张测试图片(猫图片),执行以下命令(image_path 参数指定测试图片的路径):

python3 tf_weights_to_pth.py --tf_checkpoint_path xxx/v3-large_224_1.0_float/ema/model-540000 \
    --image_path ./test/cat.jpg
cat.jpg

就可以看到对猫的分类结果:

large 模型 TensorFlow 原预训练模型和转化的 Pytorch 模型对猫图片的识别结果

显然,TensorFlow 官方和本文实现的 Pytorch 模型的预测结果也是一模一样的

2.small 模型(depth_multiplier = 0.75)

执行(output_name 指定转化来的模型的保存名字,depth_multiplier 指定卷积层的通道数乘子,model_name 指定转化的模型名):

python3 tf_weights_to_pth.py --tf_checkpoint_path xxx/v3-small_224_0.75_float/ema/model-497500 \
    --output_name mobilenet_v3_small_0.75.pth --depth_multiplier 0.75 --model_name small

得到熊猫图片的分类结果:

small-dm=0.75 模型 TensorFlow 原预训练模型和转化的 Pytorch 模型对熊猫图片的识别结果

也得到一模一样的结果,说明转化参数是正确的。

        当前支持参数转化的预训练模型如下:

本文所有支持参数转化的预训练模型

对应的模型名(由 model_name 参数指定)分别为:large, small, large_minimalistic, small_minimalistic,如果 dm=0.75,请指定参数 depth_multiplier。你可以逐一转化并验证本文定义的 MobileNet V3 模型的正确性,不出意外应该是准确的(作者未转化 8-bit 的预训练模型)。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 203,456评论 5 477
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,370评论 2 381
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 150,337评论 0 337
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,583评论 1 273
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,596评论 5 365
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,572评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,936评论 3 395
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,595评论 0 258
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,850评论 1 297
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,601评论 2 321
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,685评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,371评论 4 318
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,951评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,934评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,167评论 1 259
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 43,636评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,411评论 2 342

推荐阅读更多精彩内容