Unsupported aten::batchnorm

Marat · April 14, 2022, 4:12pm

Trying to transform model to NNAPI and getting following

Exception: Unsupported node kind ('aten::batch_norm') in node %input : Tensor = aten::batch_norm(%input.5, %self.2.running_var, %self.2.running_mean, %self.2.running_mean, %self.2.running_var, %11, %24, %23, %10), scope: __module.2 # /home/marat/anaconda3/envs/cexp/lib/python3.7/site-packages/torch/nn/functional.py:2422:0

On following

from torch import nn
import torch
from torch.backends._nnapi.prepare import convert_model_to_nnapi

model = nn.Sequential(nn.Conv2d(3, 3, kernel_size=1), nn.ReLU(), nn.BatchNorm2d(3), nn.ReLU())
# BUT this is OK
#model = nn.Sequential(nn.Conv2d(3, 3, kernel_size=1), nn.BatchNorm2d(3), nn.ReLU())
model.cpu()
model.eval()

input_float = torch.zeros(1, 3, 640, 640)
input_tensor = input_float

input_tensor = input_tensor.contiguous(memory_format=torch.channels_last)
input_tensor.nnapi_nhwc = True

with torch.no_grad():
    traced = torch.jit.trace(model, input_tensor)
nnapi_model = convert_model_to_nnapi(traced, input_tensor)

kmitsunami · April 19, 2022, 2:09pm

Looks like some layers are not supported by NNAPI yet. You can check the list of supported ops here:

github.com

pytorch/pytorch/blob/master/torch/backends/_nnapi/serializer.py#L784-L865

      
        
            ADDER_MAP = {
                "prim::GetAttr": lambda self, node:
                    self.add_getattr(node),
                "prim::Constant": lambda self, node:
                    self.add_constant_node(node),
                "prim::ListConstruct": lambda self, node:
                    self.add_list_construct(node),
                "prim::TupleConstruct": lambda self, node:
                    self.add_tuple_construct(node),
                "aten::unsqueeze": lambda self, node:
                    self.add_unsqueeze(node),
                "aten::to": lambda self, node:
                    self.add_to(node),
                "aten::detach": lambda self, node:
                    self._identity(node),
                "aten::reshape": lambda self, node:
                    self.add_reshape(node),
                "aten::flatten": lambda self, node:
                    self.add_flatten(node),
                "aten::slice": lambda self, node:

This file has been truncated. show original

I also recently found that BatchNorm is not on that list.

Marat · April 20, 2022, 8:40am

@kmitsunami you may try to remove one nn.RelU() before batchnorm and it compiles OK

# BUT this is OK
#model = nn.Sequential(nn.Conv2d(3, 3, kernel_size=1), nn.BatchNorm2d(3), nn.ReLU())

kmitsunami · April 20, 2022, 10:30am

I’ve tried NNAPI conversion of the combination of conv2d, batchnorm2d and relu before, but it failed if I remember correctly. Anyway, thanks for sharing.