Pytorch转tflite方式
目标是想把在服务器上用pytorch训练好的模型转换为可以在移动端运行的tflite模型。
最直接的思路是想把pytorch模型转换为tensorflow的模型，然后转换为tflite。但是这个转换目前没有发现比较靠谱的方法。
经过调研发现最新的tflite已经支持直接从keras模型的转换，所以可以采用keras作为中间转换的桥梁，这样就能充分利用keras高层API的便利性。
转换的基本思想就是用pytorch中的各层网络的权重取出来后直接赋值给keras网络中的对应layer层的权重。
转换为Keras模型后，再通过tf.contrib.lite.TocoConverter把模型直接转为tflite.
下面是一个例子，假设转换的是一个两层的CNN网络。
importtensorflowastf
fromtensorflowimportkeras
importnumpyasnp
importtorch
fromtorchvisionimportmodels
importtorch.nnasnn
#importtorch.nn.functionalasF
fromtorch.autogradimportVariable
classPytorchNet(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()
defforward(self,x):
returnself.feature(x)
definit_weights(self):
forminself.modules():
ifisinstance(m,nn.Conv2d):
nn.init.kaiming_normal_(
m.weight.data,mode='fan_out',nonlinearity='relu')
ifm.biasisnotNone:
m.bias.data.zero_()
ifisinstance(m,nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
defKerasNet(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)
returnmodel
classPytorchToKeras(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):
fori,layerinenumerate(self.kModel.layers):
iflen(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=[]
defadd_hooks(module):
defhook(module,input,output):
ifhasattr(module,"weight"):
#print(module)
self.__source_layers.append(module)
ifnotisinstance(module,nn.ModuleList)andnotisinstance(module,nn.Sequential)andmodule!=self.pModel:
hooks.append(module.register_forward_hook(hook))
self.pModel.apply(add_hooks)
self.pModel(input)
forhookinhooks:
hook.remove()
defconvert(self,input_size):
self.__retrieve_k_layers()
self.__retrieve_p_layers(input_size)
fori,(source_layer,target_layer)inenumerate(zip(self.__source_layers,self.__target_layers)):
print(source_layer)
weight_size=len(source_layer.weight.data.size())
transpose_dims=[]
foriinrange(weight_size):
transpose_dims.append(weight_size-i-1)
ifisinstance(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()])
elifisinstance(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()])
defsave_model(self,output_file):
self.kModel.save(output_file)
defsave_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')
#Loadthepretrainedmodel
pytorch_model=torch.load('test.pth')
##Timetotransferweights
converter=PytorchToKeras(pytorch_model,keras_model)
converter.convert((3,224,224))
##Savetheconvertedkerasmodelforlateruse
#converter.save_weights("keras.h5")
converter.save_model("keras_model.h5")
#convertkerasmodeltotflitemodel
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模型转换成tensorflowlite模型
1.把graph和网络模型打包在一个文件中
bazelbuildtensorflow/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
Forexample:
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.把第一步中生成的tensorflowpb模型转换为tflite模型
转换前需要先编译转换工具
bazelbuildtensorflow/contrib/lite/toco:toco
转换分两种，一种的转换为float的tflite，另一种可以转换为对模型进行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方式就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。