前言
在使用PyTorch进行网络训练得到.pth模型文件后,我们可能会做一些模型部署和加速的工作。这里一般会涉及到将PyTorch模型转为ONNX模型的过程。PyTorch自带了ONNX转换方法(torch.onnx.export
),可以很方便的将一些仅包含通用算子的网络的PyTorch模型转为ONNX格式。
另一方面,TensorRT是NVIDIA推出的一款高效深度学习模型推理框架,其包括了深度学习推理优化器和运行时,能够让深度学习推理应用拥有低时延和高吞吐的优点。在使用PyTorch训练得到网络模型后,我们希望在模型部署时通过TensorRT加速模型推理,那么可以先将PyTorch模型转为ONNX,然后再将ONNX转为TensorRT的engine。
而在常用的ONNX到TensorRT这一步,可能会遇到很多算子不支持的问题。例如本文提到的F.conv2d,在PyTorch到ONNX步骤能正常导出,但是从ONNX到TensorRT步骤则会报错。
问题定位
在一些通用网络结构(如StyleGAN系列)中,可能会用到F.conv2d操作,其权重是提前定义好的。
from torch.nn import functional as F
def make_kernel(k):
k = torch.tensor(k, dtype=torch.float32)
if k.ndim == 1:
k = k[None, :] * k[:, None]
k /= k.sum()
return k
class Blur(nn.Module):
def __init__(self, kernel, pad, upsample_factor=1):
super().__init__()
kernel = make_kernel(kernel)
if upsample_factor > 1:
kernel = kernel * (upsa