pytorch-模型转换与算子读取

目录

 

一、pth转换onnx，并验证转换模型

二、模型参数统计

三、pth模型参数读取

---

一、pth转换onnx，并验证转换模型

当使用用tensorRT加速的时候，需要将pytorch模型转onnx模型。下面将介绍pytorch的pth模型如何转换成onnx，并验证转换模型。

---

先给官网链接：https://pytorch.org/docs/stable/onnx.html

咱们直接用一段代码来看：

import os.path as osp


import numpy as np

import onnx

import onnxruntime as ort

import torch

import torchvision


# your test input array

test_arr = np.random.randn(10, 3, 224, 224).astype(np.float32)


dummy_input = torch.randn(10, 3, 224, 224, device='cuda')

model = torchvision.models.resnet50(pretrained=True).cuda().eval()

print('pytorch result:', model(torch.from_numpy(test_arr).cuda()))


input_names = ["input"]

output_names = ["output"]


if not osp.exists('resnet50.onnx'):

# translate your pytorch model to onnx

torch.onnx.export(model, dummy_input, "resnet50.onnx", verbose=True, input_names=input_names, output_names=output_names)


model = onnx.load("resnet50.onnx")

ort_session = ort.InferenceSession('resnet50.onnx')

outputs = ort_session.run(None, {'input': test_arr})


print('onnx result:', outputs[0])

上面代码做的操作是先加载pytorch模型(包括结构和参数，然后设置为eval())，然后通过torch.onnx.export转化成onnx模型文件。

ps：

1. dummy_input就是一个输入的实例，仅提供输入shape、type等信息，里面什么内容不重要。所以用randn, zeros, ones都是可以的，或者就打开测试集的一张图片(如果你的模型输入是图片的话)； 

2. input_names和output_names没有特别要求，可以自己定义，但是要与session里输入的一致(上面代码倒数第二行)；你上面描述的input_names为"input"，下面也只能"input"；

3. pytorch模型的输入为tensor，而onnx的输入为array；这里是一个大坑。

---

二、模型参数统计

可以使用torchsummary，用来统计PyTorch每层的参数情况。一来可以用于参数剪枝优化，二来可以用于了解模型的参数分布。

安装：

pip install torchsummary

源代码链接：https://github.com/sksq96/pytorch-summary

使用：

from torchvision.models.alexnet import alexnet

from torchsummary import summary



model = alexnet(pretrained=True).eval().cuda()

summary(model, input_size=(3, 224, 224), batch_size=-1)

效果：

结束~ 

---

三、pth模型参数读取

目的：把训练好的pth模型参数提取出来，然后用其他方式部署到边缘设备。

Pytorch给了很方便的读取参数接口：nn.Module.parameters()

直接看demo：

from torchvision.models.alexnet import alexnet

model = alexnet(pretrained=True).eval().cuda()

parameters = model.parameters()

for p in parameters:

numpy_para = p.detach().cpu().numpy()

print(type(numpy_para))

print(numpy_para.shape)

上面得到的numpy_para就是numpy参数了~

Note：model.parameters()是以一个生成器的形式迭代返回每一层的参数。所以用for循环读取到各层的参数，循环次数就表示层数。而每一层的参数都是torch.nn.parameter.Parameter类型，是Tensor的子类，所以直接用tensor转numpy（即p.detach().cpu().numpy()）的方法就可以直接转成numpy矩阵。方便又好用，爆赞~