Arm Ethos-U85: Add support for BatchNorm2d not immediately following a convolution #17241
Description
🚀 The feature, motivation and pitch
I'm trying to export a model that has a BatchNorm2d layer without a Conv layer before that and its not supported. This is the model's repository.
Supporting this should be simple given that BatchNorm2d can just be broken down to mul and add
The reason I get in the log is: "BatchNorm2D with track_running_stats==True not immediately following a convolution is not supported for quantized TOSA backends."
If I change the layer to have track_running_stats=False its still not supported with reason: "Not included in BaseTOSASupportList or a registered tosa_support_check".
Alternatives
I managed to substitute the BatchNorm2d to a module that does the exact same thing and it allows export.
class CustomBatchNorm2d(nn.Module):
def __init__(self, weight, bias, running_mean, running_var):
super(CustomBatchNorm2d, self).__init__()
self.weight = weight.view(1, -1, 1, 1)
self.bias = bias.view(1, -1, 1, 1)
self.running_mean = running_mean.view(1, -1, 1, 1)
self.var = torch.sqrt(running_var + 1e-5).view(1, -1, 1, 1)
def forward(self, x):
x_normalized = (x - self.running_mean) / self.var
out = self.weight * x_normalized + self.bias
return out
Additional context
No response
RFC (Optional)
No response