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方式就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持毛票票。