第二章:解决模型部署中的难题 — mmdeploy 0.12.0 文档
PyTorch 74.自定义操作torch.autograd.Function - 知乎
简述python中的@staticmethod作用及用法 - 腾讯云开发者社区-腾讯云
(一)指定的图片放大倍数是固定的,本节通过部署一个支持动态放大倍数的模型,来体验模型部署中可能遇到的困难。
模型部署中常见的难题
- 模型的动态化:如何让推理框架中静态的输入、输出形状和结构,泛化为动态的。
- 新算子的实现:在ONNX和推理引擎中支持新的算子。
- 中间表示与推理引擎的兼容问题:为特定推理引擎定制模型代码。
import torch
from torch import nn
from torch.nn.functional import interpolate
import torch.onnx
import cv2
import numpy as np
class SuperResolutionNet(nn.Module):#继承nn.Module类,并实现forward方法
def __init__(self):#构造函数,类的本身实例变量self
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)
self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2)
self.relu = nn.ReLU()
def forward(self, x, upscale_factor):#将原先构造函数初始化阶段nn.Upsample固化的放大倍数upscale_factor,调整为运行阶段输入的interpolate放大倍数
x = interpolate(x,
scale_factor=upscale_factor.item(),#这里是因为pytorch模型导出到ONNX模型时输入必须是torch.Tensor,int类型是不可以的,因此使用torch.Tensor.item()把torch.Tensor转换成数值。
mode='bicubic',
align_corners=False)
out = self.relu(self.conv1(x))
out = self.relu(self.conv2(out))
out = self.conv3(out)
return out
def init_torch_model():
torch_model = SuperResolutionNet()
# Please read the code about downloading 'srcnn.pth' and 'face.png' in
# https://mmdeploy.readthedocs.io/zh_CN/latest/tutorials/chapter_01_introduction_to_model_deployment.html#pytorch
state_dict = torch.load('srcnn.pth')['state_dict']
# Adapt the checkpoint
for old_key in list(state_dict.keys()):
new_key = '.'.join(old_key.split('.')[1:])
state_dict[new_key] = state_dict.pop(old_key)
torch_model.load_state_dict(state_dict)
torch_model.eval()
return torch_model
model = init_torch_model()
input_img = cv2.imread('face.png').astype(np.float32)
# HWC to NCHW
input_img = np.transpose(input_img, [2, 0, 1])
input_img = np.expand_dims(input_img, 0)
# Inference
torch_output = model(torch.from_numpy(input_img), torch.tensor(3)).detach().numpy()
# NCHW to HWC
torch_output = np.squeeze(torch_output, 0)
torch_output = np.clip(torch_output, 0, 255)
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8)
# Show image
cv2.imwrite("face_torch_2.png", torch_output)
导出动态超分辨率模型
import torch
from tensorrtdemo.dyn_torch_inference import init_torch_model
x = torch.randn(1, 3, 256, 256)#随机产生一个大小固定的输入
y = 3
model = init_torch_model()#获取带参数的模型
with torch.no_grad():#上下文管理器,禁止梯度计算,在网络推断中可以节省显存
torch.onnx.export(
model,#带参模型
(x, torch.tensor(3)),#随机的输入形式,将torch.tensor(3)代替3,是因为PyTorch输入必须全部是torch.Tensor
"srcnn2.onnx",#存储的onnx文件
opset_version=11,# opset版本是指onnx定义的算子集合
input_names=['input','factor'],#输入对象名称
output_names=['output'])#输出对象名称
错误“TracerWarning: Converting a tensor to a Python number might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
scale_factor=upscale_factor.item(),”
提示追踪失败。
Netorn查看对比,发现即便模型推理输入设置成了2个,但实际上ONNX的输入还是只有一个Input
这是因为即便是将torch.tensor(3)作为输入,但torch.Tensor.item()还是把数据从Tensor取出来了,因而ONNX导出时,操作无法被记录,因而有一个TraceWarning,因此scale_factor还是一个固定值。因此还只是一个输入。
pytorch转ONNX实际上是将PyTorch操作映射成ONNX定义的算子,如PyTorch的nn.Upsample和interpolate都转换成ONNX的Resize操作。
点击resize算子,展开scales
scales是一个长度为4的一维张量,内容是[1,1,3,3]表示在NCWH对应维度放缩系数,其类型为Initializer,表示该scales值是根据常量scale_factor初始化出来的,假定自定义一个ONNX的Resize算子,使得scales成为一个可变量而不是常量,则与x一样,该超分辨率模型就可以动态缩放。
通过运行模型来验证自定义的PyTorch算子是否正确
import torch
from torch import nn
from torch.nn.functional import interpolate
import torch.onnx
import cv2
import numpy as np
class NewInterpolate(torch.autograd.Function):#自定义求导方式(即Extending torch.autograd)
#Function只定义一个操作,因为其无法保存参数;Module是保存了参数
#Function__init__, forward, backward(需要自己写求导公式);Module只需定义__init__和forward,而backward的计算由自动求导机制构成
#Module是由一系列Function组成,因此其在forward的过程中,Function和Variable组成了计算图,在backward时,只需调用Function的backward就得到结果,因此Module不需要再定义backward。
#
@staticmethod
def symbolic(g, input, scales): #新算子映射到ONNX算子,通过算子的symbolic方法决定
# 第一个参数必须是g
#自定义输入和forward函数一样。
#具体实现由g.op实现,每个参数都可以映射ONNX的算子属性
return g.op("Resize",#参照Netron设置对应参数属性
input,
g.op("Constant",
value_t=torch.tensor([], dtype=torch.float32)),
scales,#希望scale参数由输入动态决定,因此填入ONNX的scales时可以吧 symbolic方法输入参数的scales填入。
coordinate_transformation_mode_s="pytorch_half_pixel",
cubic_coeff_a_f=-0.75,
mode_s='cubic',
nearest_mode_s="floor")
#forward()和backward()都应该是staticmethod。
#静态方法就是类对外部函数的封装,有助于优化代码结构和提高程序的可读性。
@staticmethod
def forward(ctx, input, scales):#输入:图像、放缩比例[1,1,w,h]的张量,w、h分别是图像的宽和高的放大倍数。
#forward的第一个参数必须是ctx
scales = scales.tolist()[-2:]#截取输入张量的后两个元素,把这两个元素以list格式传入到interpolate的scale_factor参数
return interpolate(input,
scale_factor=scales,
mode='bicubic',
align_corners=False)
class StrangeSuperResolutionNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)
self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2)
self.relu = nn.ReLU()
def forward(self, x, upscale_factor):
x = NewInterpolate.apply(x, upscale_factor)
out = self.relu(self.conv1(x))
out = self.relu(self.conv2(out))
out = self.conv3(out)
return out
def init_torch_model():
torch_model = StrangeSuperResolutionNet()
state_dict = torch.load('srcnn.pth')['state_dict']
# Adapt the checkpoint
for old_key in list(state_dict.keys()):
new_key = '.'.join(old_key.split('.')[1:])
state_dict[new_key] = state_dict.pop(old_key)
torch_model.load_state_dict(state_dict)
torch_model.eval()
return torch_model
model = init_torch_model()
factor = torch.tensor([1, 1, 3, 3], dtype=torch.float)
input_img = cv2.imread('face.png').astype(np.float32)
# HWC to NCHW
input_img = np.transpose(input_img, [2, 0, 1])
input_img = np.expand_dims(input_img, 0)
# Inference
torch_output = model(torch.from_numpy(input_img), factor).detach().numpy()
# NCHW to HWC
torch_output = np.squeeze(torch_output, 0)
torch_output = np.clip(torch_output, 0, 255)
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8)
# Show image
cv2.imwrite("face_torch_3.png", torch_output)
import torch
from tensorrtdemo.dyn_custom_opt_inference import init_torch_model
x = torch.randn(1, 3, 256, 256)
factor = torch.tensor([1, 1, 3, 3], dtype=torch.float)
model = init_torch_model()#获取带参数的模型
with torch.no_grad():
torch.onnx.export(model, (x, factor),
"srcnn3.onnx",
opset_version=11,
input_names=['input', 'factor'],
output_names=['output'])
可以看到输入变成了两个
import onnxruntime
import numpy as np
import cv2
input_img = cv2.imread('face.png').astype(np.float32)
# HWC to NCHW
input_img = np.transpose(input_img, [2, 0, 1])
input_img = np.expand_dims(input_img, 0)
input_factor = np.array([1, 1, 4, 4], dtype=np.float32)
ort_session = onnxruntime.InferenceSession("srcnn3.onnx")
ort_inputs = {'input': input_img, 'factor': input_factor}
ort_output = ort_session.run(None, ort_inputs)[0]
ort_output = np.squeeze(ort_output, 0)
ort_output = np.clip(ort_output, 0, 255)
ort_output = np.transpose(ort_output, [1, 2, 0]).astype(np.uint8)
cv2.imwrite("face_ort_3.png", ort_output)
总结
模型部署中的困难:模型的动态化;新算子的实现;框架间的兼容。
PyTorch转ONNX:实际上是将操作转换为ONNX定义的算子。
自定义算子映射到ONNX算子:修改继承torch.autograd.Function算子的symbolic。