How to convert darts to ONNX

username: crook52
cooked: <p>Hi.<br>
I want to convert <a href="https://github.com/quark0/darts/tree/master/cnn" rel="nofollow noopener">darts/cnn</a>’s model to TFlite, finaly.<br>
First of all, I tried to convert it to ONNX by below code.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import genotypes
from model import NetworkCIFAR as Network

genotype = eval("genotypes.%s" % 'DARTS')
model = Network(36, 10, 20, True, genotype)
model.load_state_dict(torch.load('./weights.pt'))
model = model.cuda()

onnx_model_path = './darts_model.onnx'
dummy_input = torch.randn(8,3,32,32)
input_names = ['image_array']
output_names = ['category']
torch.onnx.export(model,dummy_input, onnx_model_path,
                  input_names=input_names, output_names=output_names)
</code></pre>
<p>However, it couldn’t convert.<br>
Error is below.</p>
<pre><code class="lang-auto"> Traceback (most recent call last):
  File "&lt;stdin&gt;", line 2, in &lt;module&gt;
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/onnx/__init__.py", line 168, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 69, in export
    use_external_data_format=use_external_data_format)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 488, in _export
    fixed_batch_size=fixed_batch_size)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 334, in _model_to_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args, training)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 291, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, _force_outplace=False, _return_inputs_states=True)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/jit/__init__.py", line 278, in _get_trace_graph
    outs = ONNXTracedModule(f, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/XXXX_darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/jit/__init__.py", line 361, in forward
    self._force_outplace,
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/jit/__init__.py", line 351, in wrapper
    out_vars, _ = _flatten(outs)
RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type NoneType
</code></pre>
<p>Does onnx.export not correspond to darts?<br>
Could you please tell me how to fix it to convert to ONNX?</p>
<p>Finally, I am asking this same question in darts’s issue, sorry.</p>
<p>Thank you!</p>
<p>My environment</p>
<ul>
<li>Ubuntu 16.04</li>
<li>Python 3.6.10</li>
<li>CUDA 9.0</li>
<li>Pytorch 0.3.1(to search model), 1.5.1(to convert to ONNX)</li>
</ul>

username: ptrblck
cooked: <p>Based on the error message, it seems that an intermediate activation is <code>None</code> instead of a valid tensor.<br>
Is your model working fine in PyTorch eager mode and JIT (without the ONNX export)?</p>

username: crook52
cooked: <p>Thanks for your reply.<br>
Is it correct to understand PyTorch eager mode as normal mode?<br>
If so, the output of the model(input) is below. I think it is correct.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; out = model(dummy_input)
&gt;&gt;&gt; out
(tensor([[-0.0391, -0.0840,  0.1382, -0.0397,  0.0157, -0.0448, -0.0603, -0.0823,
          0.0025, -0.0009],
        [ 0.0257,  0.0031, -0.1880,  0.0768, -0.1047, -0.0392,  0.1393, -0.0419,
         -0.0437,  0.0032],
        [ 0.0090,  0.0097, -0.0768, -0.0383, -0.0220, -0.2048, -0.1315,  0.0117,
         -0.0538, -0.0613],
        [ 0.0438, -0.1284,  0.0325,  0.0441,  0.0736,  0.1941, -0.0407, -0.0634,
          0.1074, -0.0407],
        [ 0.0440,  0.0194,  0.0147,  0.0859,  0.2149, -0.0393,  0.1640,  0.0369,
         -0.1021, -0.1820],
        [-0.3142, -0.0726, -0.0694, -0.1064, -0.1595,  0.2461,  0.1174,  0.2102,
          0.1790,  0.2188],
        [-0.0522,  0.0327, -0.1626, -0.0955,  0.0625, -0.0061,  0.0662,  0.0667,
          0.1003,  0.0635],
        [ 0.1054, -0.0456,  0.0922,  0.0559,  0.1422, -0.1924, -0.2107, -0.0572,
         -0.0424, -0.1007]], device='cuda:0', grad_fn=&lt;AddmmBackward&gt;), tensor([[-0.0779, -0.4483,  0.5459, -0.4263,  0.3033, -0.0147,  0.1823,  0.2561,
          0.3321, -0.8131],
        [-0.1185, -0.1932,  0.2465,  0.3930, -0.0634,  0.2440, -0.0587, -0.5931,
          0.0938,  0.2163],
        [ 0.0699, -0.2207,  0.5958, -0.0778, -0.1024, -0.1841,  0.5211,  0.0760,
          0.2308,  0.1463],
        [ 0.1858, -0.0432,  0.3188, -0.0905,  0.1415, -0.6925,  0.1487, -0.2300,
          1.0883,  0.1186],
        [-0.1471,  0.1120,  0.3354, -0.3918,  0.0748, -0.5318,  0.0106, -0.4543,
          1.2513,  0.1778],
        [ 0.4499,  0.0425,  0.3949, -0.8790, -0.1463, -0.4942, -0.4362, -0.3380,
          0.3257,  0.2104],
        [ 0.2962,  0.0098,  0.6569,  0.0520,  0.1627, -0.4044,  0.2104, -0.2278,
          0.2411,  0.0337],
        [ 0.0591, -0.0795,  0.9120, -0.5483, -0.2887, -0.2304, -0.3799, -0.5769,
          0.5903, -0.6071]], device='cuda:0', grad_fn=&lt;AddmmBackward&gt;))

</code></pre>
<p>Regarding the JIT, I don’t know much about it, so could you please tell me how to check that?<br>
Here’s the <a href="https://github.com/quark0/darts/blob/master/cnn/model.py" rel="nofollow noopener">code</a>, is it possible to convert to ONNX even if ‘foward’ has ‘if’ in it?</p>
<pre><code class="lang-auto">class NetworkCIFAR(nn.Module):

  def __init__(self, C, num_classes, layers, auxiliary, genotype):
    super(NetworkCIFAR, self).__init__()
    self._layers = layers
    self._auxiliary = auxiliary

    stem_multiplier = 3 #RGB
    C_curr = stem_multiplier*C #C is input_channels
    self.stem = nn.Sequential(
      nn.Conv2d(3, C_curr, 3, padding=1, bias=False),
      nn.BatchNorm2d(C_curr)
    )
    
    C_prev_prev, C_prev, C_curr = C_curr, C_curr, C
    self.cells = nn.ModuleList()
    reduction_prev = False
    for i in range(layers):
      if i in [layers//3, 2*layers//3]:
        C_curr *= 2
        reduction = True
      else:
        reduction = False
      cell = Cell(genotype, C_prev_prev, C_prev, C_curr, reduction, reduction_prev)
      reduction_prev = reduction
      self.cells += [cell]
      C_prev_prev, C_prev = C_prev, cell.multiplier*C_curr
      if i == 2*layers//3:
        C_to_auxiliary = C_prev

    if auxiliary:
      self.auxiliary_head = AuxiliaryHeadCIFAR(C_to_auxiliary, num_classes)
    self.global_pooling = nn.AdaptiveAvgPool2d(1)
    self.classifier = nn.Linear(C_prev, num_classes)

  def forward(self, input):
    logits_aux = None
    s0 = s1 = self.stem(input)
    for i, cell in enumerate(self.cells):
      s0, s1 = s1, cell(s0, s1, self.drop_path_prob)
      if i == 2*self._layers//3:
        if self._auxiliary and self.training:
          logits_aux = self.auxiliary_head(s1)
    out = self.global_pooling(s1)
    logits = self.classifier(out.view(out.size(0),-1))
    return logits, logits_aux
</code></pre>
<p>I’m sorry for all of the questions.<br>
Thank you for your time.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="crook52" data-post="3" data-topic="89568">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/d26b3c/48.png" class="avatar"> crook52:</div>
<blockquote>
<p>Is it correct to understand PyTorch eager mode as normal mode?</p>
</blockquote>
</aside>
<p>Yes, sorry for the unclear naming. By “eager” mode I meant the normal Python usage.</p>
<p>Good to see the model is working generally.<br>
Could you , for the sake of debugging, remove the <code>logits_aux</code> from the <code>forward</code> and just return the <code>logits</code> and retry to export the model?</p>

username: ebarsoum
cooked: <p>It looks like it fails in tracing, can you try <code>torch.jit.trace</code> and see if it work or not?</p>

username: crook52
cooked: <p>I’m sorry for my late reply.<br>
After removing logits_aux  as you advised, it worked!!!<br>
I cannot thank you enough!</p>
<p>However, why it couldn’t convert with logits_aux??</p>
<pre><code class="lang-auto">cclass AuxiliaryHeadCIFAR(nn.Module):

  def __init__(self, C, num_classes):
    """assuming input size 8x8"""
    super(AuxiliaryHeadCIFAR, self).__init__()
    self.features = nn.Sequential(
      nn.ReLU(inplace=True),
      nn.AvgPool2d(5, stride=3, padding=0, count_include_pad=False), # image size = 2 x 2
      nn.Conv2d(C, 128, 1, bias=False),
      nn.BatchNorm2d(128),
      nn.ReLU(inplace=True),
      nn.Conv2d(128, 768, 2, bias=False),
      nn.BatchNorm2d(768),
      nn.ReLU(inplace=True)
    )
    self.classifier = nn.Linear(768, num_classes)

  def forward(self, x):
    x = self.features(x)
    x = self.classifier(x.view(x.size(0),-1))
    return x
</code></pre>

username: crook52
cooked: <p>Thank you for your reply, and sorry for my late reply.</p>
<aside class="quote no-group" data-username="ebarsoum" data-post="5" data-topic="89568">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ebarsoum/48/5515_2.png" class="avatar"> ebarsoum:</div>
<blockquote>
<p>It looks like it fails in tracing,</p>
</blockquote>
</aside>
<p>Yes, torch.jit.trace doesn’t work.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; torch.jit.trace(model,dummy_input)
/home/XXXX/darts/cnn/eval-EXP-20200710-150423/utils.py:105: TracerWarning: Converting a tensor to a Python index might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  mask = Variable(torch.cuda.FloatTensor(x.size(0), 1, 1, 1).bernoulli_(keep_prob))
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/jit/__init__.py", line 875, in trace
    check_tolerance, _force_outplace, _module_class)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/jit/__init__.py", line 1027, in trace_module
    module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, _force_outplace)
RuntimeError: Only tensors, lists and tuples of tensors can be output from traced functions
</code></pre>
<p>If I remove logits_aux following ptrblck’s advice, it work well.<br>
I don’t know why. So, I would appreciate it if you could inform when you find it.<br>
Thank you!</p>

username: ptrblck
cooked: <p>My best guess is that tracing the model didn’t go through the conditions where <code>aux_logits</code> is set to a tensor, so that it stayed <code>None</code> until the <code>return</code> statement. This could happen, e.g. if you called <code>model.eval()</code> before exporting it.</p>

username: crook52
cooked: <p>Thank you for your comment!<br>
The point is that I have no choice but to erase the lights_aux, tracing model.</p>