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)
Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐