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)