pkl转onnx模型

实现模型的转换代码如下

#from sklearn.externals import joblib
#import joblib
#from skl2onnx import convert_sklearn
#from skl2onnx.common.data_types 
import onnxmltools
import torchen
import torchvision
import torchvision.models as models
from lib import Model, ClassAverages
# Update the input name and path for your sklearn model
input_skl_model = 'epoch_10.pkl'
# input data type for your sklearn model
bins_no = 2
device = torch.device("cpu:0")
my_vgg = models.vgg19_bn(pretrained=True)
torch_model = Model.Model(features=my_vgg.features, bins=bins_no).to(device)
checkpoint = torch.load(input_skl_model,map_location='cpu')
torch_model.load_state_dict(checkpoint['model_state_dict'])
#torch_model.eval()
#input_data_type = [('model_state_dict', FloatTensorType([1,3,244,244]))]
dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)
# Change this path to the output name and path for the ONNX model
output_onnx_model = 'model.onnx'
# Load your sklearn model
#print(torch_model)
#skl_model = joblib.load(input_skl_model)
#print(skl_model)
#torch.onnx._export(torch_model, dummy_input, output_onnx_model, verbose=True,training=False, input_names=None, output_names=None)
input_names=["input"]
output_names=["orientation","confidence","dimension"]
torch.onnx.export(torch_model,dummy_input,output_onnx_model,input_names=input_names,training=False,output_names=output_names,opset_version=11)
# Convert the sklearn model into ONNX
#onnx_model = onnxmltools.convert_sklearn(torch_model, initial_types=input_data_type)
# Save as protobuf
#onnxmltools.utils.save_model(onnx_model, output_onnx_model)

测试模型路径：weights\epoch_10.pkl
生成模型：model.onnx
执行的指令：python Transfer.py

测试工程 测试模型和相关数据参见：https://download.csdn.net/download/weixin_40912644/74086010