【CNN】——EfficientNet转onnx转tensorrt的注意事项

EfficientNet, tensorrt
Code

想着将efficientnet转到tensorrt上推理，加入到benchmark库中，发现Efficientnet的输入上做了区别于conv2d的非规则池化，原始代码如下：

class Conv2dStaticSamePadding(nn.Conv2d):
    """2D Convolutions like TensorFlow's 'SAME' mode, with the given input image size.
       The padding mudule is calculated in construction function, then used in forward.
    """

    # With the same calculation as Conv2dDynamicSamePadding

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, image_size=None, **kwargs):
        super().__init__(in_channels, out_channels, kernel_size, stride, **kwargs)
        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2

        # Calculate padding based on image size and save it
        assert image_size is not None
        ih, iw = (image_size, image_size) if isinstance(image_size, int) else image_size
        kh, kw = self.weight.size()[-2:]
        sh, sw = self.stride
        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)
        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)
        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)
        if pad_h > 0 or pad_w > 0:
            self.static_padding = nn.ZeroPad2d((pad_w // 2, pad_w - pad_w // 2,
                                                pad_h // 2, pad_h - pad_h // 2))
        else:
            self.static_padding = nn.Identity()

    def forward(self, x):
        x = self.static_padding(x)
        x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
        return x

这里在输入的时候采用了zeropad2d,然后参考pytorch官方的document，其填充的方式为：a 4-tuple, uses (padding_left , padding_right , padding_top , padding_bottom )。

上面的模型可以转换成onnx，但是onnx——>tensorrt的时候上面的zeropad2d会有问题，所以这里考虑将zeropad2d融合到conv2d中。但是，Efficientnet中使用了非规则的zeropad2d，所以不能等值的转换成conv2d

修改后的代码：

class Conv2dStaticSamePadding(nn.Conv2d):
    """2D Convolutions like TensorFlow's 'SAME' mode, with the given input image size.
       The padding mudule is calculated in construction function, then used in forward.
    """

    # With the same calculation as Conv2dDynamicSamePadding

    def __init__(self, in_channels, out_channels, kernel_size, stride=1, image_size=None, **kwargs):
        super().__init__(in_channels, out_channels, kernel_size, stride, **kwargs)
        self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2

        # Calculate padding based on image size and save it
        assert image_size is not None
        ih, iw = (image_size, image_size) if isinstance(image_size, int) else image_size
        kh, kw = self.weight.size()[-2:]
        sh, sw = self.stride
        oh, ow = math.ceil(ih / sh), math.ceil(iw / sw)
        pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0)
        pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0)
        if pad_h > 0 or pad_w > 0:
            self.static_padding = nn.ZeroPad2d((pad_w // 2, pad_w - pad_w // 2,
                                                pad_h // 2, pad_h - pad_h // 2))
            self.padding = (pad_h // 2, pad_w // 2)             # 这种方式和原始的输出相差比较小
        else:
            self.static_padding = nn.Identity()
            self.padding = (0, 0)

    def forward(self, x):
        # x = self.static_padding(x)
        x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
        return x

采用以上方式和原始方式输出结果数值的对比：

original
tensor(309) tensor(9.8327)
tensor(309) tensor(9.0627)
tensor(599) tensor(9.0807)
tensor(304) tensor(6.9792)
tensor(310) tensor(9.6787)
tensor(310) tensor(10.4121)
tensor(310) tensor(9.1156)

修改之后的输出
tensor(309) tensor(9.7786)
tensor(309) tensor(9.0588)
tensor(599) tensor(9.0795)
tensor(304) tensor(6.9883)
tensor(310) tensor(9.6275)
tensor(310) tensor(10.3200)
tensor(310) tensor(9.1718)