【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)
更多推荐
所有评论(0)