随着移动终端的普及,以及在其上运行深度学习模型的需求,神经网络小型化越来越得到重视和关注,已经成为研究的热门之一。作为小型化模型的经典代表,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 个参数对应于批标准化的公式:
再看 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 个参数:
很容易的,你可以从 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 的预训练模型)。
网友评论