Pytorch转tflite方式


Posted in Python onMay 25, 2020

目标是想把在服务器上用pytorch训练好的模型转换为可以在移动端运行的tflite模型。

最直接的思路是想把pytorch模型转换为tensorflow的模型,然后转换为tflite。但是这个转换目前没有发现比较靠谱的方法。

经过调研发现最新的tflite已经支持直接从keras模型的转换,所以可以采用keras作为中间转换的桥梁,这样就能充分利用keras高层API的便利性。

转换的基本思想就是用pytorch中的各层网络的权重取出来后直接赋值给keras网络中的对应layer层的权重。

转换为Keras模型后,再通过tf.contrib.lite.TocoConverter把模型直接转为tflite.

下面是一个例子,假设转换的是一个两层的CNN网络。

import tensorflow as tf
from tensorflow import keras
import numpy as np

import torch
from torchvision import models
import torch.nn as nn
# import torch.nn.functional as F
from torch.autograd import Variable

class PytorchNet(nn.Module):
 def __init__(self):
 super(PytorchNet, self).__init__()
 conv1 = nn.Sequential(
  nn.Conv2d(3, 32, 3, 2),
  nn.BatchNorm2d(32),
  nn.ReLU(inplace=True),
  nn.MaxPool2d(2, 2))
 conv2 = nn.Sequential(
  nn.Conv2d(32, 64, 3, 1, groups=1),
  nn.BatchNorm2d(64),
  nn.ReLU(inplace=True),
  nn.MaxPool2d(2, 2))
 self.feature = nn.Sequential(conv1, conv2)
 self.init_weights()

 def forward(self, x):
 return self.feature(x)

 def init_weights(self):
 for m in self.modules():
  if isinstance(m, nn.Conv2d):
  nn.init.kaiming_normal_(
   m.weight.data, mode='fan_out', nonlinearity='relu')
  if m.bias is not None:
   m.bias.data.zero_()
  if isinstance(m, nn.BatchNorm2d):
  m.weight.data.fill_(1)
  m.bias.data.zero_()

def KerasNet(input_shape=(224, 224, 3)):
 image_input = keras.layers.Input(shape=input_shape)
 # conv1
 network = keras.layers.Conv2D(
 32, (3, 3), strides=(2, 2), padding="valid")(image_input)
 network = keras.layers.BatchNormalization(
 trainable=False, fused=False)(network)
 network = keras.layers.Activation("relu")(network)
 network = keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))(network)

 # conv2
 network = keras.layers.Conv2D(
 64, (3, 3), strides=(1, 1), padding="valid")(network)
 network = keras.layers.BatchNormalization(
 trainable=False, fused=True)(network)
 network = keras.layers.Activation("relu")(network)
 network = keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))(network)

 model = keras.Model(inputs=image_input, outputs=network)

 return model

class PytorchToKeras(object):
 def __init__(self, pModel, kModel):
 super(PytorchToKeras, self)
 self.__source_layers = []
 self.__target_layers = []
 self.pModel = pModel
 self.kModel = kModel
 tf.keras.backend.set_learning_phase(0)

 def __retrieve_k_layers(self):
 for i, layer in enumerate(self.kModel.layers):
  if len(layer.weights) > 0:
  self.__target_layers.append(i)

 def __retrieve_p_layers(self, input_size):

 input = torch.randn(input_size)
 input = Variable(input.unsqueeze(0))
 hooks = []

 def add_hooks(module):

  def hook(module, input, output):
  if hasattr(module, "weight"):
   # print(module)
   self.__source_layers.append(module)

  if not isinstance(module, nn.ModuleList) and not isinstance(module, nn.Sequential) and module != self.pModel:
  hooks.append(module.register_forward_hook(hook))

 self.pModel.apply(add_hooks)

 self.pModel(input)
 for hook in hooks:
  hook.remove()

 def convert(self, input_size):
 self.__retrieve_k_layers()
 self.__retrieve_p_layers(input_size)

 for i, (source_layer, target_layer) in enumerate(zip(self.__source_layers, self.__target_layers)):
  print(source_layer)
  weight_size = len(source_layer.weight.data.size())
  transpose_dims = []
  for i in range(weight_size):
  transpose_dims.append(weight_size - i - 1)
  if isinstance(source_layer, nn.Conv2d):
  transpose_dims = [2,3,1,0]
  self.kModel.layers[target_layer].set_weights([source_layer.weight.data.numpy(
  ).transpose(transpose_dims), source_layer.bias.data.numpy()])
  elif isinstance(source_layer, nn.BatchNorm2d):
  self.kModel.layers[target_layer].set_weights([source_layer.weight.data.numpy(), source_layer.bias.data.numpy(),
        source_layer.running_mean.data.numpy(), source_layer.running_var.data.numpy()])
 def save_model(self, output_file):
 self.kModel.save(output_file)

 def save_weights(self, output_file):
 self.kModel.save_weights(output_file, save_format='h5')

pytorch_model = PytorchNet()
keras_model = KerasNet(input_shape=(224, 224, 3))

torch.save(pytorch_model, 'test.pth')

#Load the pretrained model
pytorch_model = torch.load('test.pth')

# #Time to transfer weights
converter = PytorchToKeras(pytorch_model, keras_model)
converter.convert((3, 224, 224))

# #Save the converted keras model for later use
# converter.save_weights("keras.h5")
converter.save_model("keras_model.h5")

# convert keras model to tflite model
converter = tf.contrib.lite.TocoConverter.from_keras_model_file(
 "keras_model.h5")
tflite_model = converter.convert()
open("convert_model.tflite", "wb").write(tflite_model)

补充知识:tensorflow模型转换成tensorflow lite模型

1.把graph和网络模型打包在一个文件中

bazel build tensorflow/python/tools:freeze_graph && \
 bazel-bin/tensorflow/python/tools/freeze_graph \
 --input_graph=eval_graph_def.pb \
 --input_checkpoint=checkpoint \
 --output_graph=frozen_eval_graph.pb \
 --output_node_names=outputs

For example:

bazel-bin/tensorflow/python/tools/freeze_graph \ 
 --input_graph=./mobilenet_v1_1.0_224/mobilenet_v1_1.0_224_eval.pbtxt \
 --input_checkpoint=./mobilenet_v1_1.0_224/mobilenet_v1_1.0_224.ckpt \
 --output_graph=./mobilenet_v1_1.0_224/frozen_eval_graph_test.pb \
 --output_node_names=MobilenetV1/Predictions/Reshape_1

2.把第一步中生成的tensorflow pb模型转换为tf lite模型

转换前需要先编译转换工具

bazel build tensorflow/contrib/lite/toco:toco

转换分两种,一种的转换为float的tf lite,另一种可以转换为对模型进行unit8的量化版本的模型。两种方式如下:

非量化的转换:

./bazel-bin/third_party/tensorflow/contrib/lite/toco/toco \ 官网给的这个路径不对       
./bazel-bin/tensorflow/contrib/lite/toco/toco \         
 —input_file=./mobilenet_v1_1.0_224/frozen_eval_graph_test.pb \  
 —output_file=./mobilenet_v1_1.0_224/tflite_model_test.tflite \  
 --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE \       
 --inference_type=FLOAT \           
 --input_shape="1,224, 224,3" \           
 --input_array=input \            
 --output_array=MobilenetV1/Predictions/Reshape_1

量化方式的转换(注意,只有量化训练的模型才能进行量化的tf_lite转换):

./bazel-bin/third_party/tensorflow/contrib/lite/toco/toco \
./bazel-bin/tensorflow/contrib/lite/toco/toco \
 --input_file=frozen_eval_graph.pb \
 --output_file=tflite_model.tflite \
 --input_format=TENSORFLOW_GRAPHDEF --output_format=TFLITE \
 --inference_type=QUANTIZED_UINT8 \
 --input_shape="1,224, 224,3" \
 --input_array=input \
 --output_array=outputs \
 --std_value=127.5 --mean_value=127.5

以上这篇Pytorch转tflite方式就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持三水点靠木。

Python 相关文章推荐
初学Python函数的笔记整理
Apr 07 Python
django初始化数据库的实例
May 27 Python
python主线程捕获子线程的方法
Jun 17 Python
Python中常用的内置方法
Jan 28 Python
python+selenium 定位到元素,无法点击的解决方法
Jan 30 Python
python计算阶乘和的方法(1!+2!+3!+...+n!)
Feb 01 Python
python获取磁盘号下盘符步骤详解
Jun 19 Python
关于python pycharm中输出的内容不全的解决办法
Jan 10 Python
PyCharm汉化安装及永久激活详细教程(靠谱)
Jan 16 Python
python 实现字符串下标的输出功能
Feb 13 Python
jupyter notebook插入本地图片的实现
Apr 13 Python
Python+MySQL随机试卷及答案生成程序的示例代码
Feb 01 Python
Python HTMLTestRunner库安装过程解析
May 25 #Python
Anaconda+vscode+pytorch环境搭建过程详解
May 25 #Python
5行Python代码实现图像分割的步骤详解
May 25 #Python
Win10用vscode打开anaconda环境中的python出错问题的解决
May 25 #Python
keras .h5转移动端的.tflite文件实现方式
May 25 #Python
Python虚拟环境venv用法详解
May 25 #Python
将keras的h5模型转换为tensorflow的pb模型操作
May 25 #Python
You might like
php邮件发送,php发送邮件的类
2011/03/24 PHP
PHP中extract()函数的妙用分析
2012/07/11 PHP
php中定时计划任务的实现原理
2013/01/08 PHP
Yii实现简单分页的方法
2016/04/29 PHP
php批量修改表结构实例
2017/05/24 PHP
使用Firebug对js进行断点调试的图文方法
2011/04/02 Javascript
基于jquery的loading 加载提示效果实现代码
2011/09/01 Javascript
一个可拖拽列宽表格实例演示
2012/11/26 Javascript
JS 实现Json查询的方法实例
2013/04/12 Javascript
jquery禁用右键单击功能屏蔽F5刷新
2014/03/17 Javascript
jQuery实现自定义右键菜单的树状菜单效果
2015/09/02 Javascript
js实现的动画导航菜单效果代码
2015/09/10 Javascript
jquery拖拽排序简单实现方法(效果增强版)
2016/02/16 Javascript
Bootstrap多级导航栏(级联导航)的实现代码
2016/03/08 Javascript
jquery常用的12个小功能
2016/07/22 Javascript
原生js实现ajax方法(超简单)
2016/09/20 Javascript
Angular使用Md5加密的解决方法
2017/09/16 Javascript
自定义PC微信扫码登录样式写法
2017/12/12 Javascript
JS删除数组里的某个元素方法
2018/02/03 Javascript
vue组件中的数据传递方法
2018/05/14 Javascript
小程序实现列表多个批量倒计时
2021/01/29 Javascript
JS立即执行的匿名函数用法分析
2019/11/04 Javascript
[02:27]刀塔重生降临
2015/10/14 DOTA
python实现堆和索引堆的代码示例
2018/03/19 Python
pycharm重置设置,恢复默认设置的方法
2018/10/22 Python
Python编程深度学习绘图库之matplotlib
2018/12/28 Python
东芝官网商城:还原日式美学,打造美好生活
2018/12/27 全球购物
优秀的教师个人的中文求职信
2013/09/21 职场文书
行政部主管岗位职责
2013/12/28 职场文书
口才训练演讲稿范文
2014/09/16 职场文书
2014年党的群众路线整改措施思想汇报
2014/10/12 职场文书
2015学校六五普法工作总结
2015/04/22 职场文书
日本读研:怎样写好一篇日本研究计划书?
2019/07/15 职场文书
话题作文之诚信
2019/11/28 职场文书
JavaScript如何利用Promise控制并发请求个数
2021/05/14 Javascript
详解Node.js如何处理ES6模块
2021/05/15 Javascript