[Retiarii] Base execution engine, codegen and trainer (#3059)

test/ut/retiarii/debug_strategy.py

+import json
+import os
+
+from nni.retiarii import Model, submit_models, wait_models
+
+
+def single_model_startegy():
+    with open(os.path.join(os.path.dirname(__file__), 'mnist_pytorch.json')) as f:
+        ir = json.load(f)
+    model = Model._load(ir)
+    submit_models(model)
+    wait_models(model)
+    print('Strategy says:', model.metric)
+
+
+if __name__ == '__main__':
+    single_model_startegy()

test/ut/retiarii/engine.py

+import json
+import os
+import sys
+import threading
+import unittest
+
+from nni.retiarii import Model, submit_models
+from nni.retiarii.codegen import model_to_pytorch_script
+from nni.retiarii.integration import RetiariiAdvisor, register_advisor
+from nni.retiarii.trainer import PyTorchImageClassificationTrainer
+from nni.retiarii.utils import import_
+
+
+class CodeGenTest(unittest.TestCase):
+    def test_mnist_example_pytorch(self):
+        with open('mnist_pytorch.json') as f:
+            model = Model._load(json.load(f))
+            script = model_to_pytorch_script(model)
+        with open('debug_mnist_pytorch.py') as f:
+            reference_script = f.read()
+        self.assertEqual(script.strip(), reference_script.strip())
+
+
+class TrainerTest(unittest.TestCase):
+    def test_trainer(self):
+        Model = import_('debug_mnist_pytorch._model')
+        trainer = PyTorchImageClassificationTrainer(
+            Model(),
+            dataset_kwargs={'root': 'data/mnist', 'download': True},
+            dataloader_kwargs={'batch_size': 32},
+            optimizer_kwargs={'lr': 1e-3},
+            trainer_kwargs={'max_epochs': 1}
+        )
+        trainer.fit()
+
+
+class EngineTest(unittest.TestCase):
+
+    def test_submit_models(self):
+        os.makedirs('generated', exist_ok=True)
+        from nni.runtime import protocol
+        protocol._out_file = open('generated/debug_protocol_out_file.py', 'wb')
+        anything = lambda: None
+        advisor = RetiariiAdvisor(anything)
+        with open('mnist_pytorch.json') as f:
+            model = Model._load(json.load(f))
+        submit_models(model, model)
+
+        advisor.stopping = True
+        advisor.default_worker.join()
+        advisor.assessor_worker.join()
+
+    def test_execution_engine(self):
+        pass

test/ut/retiarii/fake_search_space.json

+{}

test/ut/retiarii/mnist.json

+{
+    "_model": {
+        "inputs": ["image"],
+        "outputs": ["metric"],
+
+        "nodes": {
+            "stem": {"cell": "stem"},
+            "flatten": {"type": "Flatten"},
+            "fc1": {"type": "Dense", "units": 1024, "activation": "relu"},
+            "fc2": {"type": "Dense", "units": 10},
+            "softmax": {"type": "Softmax"}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["stem", 0]},
+            {"head": ["stem", 0], "tail": ["flatten", null]},
+            {"head": ["flatten", null], "tail": ["fc1", null]},
+            {"head": ["fc1", null], "tail": ["fc2", null]},
+            {"head": ["fc2", null], "tail": ["softmax", null]},
+            {"head": ["softmax", null], "tail": ["_outputs", 0]}
+        ]
+    },
+
+    "stem": {
+        "nodes": {
+            "conv1": {"type": "Conv2D", "filters": 32, "kernel_size": 5, "activation": "relu"},
+            "pool1": {"type": "MaxPool2D", "pool_size": 2},
+            "conv2": {"type": "Conv2D", "filters": 64, "kernel_size": 5, "activation": "relu"},
+            "pool2": {"type": "MaxPool2D", "pool_size": 2}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["conv1", null]},
+            {"head": ["conv1", null], "tail": ["pool1", null]},
+            {"head": ["pool1", null], "tail": ["conv2", null]},
+            {"head": ["conv2", null], "tail": ["pool2", null]},
+            {"head": ["pool2", null], "tail": ["_outputs", 0]}
+        ]
+    }
+}

test/ut/retiarii/mnist_pytorch.json

+{
+    "_model": {
+        "inputs": ["image"],
+        "outputs": ["metric"],
+
+        "nodes": {
+            "stem": {"type": "_cell", "cell": "stem"},
+            "flatten": {"type": "Flatten"},
+            "fc1": {"type": "Dense", "parameters": {"out_features": 256, "in_features": 1024}},
+            "fc2": {"type": "Dense", "parameters": {"out_features": 10, "in_features": 256}},
+            "softmax": {"type": "Softmax"}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["stem", null]},
+            {"head": ["stem", null], "tail": ["flatten", null]},
+            {"head": ["flatten", null], "tail": ["fc1", null]},
+            {"head": ["fc1", null], "tail": ["fc2", null]},
+            {"head": ["fc2", null], "tail": ["softmax", null]},
+            {"head": ["softmax", null], "tail": ["_outputs", 0]}
+        ]
+    },
+
+    "stem": {
+        "nodes": {
+            "conv1": {"type": "Conv2d", "parameters": {"out_channels": 32, "in_channels": 1, "kernel_size": 5}},
+            "pool1": {"type": "MaxPool2d", "parameters": {"kernel_size": 2}},
+            "conv2": {"type": "Conv2d", "parameters": {"out_channels": 64, "in_channels": 32, "kernel_size": 5}},
+            "pool2": {"type": "MaxPool2d", "parameters": {"kernel_size": 2}}
+        },
+
+        "edges": [
+            {"head": ["_inputs", 0], "tail": ["conv1", null]},
+            {"head": ["conv1", null], "tail": ["pool1", null]},
+            {"head": ["pool1", null], "tail": ["conv2", null]},
+            {"head": ["conv2", null], "tail": ["pool2", null]},
+            {"head": ["pool2", null], "tail": ["_outputs", 0]}
+        ]
+    },
+
+    "_training_config": {
+        "module": "nni.retiarii.trainer.PyTorchImageClassificationTrainer",
+        "kwargs": {
+            "dataset_kwargs": {
+                "root": "data/mnist",
+                "download": true
+            },
+            "dataloader_kwargs": {
+                "batch_size": 32
+            },
+            "optimizer_kwargs": {
+                "lr": 1e-3
+            },
+            "trainer_kwargs": {
+                "max_epochs": 1
+            }
+        }
+    }
+}

test/ut/retiarii/nni.yaml

+authorName: nni
+experimentName: naive
+trialConcurrency: 3
+maxExecDuration: 1h
+maxTrialNum: 10
+trainingServicePlatform: local
+searchSpacePath: fake_search_space.json
+useAnnotation: false
+advisor:
+  codeDir: .
+  classFileName: advisor_entry.py
+  className: RetiariiAdvisor
+  classArgs:
+    strategy: debug_strategy.single_model_startegy
+trial:
+  command: python -m nni.retiarii.trial_entry
+  codeDir: ../..
+  gpuNum: 0