Repository: autogluon/autogluon

Issue:
How to save the model created with an ENAS_Scheduler?
We have this nice tutorial showing how to use the ENAS_Scheduler to train a model.

[ENAS Tutorial](https://autogluon.mxnet.io/tutorials/nas/enas_proxylessnas.html) 

My question is:  How do we save and load trained models created that way?  I've tried 5 different ways and none of them seem to work.  

Based on what I read, the one I think I'm supposed to use is this:

`
mynet.export(modelFile, epoch=numEpochsTrained)
`

However, that produces this error message when I run it:
```
Traceback (most recent call last):
  File "/usr/lib/python3.6/pdb.py", line 1667, in main
    pdb._runscript(mainpyfile)
  File "/usr/lib/python3.6/pdb.py", line 1548, in _runscript
    self.run(statement)
  File "/usr/lib/python3.6/bdb.py", line 434, in run
    exec(cmd, globals, locals)
  File "<string>", line 1, in <module>
  File "/app/src/cli/boundryDetector.py", line 2, in <module>
    import sys,os;
  File "/app/src/code/boundryDetector/../lib/click/src/click/core.py", line 857, in __call__
    return self.main(*args, **kwargs)
  File "/app/src/code/boundryDetector/../lib/click/src/click/core.py", line 810, in main
    rv = self.invoke(ctx)
  File "/app/src/code/boundryDetector/../lib/click/src/click/core.py", line 1099, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/app/src/code/boundryDetector/../lib/click/src/click/core.py", line 613, in invoke
    return callback(*args, **kwargs)
  File "/app/src/cli/boundryDetector.py", line 108, in boundry_detector
    dnn.save(modelFile)
  File "/app/src/code/boundryDetector/./dnn/save.py", line 12, in save
    s.model.export(modelFile, epoch=s.numEpochsTrained)
  File "/usr/local/lib/python3.6/dist-packages/mxnet/gluon/block.py", line 1106, in export
    assert name in aux_names
AssertionError
```

I'm using:
```
autogluon==0.0.11
mxnet==1.6.0
```

Your kind assistance would be greatly appreciated. :)


Discussion:
@zhreshold Could you take a look at this?

Resolving patch:
diff --git a/autogluon/contrib/enas/enas.py b/autogluon/contrib/enas/enas.py
--- a/autogluon/contrib/enas/enas.py
+++ b/autogluon/contrib/enas/enas.py
@@ -70,7 +70,7 @@ class ENAS_Net(Cls):
             def __init__(self, *args, **kwargs):
                 kwvars.update(kwargs)
                 super().__init__(*args, **kwvars)
-                # 
+                #
                 self._modules = {}
                 self._kwspaces = collections.OrderedDict()
                 for k, module in kwvars.items():
@@ -242,7 +242,7 @@ def graph(self):
                     e.edge(pre_node, i)
                 pre_node = i
         return e
- 
+
     @property
     def kwspaces(self):
         return self._kwspaces
@@ -303,6 +303,51 @@ def __repr__(self):
         reprstr += ')\n'
         return reprstr
 
+    def export(self, path):
+        """Export HybridBlock to json format that can be loaded by
+        `gluon.SymbolBlock.imports`, `mxnet.mod.Module` or the C++ interface.
+
+        .. note:: When there are only one input, it will have name `data`. When there
+                  Are more than one inputs, they will be named as `data0`, `data1`, etc.
+
+        Parameters
+        ----------
+        path : str
+            Path to save model. Two files `path-symbol.json` and `path.params`
+            will be created.
+
+        Examples
+        --------
+
+        >>> mynet.export('enas')
+        >>> mynet_static = mx.gluon.nn.SymbolBlock.imports(
+            "enas-symbol.json", ['data'], "enas.params")
+        >>> x = mx.nd.random.uniform(shape=(1, 1, 28, 28))
+        >>> y = mynet_static(x)
+        """
+        from mxnet import npx as _mx_npx
+        from mxnet.util import is_np_array
+        from mxnet import ndarray
+        if not self._cached_graph:
+            raise RuntimeError(
+                "Please first call block.hybridize() and then run forward with "
+                "this block at least once before calling export.")
+        sym = self._cached_graph[1]
+        sym.save('%s-symbol.json'%path, remove_amp_cast=True)
+
+        arg_names = set(sym.list_arguments())
+        aux_names = set(sym.list_auxiliary_states())
+        arg_dict = {}
+        for name, param in self.collect_params().items():
+            if name in arg_names:
+                arg_dict['arg:%s'%name] = param._reduce()
+            elif name in aux_names:
+                arg_dict['aux:%s'%name] = param._reduce()
+            else:
+                pass
+        save_fn = _mx_npx.save if is_np_array() else ndarray.save
+        save_fn('%s.params'%(path), arg_dict)
+
 class Zero_Unit(gluon.HybridBlock):
     def hybrid_forward(self, F, x):
         return x


Test patch:
diff --git a/tests/unittests/test_contrib_enas.py b/tests/unittests/test_contrib_enas.py
new file mode 100644
--- /dev/null
+++ b/tests/unittests/test_contrib_enas.py
@@ -0,0 +1,65 @@
+import os
+
+import pytest
+import numpy as np
+
+import autogluon as ag
+import mxnet as mx
+import mxnet.gluon.nn as nn
+from autogluon.contrib.enas import *
+
+class Identity(mx.gluon.HybridBlock):
+    def hybrid_forward(self, F, x):
+        return x
+
+class ConvBNReLU(mx.gluon.HybridBlock):
+    def __init__(self, in_channels, channels, kernel, stride):
+        super().__init__()
+        padding = (kernel - 1) // 2
+        self.conv = nn.Conv2D(channels, kernel, stride, padding, in_channels=in_channels)
+        self.bn = nn.BatchNorm(in_channels=channels)
+        self.relu = nn.Activation('relu')
+    def hybrid_forward(self, F, x):
+        return self.relu(self.bn(self.conv(x)))
+
+@enas_unit()
+class ResUnit(mx.gluon.HybridBlock):
+    def __init__(self, in_channels, channels, hidden_channels, kernel, stride):
+        super().__init__()
+        self.conv1 = ConvBNReLU(in_channels, hidden_channels, kernel, stride)
+        self.conv2 = ConvBNReLU(hidden_channels, channels, kernel, 1)
+        if in_channels == channels and stride == 1:
+            self.shortcut = Identity()
+        else:
+            self.shortcut = nn.Conv2D(channels, 1, stride, in_channels=in_channels)
+    def hybrid_forward(self, F, x):
+        return self.conv2(self.conv1(x)) + self.shortcut(x)
+
+
+def test_enas_net():
+    mynet = ENAS_Sequential(
+        ResUnit(1, 8, hidden_channels=ag.space.Categorical(4, 8), kernel=ag.space.Categorical(3, 5), stride=2),
+        ResUnit(8, 8, hidden_channels=8, kernel=ag.space.Categorical(3, 5), stride=2),
+        ResUnit(8, 16, hidden_channels=8, kernel=ag.space.Categorical(3, 5), stride=2),
+        ResUnit(16, 16, hidden_channels=8, kernel=ag.space.Categorical(3, 5), stride=1, with_zero=True),
+        ResUnit(16, 16, hidden_channels=8, kernel=ag.space.Categorical(3, 5), stride=1, with_zero=True),
+        nn.GlobalAvgPool2D(),
+        nn.Flatten(),
+        nn.Activation('relu'),
+        nn.Dense(10, in_units=16),
+    )
+
+    mynet.initialize()
+    mynet.hybridize()
+
+    x = mx.nd.random.uniform(shape=(1, 1, 28, 28))
+    xx = mynet.evaluate_latency(x)
+    y = mynet(x)
+    assert mynet.nparams == 8714
+    mynet.export('enas')
+    mynet_static = mx.gluon.nn.SymbolBlock.imports("enas-symbol.json", ['data'], "enas.params")
+    yy = mynet_static(x)
+    np.testing.assert_almost_equal(y.asnumpy(), yy.asnumpy())
+
+if __name__ == "__main__":
+    test_enas_net()