pytorch模型转onnx转ncnn
一、pth转onnx
训练生成好.pth文件后，参照网上的代码转onnx,网上代码如下
import torch
import torchvision
#define resnet18 model
model = torchvision.models.resnet18(pretrained=True)
#define input shape
x = torch.rand(1, 3, 224, 224)
#define input and output nodes, can be customized
input_names = ["x"]
output_names = ["y"]
#convert pytorch to onnx
torch_out = torch.onnx.export(model, x, "resnet18.onnx", input_names=input_names, output_names=output_names)
主要是model这个变量不好操作，我们把这段代码，插入到我们的模型测试文件中生成model的地方，下面这段代码是测试文件中生成model的代码：
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone=backbone, pretrained=False, result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=scale, gpu_id=gpu_id)
total_frame = 0.0
total_time = 0.0
for img_path in tqdm(img_paths):
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(save_txt_folder, 'res_' + img_name + '.txt')
_, boxes_list, t = model.predict(img_path)
total_frame += 1
total_time += t
img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
cv2.imwrite(os.path.join(save_img_folder, '{}.jpg'.format(img_name)), img)
np.savetxt(save_name, boxes_list.reshape(-1, 8), delimiter=',', fmt='%d')
print('fps:{}'.format(total_frame / total_time))
return save_txt_folder
我们在model变量后插入代码，并注释掉model，因为这里的model不能利用，重新写了代码插入
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone=backbone, pretrained=False, result_num=config.n)
#model = Pytorch_model(model_path, net=net, scale=scale, gpu_id=gpu_id)
##################插入的代码#############################################
net1 = load_model(net, './output/Best_926_r0.424205_p0.406323_f10.415072.pth')
# define input shape
x = torch.rand(1, 3, 224, 224)#这里改成自己的前处理操作
# define input and output nodes, can be customized
input_names = ["x"]
output_names = ["y"]
# convert pytorch to onnx
torch_out = torch.onnx.export(net1, x, "resnet18.onnx", input_names=input_names, output_names=output_names)
exit()
#########################################################################
total_frame = 0.0
total_time = 0.0
load_model函数的代码如下：
def load_model(model, model_path):
start_epoch = 0
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
# print('loaded {}, epoch {}'.format(model_path, checkpoint['epoch']))
try:
state_dict_ = checkpoint["state_dict"]
state_dict = {}
except:
state_dict_ = checkpoint
state_dict = {}
# convert data_parallal to model
for k in state_dict_:
if k.startswith('module') and not k.startswith('module_list'):
state_dict[k[7:]] = state_dict_[k]
else:
state_dict[k] = state_dict_[k]
model_state_dict = model.state_dict()
# check loaded parameters and created model parameters
msg = 'If you see this, your model does not fully load the ' + \
'pre-trained weight. Please make sure ' + \
'you have correctly specified --arch xxx ' + \
'or set the correct --num_classes for your own dataset.'
for k in state_dict:
if k in model_state_dict:
if state_dict[k].shape != model_state_dict[k].shape:
print('Skip loading parameter {}, required shape{}, ' \
'loaded shape{}. {}'.format(
k, model_state_dict[k].shape, state_dict[k].shape, msg))
state_dict[k] = model_state_dict[k]
else:
print('Drop parameter {}.'.format(k) + msg)
for k in model_state_dict:
if not (k in state_dict):
print('No param {}.'.format(k) + msg)
state_dict[k] = model_state_dict[k]
model.load_state_dict(state_dict, strict=False)
return model
运行后生成onnx文件，会有提示：
/home/admin1/anaconda3/lib/python3.7/site-packages/torch/onnx/symbolic_helper.py:243: UserWarning: You are trying to export the model with onnx:Upsample for ONNX opset version 9. This operator might cause results to not match the expected results by PyTorch.
ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. Attributes to determine how to transform the input were added in onnx:Resize in opset 11 to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).
We recommend using opset 11 and above for models using this operator.
"" + str(_export_onnx_opset_version) + ". "
我们在torch.onnx.export函数的参数里加上opset_version=11
torch_out = torch.onnx.export(net1, x, "resnet18.onnx", input_names=input_names, output_names=output_names,opset_version=11)
提示就会消失。
导出的onnx模型包含许多冗余的维度，这是不支持ncnn的，所以需要进行去掉冗余的维度。
python -m onnxsim resnet18.onnx resnet18-sim.onnx
二、onnx转ncnn
1、编译ncnn
参考文章：
https://blog.csdn.net/xiao13mm/article/details/106165477
2、转换成ncnn
ncnn编译完后，在build/tools/onnx里会生成个可执行文件onnx2ncnn
./onnx2ncnn resnet18-sim.onnx resnet18.param resnet18.bin
最后得到的parm和bin文件即是ncnn所需要的模型。
QQ交流群：1080729300