溫馨提示×
您好，登錄后才能下訂單哦！
點擊 登錄注冊 即表示同意《億速云用戶服務條款》
您好，登錄后才能下訂單哦！
這篇文章主要為大家展示了Pytorch如何轉tflite，內容簡而易懂，希望大家可以學習一下，學習完之后肯定會有收獲的，下面讓小編帶大家一起來看看吧。
目標是想把在服務器上用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='h6')
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.h6")
converter.save_model("keras_model.h6")
# convert keras model to tflite model
converter = tf.contrib.lite.TocoConverter.from_keras_model_file(
"keras_model.h6")
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的內容，如果你們有學習到知識或者技能，可以把它分享出去讓更多的人看到。
免責聲明：本站發布的內容（圖片、視頻和文字）以原創、轉載和分享為主，文章觀點不代表本網站立場，如果涉及侵權請聯系站長郵箱：is@yisu.com進行舉報，并提供相關證據，一經查實，將立刻刪除涉嫌侵權內容。
