Pytorch训练的模型转.onnx、.pb、.h5、.tftile

Pytorch训练的模型转.onnx、.pb、.h5、.tftile

由于深度学习训练出后的参数往往是保存在指定的模型中的，这在使用时将不是那么方便，同时为了减小训练出的模型所占的内存空间，通常只保存了模型的参数，而未将整个网络一起保存，这样便直接导致了保存的模型无法直接转换成其他格式的模型。在此，我们提供了一个将pytorch模型转换成tftile的方法，希望能帮助大家任意转换各种模型以便使用！

Pytorch模型导入

这一部分我们根据pytorch保存整个模型以及只保存模型参数进行了不同的导入方式分析，并对将模型导入到CPU或GPU的方式进行了说明。
当Pytorch保存的是整个模型的时候，可以直接对模型进行导入，如：

import torch
torch_model = torch.load(model_path) #model_path是模型所在路径

要判断模型是否导入成功，只用运行一下下列语句，若不报错即导入成功：

torch_model.eval()

若上述语句运行报错，则极大可能是保存的为模型参数，此时需要结合网络结构一起将模型参数进行导入，以ResNet18为例：

import torch
from models.resnet import ResNet18 #将模型对应的网咯导入进来，大部分测试代码中能够找到

## 根据将模型导入至CPU或GPU对模型导入方式进行定义
def load_model(model, pretrained_path, load_to_cpu):
    print('Loading pretrained model from {}'.format(pretrained_path))
    if load_to_cpu:
        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
    else:
        device = torch.cuda.current_device()
        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
    model.load_state_dict(pretrained_dict, strict=False)
    return model

net = ResNet18(phase = 'test') #这里需要根据模型的传入参数进行修改
net = load_model(net, model_path, True)

导入完成后运行下列代码，查看模型是否导入成功

net.eval() # 查看模型是否导入成功
print(net) # 可以将网络结构进行打印，方便查看

当然，当我们直接运行别人的工程时，可在测试代码中找到net.eval()，在其后一般认为其模型导入结束

Pytorch模型转.onnx、.pb、.h5

这一部分的转换我们参考了博客的方法，亲测可行，这里直接附上代码，当然使用前需要自己提前安装相应模块，可直接通过pip install即可

import torch
import onnx
from onnx_tf.backend import prepare
from onnx2keras import onnx_to_keras
import keras
import tensorflow as tf

def pth_to_onnx(input_path,output_path):
    '''
    1)声明：使用本函数之前，必须保证你手上已经有了.pth模型文件.
    2)功能：本函数功能四将pytorch训练得到的.pth文件转化为onnx文件。
    '''
    torch_model = torch.load(input_path)       # pytorch模型加载,此处加载的模型包含图和参数
    # torch_model = selfmodel()  # 若只保存参数，selfmodel参考上述loadmodel进行改写
    torch_model.eval()
    x = torch.randn(1,1,28,28)          # 输入一张28*28的灰度图像并生成张量
    export_onnx_file = output_path         #输出.onnx文件的文件路径及文件名
    torch.onnx.export(torch_model,
                      x,
                      export_onnx_file,
                      opset_version=9,    #操作的版本，稳定操作集为9
                      do_constant_folding=True,          # 是否执行常量折叠优化
                      input_names=["input"],        # 输入名
                      output_names=["output"],       # 输出名
                      dynamic_axes={"input": {0: "batch_size"},         # 批处理变量
                                    "output": {0: "batch_size"}}
                      )
    # onnx_model = onnx.load('model_all.onnx')    #加载.onnx文件
    # onnx.checker.check_model(onnx_model)
    # print(onnx.helper.printable_graph(onnx_model.graph))       #打印.onnx文件信息

def onnx_to_pb(output_path):
    '''
    将.onnx模型保存为.pb文件模型
    '''
    model = onnx.load(output_path) #加载.onnx模型文件
    tf_rep = prepare(model)
    tf_rep.export_graph('model_all.pb')    #保存最终的.pb文件

def onnx_to_h5(output_path ):
    '''
    将.onnx模型保存为.h5文件模型,并打印出模型的大致结构
    '''
    onnx_model = onnx.load(output_path)
    k_model = onnx_to_keras(onnx_model, ['input'])
    keras.models.save_model(k_model, 'kerasModel.h5', overwrite=True, include_optimizer=True)    #第二个参数是新的.h5模型的保存地址及文件名
    # 下面内容是加载该模型，然后将该模型的结构打印出来
    model = tf.keras.models.load_model('kerasModel.h5')
    model.summary()
    print(model)
    
if __name__=='__main__':
    input_path = "model_all.pth"    #输入需要转换的.pth模型路径及文件名
    output_path = "model_all.onnx"  #转换为.onnx后文件的保存位置及文件名
    pth_to_onnx(input_path,output_path)  #执行pth转onnx函数，具体转换参数去该函数里面修改
    # onnx_pre(output_path)   #【可选项】若有需要，可以使用onnxruntime进行部署测试，看所转换模型是否可用，其中，output_path指加载进去的onnx格式模型所在路径及文件名
    # onnx_to_pb(output_path)   #将onnx模型转换为pb模型
    # onnx_to_h5(output_path )   #将onnx模型转换为h5模型

以上即实现了将Pytorch下训练保存的模型转换为.onnx、.pb、.h5。

.h5文件转.tftile文件

这一部分的转换我们参考了网址相关内容，转换方式如下：
首选使用pip install安装相应模块，并将刚才转换好的.h5模型进行导入及预测

import os
import time
import tensorflow as tf

# 恢复 keras 模型，并预测
keras_file = '../Models/model.h5'
model = tf.keras.models.load_model(keras_file)
# model.summary()
# tf.autograph.set_verbosity(0)

start_time = time.time()
pred = model.predict(image_bn)
stop_time = time.time()

print(f"prediction: {pred}")
print('time: {:.3f}ms'.format((stop_time - start_time) * 1000))
print("model size: {:.2f} MB".format(os.path.getsize(keras_file)/1024/1024))

输出为：
prediction: [[0.9730365 0.02696355]]
time: 126.018ms
model size: 7.46 MB
其次是将刚才导入的.h5模型转换成.tftile格式

# 最直接的保存为tflite，没有任何量化
converter = tf.lite.TFLiteConverter.from_keras_model(model)

tflite_model = converter.convert()

tflite_file = Path("../Models/Tflites/model.tflite")
tflite_file.write_bytes(tflite_model)

最后对导出的.tftile模型进行导入，查看模型相关信息，代码如下：

tflite_file = Path("../Models/Tflites/model.tflite")
# tflite 模型推理
interpreter = tf.lite.Interpreter(model_path=str(tflite_file))
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()[0]
output_details = interpreter.get_output_details()[0]

interpreter.set_tensor(input_details['index'], image_bn)

start_time = time.time()
interpreter.invoke()
stop_time = time.time()

output_data = interpreter.get_tensor(output_details['index'])
print(f"prediction: {output_data}")
print('time: {:.3f}ms'.format((stop_time - start_time) * 1000))
print("model size: {:.2f} MB".format(os.path.getsize(tflite_file)/1024/1024))

输出结果如：
prediction: [[0.9730364 0.02696356]]
time: 1.637ms
model size: 2.47 MB
至此，从Pytorch中训练保存的模型已顺利完成到.onnx、.pb、.h5、.tftile的转化，欢迎大家收藏转载，也感谢文中提及的博客及网址对本文的贡献！