PyTorch中F.conv2d从ONNX到TensorRT的支持

前言

在使用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 *