Pytorch如何转tflite-创新互联
这篇文章主要为大家展示了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)
分享标题:Pytorch如何转tflite-创新互联
本文网址:http://scyanting.com/article/dcjgip.html