Unsupported aten::batchnorm

Marat: <p>Trying to transform model to NNAPI and getting following</p>
<pre><code class="lang-auto">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
</code></pre>
<p>On following</p>
<pre><code class="lang-auto">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)
</code></pre>

kmitsunami: <p>Looks like some layers are not supported by NNAPI yet. You can check the list of supported ops here:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/master/torch/backends/_nnapi/serializer.py#L784-L865">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/master/torch/backends/_nnapi/serializer.py#L784-L865" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/backends/_nnapi/serializer.py#L784-L865" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/torch/backends/_nnapi/serializer.py#L784-L865</a></h4>



    <pre class="onebox">      <code class="lang-py">
        <ol class="start lines" start="784" style="counter-reset: li-counter 783 ;">
            <li>ADDER_MAP = {</li>
            <li>    "prim::GetAttr": lambda self, node:</li>
            <li>        self.add_getattr(node),</li>
            <li>    "prim::Constant": lambda self, node:</li>
            <li>        self.add_constant_node(node),</li>
            <li>    "prim::ListConstruct": lambda self, node:</li>
            <li>        self.add_list_construct(node),</li>
            <li>    "prim::TupleConstruct": lambda self, node:</li>
            <li>        self.add_tuple_construct(node),</li>
            <li>    "aten::unsqueeze": lambda self, node:</li>
            <li>        self.add_unsqueeze(node),</li>
            <li>    "aten::to": lambda self, node:</li>
            <li>        self.add_to(node),</li>
            <li>    "aten::detach": lambda self, node:</li>
            <li>        self._identity(node),</li>
            <li>    "aten::reshape": lambda self, node:</li>
            <li>        self.add_reshape(node),</li>
            <li>    "aten::flatten": lambda self, node:</li>
            <li>        self.add_flatten(node),</li>
            <li>    "aten::slice": lambda self, node:</li>
        </ol>
      </code>
    </pre>


  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/master/torch/backends/_nnapi/serializer.py#L784-L865" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
I also recently found that BatchNorm is not on that list.</p>

Marat: <p><a class="mention" href="/u/kmitsunami">@kmitsunami</a> you may try to remove one nn.RelU() before batchnorm and it compiles OK</p>
<pre><code class="lang-auto"># BUT this is OK
#model = nn.Sequential(nn.Conv2d(3, 3, kernel_size=1), nn.BatchNorm2d(3), nn.ReLU())
</code></pre>

kmitsunami: <p>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.</p>